Prem 1B SQL

🚀 Prem-1B-SQL (HuggingFace)

Prem-1B-SQL 是由 Prem AI 开发的首批完全本地化的文本转 SQL 模型系列之一。作为一个拥有 10 亿参数的模型，它可以轻松适配低配置 GPU 设备（量化后也可在 CPU 设备上运行）。我们认为，人工智能辅助数据分析应优先考虑本地化。因为将数据库暴露给第三方闭源模型可能会导致数据安全漏洞。我们很快将公布该模型的一些公开基准测试结果，也会对该模型进行迭代优化以获得更好的效果。

• 阅读博客文章 点击此处

• PremSQL 库 | GitHub

• 开发者： Prem AI

• 许可证： [MIT]

✨ 主要特性

• 完全本地化：避免使用第三方闭源模型，保障数据安全。
• 低资源需求：10 亿参数模型，适合低配置 GPU 设备，量化后可在 CPU 设备运行。
• 持续优化：将公布公开基准测试结果，并对模型进行迭代优化。

📦 安装指南

要安装 PremSQL，只需创建一个新环境并输入以下命令：

pip install -U premsql

💻 使用示例

基础用法

使用 PremSQL BaseLine Agent 运行 Prem-1B-SQL：

from premsql.agents import BaseLineAgent
from premsql.generators import Text2SQLGeneratorOllama
from premsql.agents.tools import SimpleMatplotlibTool
from premsql.executors import SQLiteExecutor

text2_sqlmodel = Text2SQLGeneratorHF(
    model_or_name_or_path="premai-io/prem-1B-SQL",
    experiment_name="test_generators",
    device="cuda:0",
    type="test"
)

analyser_and_plotter = Text2SQLGeneratorHF(
    model_or_name_or_path="meta-llama/Llama-3.2-1B-Instruct",
    experiment_name="test_generators",
    device="cuda:0",
    type="test"
)

agent = BaseLineAgent(
    session_name="testing_hf",
    db_connection_uri="sqlite:////path/to/your/database.sqlite",
    specialized_model1=model,
    specialized_model2=model,
    plot_tool=SimpleMatplotlibTool(),
    executor=SQLiteExecutor()
)

response = agent(
    "/query what all tables are present inside the database"
)
response.show_dataframe()

在底层，它会自动连接到你的数据库，并为你完成所有繁重的工作，如创建提示、执行查询等。

高级用法

使用 PremSQL Generators 运行 Prem-1B-SQL：

from premsql.generators import Text2SQLGeneratorHF
from premsql.datasets import Text2SQLDataset

# Define a dataset
dataset = bird_dataset = Text2SQLDataset(
    dataset_name='bird', split="validation", force_download=False,
    dataset_folder="/path/to/dataset"
).setup_dataset(num_rows=10, num_fewshot=3)

# Define a generator
generator = Text2SQLGeneratorHF(
    model_or_name_or_path="premai-io/prem-1B-SQL",
    experiment_name="test_generators",
    device="cuda:0",
    type="test"
)

# Generate on the full dataset
responses = generator.generate_and_save_results(
    dataset=bird_dataset,
    temperature=0.1,
    max_new_tokens=256
)

print(responses)

这在你想对某个数据集进行批量生成时非常有用。

使用执行引导解码

这种策略会将生成的 SQL 在数据库上执行，如果执行失败，则使用错误消息进行修正，重复此过程直到得到有效结果或达到最大重试次数。

from premsql.executors import SQLiteExecutor

executor = SQLiteExecutor()
response = generator.generate_and_save_results(
    dataset=bird_dataset,
    temperature=0.1,
    max_new_tokens=256,
    force=True,
    executor=executor,
    max_retries=5 # this is optional (default is already set to 5)
)

📚 详细文档

你还可以使用 HuggingFace Transformers 和 PremSQL Tuners 对 Prem-1B-SQL 进行微调。请 查看我们的文档 以了解更多关于 PremSQL 的信息和我们提供的所有功能。

🔧 技术细节

评估结果

我们在两个流行的基准数据集上对模型进行了评估：BirdBench 和 Spider。BirdBench 包含一个公共验证数据集（有 1534 个数据点）和一个私有测试数据集，Spider 只有一个公共验证数据集。以下是评估结果：

BirdBench 数据集分布在不同的难度级别上。以下是不同难度级别的私有测试结果详情：

以下是与流行的闭源和开源模型的更详细比较：

训练数据集

Prem-1B-SQL 使用以下数据集进行训练：

1. BirdBench 训练数据集 | 上传至 HF 上的 PremSQL 数据集
2. Spider 数据集 | 上传至 HF 上的 PremSQL 数据集
3. 领域特定数据集，收集并上传至 PremSQL 数据集
4. Gretel AI 合成数据集

此外，我们在这些数据集的基础上创建了错误处理数据集，以使模型能够从错误中学习并自我修正。

未来计划

Prem-1B-SQL 在一些公共基准测试中的结果将很快公布。

📄 许可证

本项目采用 [MIT] 许可证。