duckdb-nsql

MotherDuck 和 Numbers Station 製作的 70 億參數文字轉 SQL 模型。

7b

28.5K 12 個月前

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 個 epoch。

預期用途與限制

該模型專為從給定的表格 schema 和自然語言提示中生成文字轉 SQL 的任務而設計。該模型最適合使用以下定義的提示格式和輸出。與現有的文字轉 SQL 模型相比，SQL 生成不限於 SELECT 語句，而是可以生成任何有效的 DuckDB SQL 語句，包括用於官方 DuckDB 擴展的語句。

參考文獻

Hugging Face