duckdb-nsql

7B 參數的 text-to-SQL 模型，由 MotherDuck 和 Numbers Station 製作。

7b

28.5K 12 months ago

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

設定系統提示

此模型預期系統提示中會輸入結構描述

/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,
);"""

一旦在系統提示中提供結構描述，模型將在後續回應中使用它。

對於以下提示

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 轉譯為 DuckDB SQL 的 text-to-SQL 配對組（使用 sqlglot）。

訓練程序

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

預期用途與限制

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

參考文獻

Hugging Face