Nsql 6B

🚀 NSQL (NSQL-6B)

NSQL是专门为SQL生成任务设计的自回归开源大型基础模型（FMs）系列。本模型能根据给定的表结构和自然语言提示生成SQL查询，为SQL相关任务提供了高效的解决方案。

🚀 快速开始

模型推理参数

使用示例

以下是一些使用NSQL模型进行SQL生成的示例：

示例1：计算体育场总数

from transformers import AutoTokenizer, AutoModelForCausalLM
tokenizer = AutoTokenizer.from_pretrained("NumbersStation/nsql-6B")
model = AutoModelForCausalLM.from_pretrained("NumbersStation/nsql-6B")

text = """CREATE TABLE stadium (
    stadium_id number,
    location text,
    name text,
    capacity number,
    highest number,
    lowest number,
    average number
)

CREATE TABLE singer (
    singer_id number,
    name text,
    country text,
    song_name text,
    song_release_year text,
    age number,
    is_male others
)

CREATE TABLE concert (
    concert_id number,
    concert_name text,
    theme text,
    stadium_id text,
    year text
)

CREATE TABLE singer_in_concert (
    concert_id number,
    singer_id text
)

-- Using valid SQLite, answer the following questions for the tables provided above.

-- What is the maximum, the average, and the minimum capacity of stadiums ?

SELECT"""

input_ids = tokenizer(text, return_tensors="pt").input_ids

generated_ids = model.generate(input_ids, max_length=500)
print(tokenizer.decode(generated_ids[0], skip_special_tokens=True))

示例2：计算体育场总数（简化表结构）

from transformers import AutoTokenizer, AutoModelForCausalLM
tokenizer = AutoTokenizer.from_pretrained("NumbersStation/nsql-6B")
model = AutoModelForCausalLM.from_pretrained("NumbersStation/nsql-6B")

text = """CREATE TABLE stadium (
    stadium_id number,
    location text,
    name text,
    capacity number,
)

-- Using valid SQLite, answer the following questions for the tables provided above.

-- how many stadiums in total?

SELECT"""

input_ids = tokenizer(text, return_tensors="pt").input_ids

generated_ids = model.generate(input_ids, max_length=500)
print(tokenizer.decode(generated_ids[0], skip_special_tokens=True))

示例3：计算未完成的工作订单数量

from transformers import AutoTokenizer, AutoModelForCausalLM
tokenizer = AutoTokenizer.from_pretrained("NumbersStation/nsql-6B")
model = AutoModelForCausalLM.from_pretrained("NumbersStation/nsql-6B")

text = """CREATE TABLE work_orders (
    ID NUMBER,
    CREATED_AT TEXT,
    COST FLOAT,
    INVOICE_AMOUNT FLOAT,
    IS_DUE BOOLEAN,
    IS_OPEN BOOLEAN,
    IS_OVERDUE BOOLEAN,
    COUNTRY_NAME TEXT,
)

-- Using valid SQLite, answer the following questions for the tables provided above.

-- how many work orders are open?

SELECT"""

input_ids = tokenizer(text, return_tensors="pt").input_ids

generated_ids = model.generate(input_ids, max_length=500)
print(tokenizer.decode(generated_ids[0], skip_special_tokens=True))

更多信息（例如，在本地数据库上运行），请参考此仓库中的示例。

✨ 主要特性

• 专为SQL生成设计：NSQL是专门为SQL生成任务打造的自回归开源大型基础模型，能够根据自然语言提示准确生成SQL查询。
• 基于预训练和微调：该模型基于Salesforce的CodeGen - Multi 6B进行预训练，并在通用SQL查询数据集上进一步预训练，然后在文本到SQL对的数据集上进行微调。

📦 安装指南

文档未提及安装步骤，故跳过此章节。

📚 详细文档

模型描述

NSQL是专门为SQL生成任务设计的自回归开源大型基础模型（FMs）系列。本仓库中的检查点基于Salesforce的CodeGen - Multi 6B，并在通用SQL查询数据集上进一步预训练，然后在由文本到SQL对组成的数据集上进行微调。

训练数据

• 通用SQL查询：来自The Stack的SQL子集，包含100万个训练样本。
• 文本到SQL对：来自网络上20多个公共来源的标准数据集。我们预留了Spider和GeoQuery数据集用于评估。

评估数据

我们在两个文本到SQL基准测试上评估模型：Spider和GeoQuery。

训练过程

NSQL使用交叉熵损失进行训练，以最大化序列输入的似然性。对于文本到SQL对的微调，我们仅计算对中SQL部分的损失。该系列模型使用80GB A100 GPU进行训练，利用数据和模型并行性。我们进行了3个周期的预训练和10个周期的微调。

预期用途和限制

该模型专为根据给定的表结构和自然语言提示进行文本到SQL生成任务而设计。该模型在以下定义的提示格式下效果最佳，并输出SELECT查询。

🔧 技术细节

NSQL模型的训练基于交叉熵损失，通过在通用SQL查询数据集上预训练和在文本到SQL对数据集上微调，不断优化模型的性能。在训练过程中，利用了数据和模型并行性，使用80GB A100 GPU进行高效训练。预训练和微调的周期设置分别为3个周期和10个周期，以确保模型能够学习到足够的知识来完成SQL生成任务。

📄 许可证

本项目采用BSD 3 - 条款许可证（BSD-3-Clause）。