【翻译】OmniSQL——大规模生成高质量文本转SQL数据

OmniSQL——大规模生成高质量文本转SQL数据

简介

我们提出了一套自动化、可扩展的文本转SQL数据合成框架（如下图所示），并基于该框架发布了首个百万量级文本转SQL数据集SynSQL-2.5M。该数据集包含超过250万条多样化的高质量数据样本，覆盖16,000余个跨领域数据库。

在SynSQL-2.5M基础上，我们推出了OmniSQL系列模型（7B/14B/32B三个版本）。微调过程中还融合了Spider和BIRD提供的人类标注高质量数据。

论文链接: https://arxiv.org/abs/2503.02240
GitHub链接:https://github.com/RUCKBReasoning/OmniSQL

更新(2025.3.11): 已在train_and_evaluate文件夹上传训练与评估脚本，可轻松复现实验结果
更新(2025.3.12): 在data_synthesis文件夹开源数据合成框架源码

下载资源

模型与数据集	下载地址
SynSQL-2.5M	✨ Modelscope / 🤗 HuggingFace
OmniSQL-7B	✨ Modelscope / 🤗 HuggingFace
OmniSQL-14B	✨ Modelscope / 🤗 HuggingFace
OmniSQL-32B	✨ Modelscope / 🤗 HuggingFace

SynSQL-2.5M数据集特征

本数据集完全由开源大模型生成，遵循Apache 2.0协议，具有以下特点：

• 规模庞大：包含2,544,390个<数据库、问题、SQL查询、思维链解答>四元组
• 领域广泛：覆盖16,583个现实场景的合成数据库
• 复杂度分层：从单表查询到多表连接/函数/公共表表达式，涵盖简单、中等、复杂、极复杂四个层级
• 语言风格多样：问题表述包含正式、口语、命令式、疑问式、描述性、简略、模糊、隐喻、对话等形态
• 全量思维链：所有样本均附带推理过程

截至2025年3月，SynSQL-2.5M是当前规模最大、多样性最丰富的合成文本转SQL数据集。我们鼓励研究者利用该数据集开展模型研发，如果觉得有用请为我们点亮星标或引用论文，您的反馈是我们持续优化的动力。

性能评估

我们在Spider、BIRD等标准基准，以及Spider2.0-SQLite、ScienceBenchmark、EHRSQL等专业领域基准和三个鲁棒性测试集（Spider-DK/Spider-Syn/Spider-Realistic）上进行了全面评估：

（注："Gre"表示贪婪解码，"Maj"为8次多数投票。Spider(dev)/Spider-Syn/Spider-Realistic采用测试套件准确率(TS)，其余采用执行准确率(EX)）

OmniSQL在同等规模基线上显著领先，甚至在多个数据集上超越GPT-4o和DeepSeek-V3。这些成绩由单一模型实现，未引入模式链接/SQL修正等额外设计，集成这些技术后精度有望进一步提升。

快速入门

提示词模板

