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为基座进行领域微调。
数据合成框架
本目录包含数据合成框架的源码及提示词模板
四步合成流程
- 网络表格驱动的数据库合成（见
database_synthesis） - 复杂度感知的SQL查询生成（见
sql_synthesis） - 风格化自然语言问题生成（见
question_synthesis） - 思维链解决方案生成（见
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文件）。如需重新预处理：
- 配置环境：
conda create -n omnisql_process_data python=3.9.5
conda activate omnisql_process_dataapt-get update
apt-get install -y openjdk-11-jdkpip3 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
- 运行预处理脚本：
# 构建数据库内容的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
关键说明：
- 完整版训练需手动合并三个训练集（
train_synsql.json、train_bird.json、train_spider.json）并更新脚本中的DATASET_DIR - OmniSQL-32B训练后需运行
merge_lora_adapter.py合并LoRA适配器 - 训练耗时参考（8xA800 80GB）：
- 7B模型约6天
- 14B模型约12天
- 32B模型约20天
建议：若无特殊需求，可直接使用我们开源的预训练模型，或基于OmniSQL进行小规模增量训练
