prem-1B-SQL

Prem-1B-SQL项目介绍

Prem-1B-SQL是由Prem AI开发的一款先进的本地化文本到SQL模型。它在数据安全方面秉持“本地优先”的理念，避免了将数据库暴露给第三方封闭源模型的潜在风险。这款模型拥有10亿个参数，适合在低GPU甚至CPU设备上运行，并支持量化处理。未来，官方计划持续优化其表现，这一切使得Prem-1B-SQL在数据分析领域展现出强大的应用潜力。

项目背景

Prem AI坚信，AI辅助的数据分析应当以本地化为核心，因为将数据库暴露给第三方封闭源模型可能导致数据安全问题。基于此理念，他们推出了Prem-1B-SQL，一款适合在本地设备上高效运行的文本到SQL模型。

项目成果

Prem-1B-SQL在两个常用的数据集上进行了评估：BirdBench和Spider。具体结果如下：

其中，BirdBench数据集根据难度分为三个等级，私密测试结果如下：

此外，通过与其他流行的模型进行对比，可以看出Prem-1B-SQL在保持紧凑的参数量的同时，能达到相当不错的测试分数：

如何使用Prem-1B-SQL

Prem-1B-SQL构建在transformers库上，因此可以直接与transformers库结合使用。也正因为如此，它与数据库中输入提示的形成紧密关联。为了简化使用，Prem AI开发了PremSQL库，这是一款完全开源的软件库，它具有以下特性：

• 本地优先：避免使用第三方封闭来源提供商，确保数据安全。
• 可定制数据集：能够创建、微调和评估内置或自定义的数据集。
• 坚固的执行器和评估器：轻松连接数据库并评估模型性能。
• 高级生成器：将自然语言提示转换为可执行的SQL查询。
• 错误处理和自我纠正：在推理过程中自动纠正SQL查询。
• 微调支持：支持通过LoRA、QLoRA或全微调策略进行模型微调。
• 端到端管道：无缝集成所有组件，实现自主数据分析。

Prem-1B-SQL可以通过PremSQL管道或生成器进行部署。以下是一些使用示例：

通过PremSQL管道运行Prem-1B-SQL

最简单的方式就是通过PremSQL管道，只需提供数据库路径或连接URI即可。以下是示例代码：

from premsql.pipelines import SimpleText2SQLAgent
from premsql.generators import Text2SQLGeneratorHF
from premsql.executors import SQLiteExecutor

dsn_or_db_path = "./data/db/california_schools.sqlite"

agent = SimpleText2SQLAgent(
    dsn_or_db_path=dsn_or_db_path,
    generator=Text2SQLGeneratorHF(
        model_or_name_or_path="premai-io/prem-1B-SQL",
        experiment_name="simple_pipeline",
        device="cuda:0",
        type="test"
    ),
)

question = "please list the phone numbers of the direct charter-funded schools that are opened after 2000/1/1"

response = agent.query(question)
response["table"]

通过PremSQL生成器运行Prem-1B-SQL

当需要在某个数据集上进行批量生成时，可以使用PremSQL生成器：

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

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)

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

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 
)

模型训练使用的数据集

Prem-1B-SQL的训练使用了以下数据集：

1. BirdBench训练数据集：可在此处找到。
2. Spider数据集：可在此处找到。
3. 域专属数据集：由PremAI收集并上传。
4. Gretel AI合成数据集：可在此处找到。

此外，开发团队通过结合上述数据集创建了错误处理数据集，将模型的学习扩展到自我纠正能力上。

未来，Prem AI计划在公共基准上发布更多的评估结果，以展示Prem-1B-SQL的能力。有关更多的详细信息，请参阅文档。