input_prompt_template = '''任务概述：  
您是一名数据科学专家。请根据提供的数据库模式和自然语言问题，生成有效的SQL查询语句。  

数据库引擎：SQLite  

数据库模式：  
{db_details}  
（包含表结构、主外键关系等约束条件）  

问题：  
{question}  

注意事项：  
- 仅输出问题要求的列  
- 确保查询结果完整匹配问题需求  
- 生成SQL前请逐步思考  

输出格式：  
```sql  
-- 您的SQL查询  

请深呼吸并逐步思考…‘’’

（"db_details"需替换为CREATE TABLE格式的DDL语句，"question"可拼接外部知识。当前仅支持SQLite语法）  

#### vLLM调用示例  
```python  
from vllm import LLM, SamplingParams  
# 初始化模型/参数后调用  
chat_prompt = tokenizer.apply_chat_template(  
    [{"role": "user", "content": prompt}],  
    add_generation_prompt = True, tokenize = False  
)  
outputs = llm.generate([chat_prompt], sampling_params)  

Transformers调用示例

model = AutoModelForCausalLM.from_pretrained("seeklhy/OmniSQL-7B", torch_dtype=torch.bfloat16).to("cuda:0")  
output_ids = model.generate(**inputs, max_new_tokens=2048)  

局限性

当前数据集仅支持英语和SQLite引擎，在多语言/多SQL方言场景可能存在局限。但用户可通过我们的框架合成新数据，并以OmniSQL为基座进行领域微调。

数据合成框架

本目录包含数据合成框架的源码及提示词模板

四步合成流程

1. 网络表格驱动的数据库合成（见database_synthesis）
2. 复杂度感知的SQL查询生成（见sql_synthesis）
3. 风格化自然语言问题生成（见question_synthesis）
4. 思维链解决方案生成（见cot_synthesis）

流程虽为顺序执行，但可从任意中间步骤启动合成。例如：若已有现成数据库，可跳过步骤1，直接为您的数据库生成<问题，SQL查询，思维链解答>三元组。

环境配置（Anaconda）

conda create -n omnisql_data_synthesis python=3.9.5  
conda activate omnisql_data_synthesis  

pip install -U sentence-transformers  
pip install json-repair ijson matplotlib func_timeout  

OmniSQL训练与评估指南

环境配置

所有实验均在以下环境运行：

• Anaconda 3
• Python 3.9.5
• 8张NVIDIA A800 80GB显卡

注：推理和评估仅需单张A800 80GB显卡，完整训练OmniSQL建议使用8张A800 80GB显卡

---

数据集准备

下载资源

从以下平台获取数据集：

• ModelScope-OmniSQL数据集
• HuggingFace-OmniSQL数据集

包含BIRD、Spider、ScienceBenchmark、EHRSQL、Spider2-SQLite、Spider-DK、Spider-Realistic、Spider-Syn及SynSQL-2.5M。解压本目录下的data.zip即可。

数据预处理

预处理后的数据集已包含在data.zip（参见*.json文件）。如需重新预处理：

1. 配置环境：

conda create -n omnisql_process_data python=3.9.5
conda activate omnisql_process_data

apt-get update
apt-get install -y openjdk-11-jdk

pip3 install func_timeout ijson pyserini==0.22.1 faiss-cpu torch==2.1.0 numpy==1.24.3 nltk==3.8.1
python3 nltk_downloader.py

2. 运行预处理脚本：

# 构建数据库内容的BM25索引
python3 build_contents_index.py
# 生成输入-输出序列
sh process_dataset.sh

注意：SynSQL-2.5M因规模庞大（约250万样本），预处理可能需要超过24小时

---

复现评估结果

环境配置

conda create -n omnisql_eval python=3.9.5
conda activate omnisql_eval
pip3 install vllm==0.6.3.post1 func_timeout tqdm matplotlib nltk==3.8.1 sqlparse
python3 nltk_downloader.py

下载评估材料

从test_suite_sql_eval.zip获取Spider测试套件数据库和评估脚本，解压至当前目录

执行评估

python3 eval_open_source_models.py

• 预测的SQL语句保存在results文件夹
• 评估结果（如模型准确率）存储在evaluation_results文件夹

---

完整训练OmniSQL

环境配置

conda create -n omnisql_train python=3.9.5
conda activate omnisql_train
pip3 install torch==2.1.0 torchvision==0.16.0 torchaudio==2.1.0 transformers==4.45.1 accelerate==0.34.2 deepspeed==0.10.3 numpy==1.24.3 peft datasets tensorboard ijson

安装Flash-Attention加速

建议下载预编译的Wheel文件（根据实际环境选择）：

pip3 install flash_attn-2.5.8+cu{cuda版本}torch{torch版本}cxx11abiFALSE-cp{python版本}-cp{python版本}-linux_x86_64.whl

例如CUDA 12.2 + PyTorch 2.1 + Python 3.9.5需安装：
flash_attn-2.5.8+cu122torch2.1cxx11abiFALSE-cp39-cp39-linux_x86_64.whl

启动训练

# 训练OmniSQL-7B
sh train_omnisql_7b.sh
# 训练OmniSQL-14B  
sh train_omnisql_14b.sh
# 训练OmniSQL-32B
sh train_omnisql_32b.sh

关键说明：

1. 完整版训练需手动合并三个训练集（train_synsql.json、train_bird.json、train_spider.json）并更新脚本中的DATASET_DIR
2. OmniSQL-32B训练后需运行merge_lora_adapter.py合并LoRA适配器
3. 训练耗时参考（8xA800 80GB）：
   • 7B模型约6天
   • 14B模型约12天
   • 32B模型约20天
     建议：若无特殊需求，可直接使用我们开源的预训练模型，或基于OmniSQL进行小规模增量训练