Duckdb NSQL 7B V0.1

🚀 DuckDB-NSQL-7B

NSQLは、SQL生成タスクに特化して設計された自己回帰型のオープンソース大規模基礎モデル（FM）のファミリーです。

このリポジトリでは、NSQLの新しいメンバーであるDuckDB-NSQLを紹介します。このモデルはMetaのオリジナルLlama-2 7Bモデルをベースに、一般的なSQLクエリのデータセットで事前学習され、その後DuckDBのテキストからSQLへのペアで構成されるデータセットで微調整されています。

🚀 クイックスタート

モデルの概要

NSQLは、SQL生成タスクに特化した自己回帰型のオープンソース大規模基礎モデル（FM）のファミリーです。このリポジトリでは、NSQLの新しいメンバーであるDuckDB-NSQLを紹介します。このモデルはMetaのオリジナルLlama-2 7Bモデルをベースに、一般的なSQLクエリのデータセットで事前学習され、その後DuckDBのテキストからSQLへのペアで構成されるデータセットで微調整されています。

学習データ

200kのDuckDBのテキストからSQLへのペアです。これはMixtral-8x7B-Instruct-v0.1を使用して合成生成され、DuckDB v0.9.2のドキュメントに沿って作成されています。また、NSText2SQLからのテキストからSQLへのペアがsqlglotを使用してDuckDB SQLに変換されています。

評価データ

モデルは、75のテキストからSQLへのペアを含むDuckDB固有のベンチマークで評価されています。このベンチマークはこちらで入手できます。

学習手順

DuckDB-NSQLは、クロスエントロピー損失を使用して学習され、逐次入力の尤度を最大化しています。テキストからSQLへのペアでの微調整では、ペアのSQL部分に対してのみ損失を計算します。モデルは80GBのA100を使用して学習され、データ並列とモデル並列を活用しています。10エポックで微調整されています。

想定される使用方法と制限

このモデルは、与えられたテーブルスキーマと自然言語のプロンプトからテキストからSQLへの生成タスクに設計されています。以下に定義されたプロンプト形式と出力で最適に動作します。既存のテキストからSQLへのモデルとは異なり、SQL生成はSELECT文に限定されず、任意の有効なDuckDB SQL文を生成でき、公式のDuckDB拡張機能の文も含みます。

✨ 主な機能

• 一般的なSQLクエリの事前学習とDuckDBのテキストからSQLへのペアでの微調整
• 任意の有効なDuckDB SQL文の生成（SELECT文に限定されない）

📦 インストール

このREADMEには具体的なインストール手順が記載されていないため、このセクションは省略されます。

💻 使用例

基本的な使用法

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM
tokenizer = AutoTokenizer.from_pretrained("motherduckdb/DuckDB-NSQL-7B-v0.1")
model = AutoModelForCausalLM.from_pretrained("motherduckdb/DuckDB-NSQL-7B-v0.1", torch_dtype=torch.bfloat16)

text = """### Instruction:
Your task is to generate valid duckdb SQL to answer the following question.

### Input:

### Question:
create a new table called tmp from test.csv

### Response (use duckdb shorthand if possible):
"""

input_ids = tokenizer(text, return_tensors="pt").input_ids

generated_ids = model.generate(input_ids, max_length=500)
print(tokenizer.decode(generated_ids[0], skip_special_tokens=True))

高度な使用法

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM
tokenizer = AutoTokenizer.from_pretrained("motherduckdb/DuckDB-NSQL-7B-v0.1")
model = AutoModelForCausalLM.from_pretrained("motherduckdb/DuckDB-NSQL-7B-v0.1", torch_dtype=torch.bfloat16)

text = """### Instruction:
Your task is to generate valid duckdb SQL to answer the following question, given a duckdb database schema.

### Input:
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,
);

### Question:
get all columns ending with _amount from taxi table

### Response (use duckdb shorthand if possible):"""

input_ids = tokenizer(text, return_tensors="pt").input_ids

generated_ids = model.generate(input_ids, max_length=500)
print(tokenizer.decode(generated_ids[0], skip_special_tokens=True))

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM
tokenizer = AutoTokenizer.from_pretrained("motherduckdb/DuckDB-NSQL-7B-v0.1")
model = AutoModelForCausalLM.from_pretrained("motherduckdb/DuckDB-NSQL-7B-v0.1", torch_dtype=torch.bfloat16)

text = """### Instruction:
Your task is to generate valid duckdb SQL to answer the following question, given a duckdb database schema.

### Input:
Here is the database schema that the SQL query will run on:
CREATE TABLE rideshare (
    hvfhs_license_num varchar,
    dispatching_base_num varchar,
    originating_base_num varchar,
    request_datetime timestamp,
    on_scene_datetime timestamp,
    pickup_datetime timestamp,
    dropoff_datetime timestamp,
    trip_miles double,
    trip_time bigint,

);

### Question:
get longest trip in december 2022

### Response (use duckdb shorthand if possible):
"""

input_ids = tokenizer(text, return_tensors="pt").input_ids

generated_ids = model.generate(input_ids, max_length=500)
print(tokenizer.decode(generated_ids[0], skip_special_tokens=True))

詳細情報（例えば、ローカルデータベースでの実行）については、このリポジトリの例を参照してください。

📚 ドキュメント

このセクションには、モデルの説明、学習データ、評価データ、学習手順、想定される使用方法と制限、使用例などの詳細情報が含まれています。

🔧 技術詳細

このREADMEには具体的な技術詳細が記載されていないため、このセクションは省略されます。

📄 ライセンス

このモデルはLlama2のライセンスの下で提供されています。