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 |