Rain SQLCoder

🚀 レインのSQLCoder

Rain's SQLCoderは、自然言語からSparkSQLを生成するために設計された最先端の大規模言語モデルです。このモデルは、自然言語からSparkSQLへの変換タスクに最適化されており、最大32kトークンのコンテキストを効果的に処理でき、複雑で大規模なSQLクエリの生成に特に適しています。

🤗 Hugging Face | 🖥️ デモ | 💬 WeChat | GitHub

English | 中文

🚀 クイックスタート

ここでは、モデルの読み込みと使用方法をすぐに学べるように、サンプルコードを提供します。

💡 使用ヒント

Rain's SQLCoderはSELECT文の生成のみを目的として学習されています。また、テーブルスキーマがユーザーの質問に答えるのに不十分な場合、モデルは応答を拒否します。

基本的な使用法

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from utils.prompt import SQLGeneratePrompt

model_name = "SuanChang/rain-SQLCoder"

model = AutoModelForCausalLM.from_pretrained(
    model_name,
    torch_dtype=torch.bfloat16,
    device_map="auto",
)
tokenizer = AutoTokenizer.from_pretrained(model_name)

question = "What is the name of the department that offers a course that has a description including the word 'Statistics'?"
schemas = [
'''CREATE TABLE `course` (
    `crs_code` STRING,
    `dept_code` STRING,
    `crs_description` STRING,
    `crs_credit` DOUBLE
);''',
'''CREATE TABLE `department` (
    `dept_code` STRING,
    `dept_name` STRING,
    `school_code` STRING,
    `emp_num` INT,
    `dept_address` STRING,
    `dept_extension` INT
);''',
'''CREATE TABLE `student` (
    `stu_num` INT,
    `stu_lname` STRING,
    `stu_fname` STRING,
    `stu_init` STRING,
    `stu_dob` STRING,
    `stu_hrs` INT,
    `stu_class` STRING,
    `stu_gpa` DOUBLE,
    `stu_transfer` INT,
    `dept_code` STRING,
    `stu_phone` INT,
    `prof_num` INT
);'''
]
hint = "- Today is 2025-02-01."
data = dict(
    question=question,
    schema="\n\n".join(schemas),
    hint=hint,
    related_question_sqls=None,
)
text, _, _ = SQLGeneratePrompt.prompt(data)

model_inputs = tokenizer([text], return_tensors="pt").to(model.device)

generated_ids = model.generate(
    **model_inputs,
    max_new_tokens=32768
)
generated_ids = [
    output_ids[len(input_ids):] for input_ids, output_ids in zip(model_inputs.input_ids, generated_ids)
]
response = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]

print(response)

'''
```sql
SELECT d.dept_name FROM department d JOIN course c ON d.dept_code = c.dept_code WHERE c.crs_description LIKE '%Statistics%';

'''

## ✨ 主な機能
- 自然言語からSparkSQLへの変換タスクに最適化されています。
- 最大32kトークンのコンテキストを効果的に処理できます。
- 複雑で大規模なSQLクエリの生成に特に適しています。

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

## 📚 ドキュメント

### プロンプト
Rain's SQLCoderは[Alpaca](https://github.com/tatsu-lab/stanford_alpaca)テンプレートを採用しており、以下のプロンプト形式を使用します。

Below is an instruction that describes a task. Write a response that appropriately completes the request.

Instruction:

[BEGIN OF TASK INSTRUCTION] You are an expert in composing Spark SQL queries. You are given a user query and a set of table schemas. Based on the user query, you need to generate one Spark SQL query to achieve the purpose. {task description for date hint and related question and sqls} [END OF TASK INSTRUCTION]

[BEGIN OF TABLE SCHEMAS] {schemas} [END OF TABLE SCHEMAS]

[BEGIN OF GENERATION HINT] {date hint} [END OF GENERATION HINT]

[BEGIN OF RELATED QUERIES] {related question and sqls} [END OF RELATED QUERIES]

[BEGIN OF FORMAT INSTRUCTION] The output MUST strictly adhere to the following format, and NO other text MUST be included.

your output Spark SQL query

[END OF FORMAT INSTRUCTION]

[BEGIN OF QUERY] User Query: {user question} [END OF QUERY]

Response:

### 評価
[SQL-Eval](https://github.com/defog-ai/sql-eval)の評価ロジックに従って、予測結果と正解データを比較します。
1. 予測されたデータフレームが正解データフレームと完全に一致する場合、予測は正しいとみなされます。
2. 正解のSQLにソートロジックが含まれておらず、予測されたデータフレームをソートした後に正解データフレームと一致する場合、予測は正しいとみなされます。
3. 正解データフレームの列が予測されたデータフレームの列のサブセットである場合、予測は正しいとみなされます。
4. それ以外の場合は、予測は誤りとみなされます。

### 実験結果
2つのテストデータセットで、Rain's SQLCoderの生成精度を国内外の最先端の自然言語大規模モデルと比較しました。ベンチマークデータセットにはベースラインサンプルが含まれており、拡張データセットはベンチマークデータセットの20%に層化サンプリングを適用し、関連するユーザー質問と対応するSparkSQL文を補足することで構築され、拡張されたコンテキスト情報下でのモデルのパフォーマンスを評価します。実験結果から、Rain's SQLCoderはクエリ意図の理解、SQL構文の正確性、複雑なクエリの処理において大きな優位性を示しています。

#### ベンチマークデータセット
<img src="./figures/benchmark_dataset_result.png" alt="benchmark" width=800>

#### 拡張データセット
<img src="./figures/enhanced_dataset_result.png" alt="enhanced" width=800>

## 🔧 技術詳細
このREADMEには技術的な詳細に関する具体的な内容（50文字以上）が記載されていないため、このセクションを省略します。

## 📄 ライセンス
このプロジェクトはApache-2.0ライセンスの下で提供されています。