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] 許可證。