Rain SQLCoder

🚀 [Rain's SQLCoder]

Rain's SQLCoder 是一款先進的大語言模型（LLM），專為自然語言到 SparkSQL 的生成而設計。它基於 Qwen2.5-Coder-32B-Instruct 微調而來，擁有 320 億參數。該模型針對自然語言到 SparkSQL 的轉換任務進行了優化，能夠有效處理長達 32k 標記的上下文，尤其適合生成複雜和大規模的 SQL 查詢。

🤗 Hugging Face | 🖥️ Demo | 💬 微信 | 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 生成設計**：Rain's SQLCoder 是一款先進的大語言模型，專注於將自然語言轉換為 SparkSQL 查詢。
- **基於強大的基礎模型微調**：它基於 [Qwen2.5-Coder-32B-Instruct](https://huggingface.co/Qwen/Qwen2.5-Coder-32B-Instruct) 進行微調，擁有 320 億參數。
- **處理長上下文能力**：能夠有效處理長達 32k 標記的上下文，適合生成複雜和大規模的 SQL 查詢。

## 📚 詳細文檔
### 提示模板
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. 在所有其他情況下，預測被認為是不正確的。

### 實驗結果
我們在兩個測試數據集上比較了 Rain's SQLCoder 與國內外最先進的自然語言大模型的生成準確率。基準數據集包含基線樣本，而增強數據集是通過對基準數據集的 20% 進行分層抽樣，並補充相關的用戶問題和相應的 SparkSQL 語句構建的，以評估模型在增強上下文信息下的性能。實驗結果表明，Rain's SQLCoder 在查詢意圖理解、SQL 語法準確性和複雜查詢處理方面具有顯著優勢。

#### 基準數據集
![基準數據集結果](./figures/benchmark_dataset_result.png)

#### 增強數據集
![增強數據集結果](./figures/enhanced_dataset_result.png)

## 📄 許可證
本項目採用 Apache-2.0 許可證。

| 屬性 | 詳情 |
|------|------|
| 模型類型 | 自然語言到 SparkSQL 生成的大語言模型 |
| 基礎模型 | Qwen/Qwen2.5-Coder-32B-Instruct |
| 管道標籤 | 文本生成 |
| 庫名稱 | transformers |
| 許可證 | Apache-2.0 |