duckdb-nsql

MotherDuck 和 Numbers Station 制作的 7B 参数文本转 SQL 模型。

7b

30.5K 13 个月前

17 个标签

自述文件

duckdb-nsql model

DuckDB-NSQL 是一个 70 亿参数的文本转 SQL 模型，专门为 SQL 生成任务而设计。

此模型基于 Meta 的原始 Llama-2 7B 模型，并在通用 SQL 查询数据集上进一步预训练，然后在由 DuckDB 文本转 SQL 对组成的数据集上进行微调。

使用方法

示例提示

Provided this schema:

CREATE TABLE taxi (
    VendorID bigint,
    tpep_pickup_datetime timestamp,
    tpep_dropoff_datetime timestamp,
    passenger_count double,
    trip_distance double,
    fare_amount double,
    extra double,
    tip_amount double,
    tolls_amount double,
    improvement_surcharge double,
    total_amount double,
);

Give me taxis with more than 2 passengers

示例输出

SELECT * FROM taxi WHERE passenger_count > 2

设置系统提示

此模型期望在系统提示中输入 schema

/set system """Here is the database schema that the SQL query will run on:
CREATE TABLE taxi (
    VendorID bigint,
    tpep_pickup_datetime timestamp,
    tpep_dropoff_datetime timestamp,
    passenger_count double,
    trip_distance double,
    fare_amount double,
    extra double,
    tip_amount double,
    tolls_amount double,
    improvement_surcharge double,
    total_amount double,
);"""

一旦在系统提示中提供了 schema，该模型将在后续响应中使用它。

对于以下提示

get all columns ending with _amount from taxi table

该模型将输出如下内容

SELECT COLUMNS('.*_amount') FROM taxi;

API 示例

$ curl https://127.0.0.1:11434/api/generate -d '{
    "model": "duckdb-nsql:7b-q4_0",
    "system": "Here is the database schema that the SQL query will run on: CREATE TABLE taxi (VendorID bigint, tpep_pickup_datetime timestamp, tpep_dropoff_datetime timestamp, passenger_count double, trip_distance double, fare_amount double, extra double, tip_amount double, tolls_amount double, improvement_surcharge double, total_amount double,);",
    "prompt": "get all columns ending with _amount from taxi table"
}'

Python 库示例

pip install ollama

import ollama

r = ollama.generate(
    model='duckdb-nsql:7b-q4_0',
    system='''Here is the database schema that the SQL query will run on:
CREATE TABLE taxi (
    VendorID bigint,
    tpep_pickup_datetime timestamp,
    tpep_dropoff_datetime timestamp,
    passenger_count double,
    trip_distance double,
    fare_amount double,
    extra double,
    tip_amount double,
    tolls_amount double,
    improvement_surcharge double,
    total_amount double,
);''',
    prompt='get all columns ending with _amount from taxi table',
)

print(r['response'])

训练数据

20 万个 DuckDB 文本转 SQL 对，使用 Mixtral-8x7B-Instruct-v0.1 合成生成，并以 DuckDB v0.9.2 文档为指导。以及从 NSText2SQL 转换到 DuckDB SQL 的文本转 SQL 对，使用 sqlglot。

训练过程

DuckDB-NSQL 使用交叉熵损失进行训练，以最大化顺序输入的可能性。 对于文本转 SQL 对的微调，我们仅计算该对的 SQL 部分的损失。 该模型使用 80GB A100 进行训练，利用数据和模型并行性。 我们进行了 10 个周期的微调。

预期用途和限制

该模型专为从给定的表 schema 和自然语言提示生成文本转 SQL 任务而设计。 该模型最适合以下定义的提示格式和输出。 与现有的文本转 SQL 模型相比，SQL 生成不限于 SELECT 语句，而是可以生成任何有效的 DuckDB SQL 语句，包括官方 DuckDB 扩展的语句。

参考文献

Hugging Face