duckdb-nsql

由 MotherDuck 和 Numbers Station 製作的 70 億參數 text-to-SQL 模型。

7b

2.85 萬 12 個月前

17 個標籤

讀我檔案

duckdb-nsql model

DuckDB-NSQL 是一個 70 億參數的 text-to-SQL 模型，專為 SQL 生成任務設計。

此模型基於 Meta 的原始 Llama-2 7B 模型，並在通用 SQL 查詢資料集上進一步預訓練，然後在由 DuckDB text-to-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 text-to-SQL 配對，使用 Mixtral-8x7B-Instruct-v0.1 合成生成，並以 DuckDB v0.9.2 文件為指導。以及來自 NSText2SQL 的 text-to-SQL 配對，這些配對已使用 sqlglot 轉譯為 DuckDB SQL。

訓練程序

DuckDB-NSQL 使用交叉熵損失進行訓練，以最大化序列輸入的可能性。對於 text-to-SQL 配對的微調，我們僅計算配對中 SQL 部分的損失。該模型使用 80GB A100 進行訓練，利用資料和模型並行性。我們微調了 10 個 epoch。

預期用途和限制

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

參考文獻

Hugging Face