Prem 1B SQL

🚀 Prem-1B-SQL (HuggingFace)

Prem-1B-SQLは、Prem AIによって開発された完全にローカルなText-to-SQLモデルの最初のシリーズの1つです。10億個のパラメータを持つモデルであり、低グレードのGPUデバイス（量子化するとCPUデバイスでも）に簡単に収まります。AIによるデータ分析は、まずローカルで行うべきだと考えています。なぜなら、データベースをサードパーティのクローズドソースモデルに公開すると、データセキュリティ上の問題につながる可能性があるからです。このモデルのいくつかの公開ベンチマーク結果を近日中に公開する予定です。また、より良い結果を得るためにこのモデルを改善していきます。

• ブログ記事はこちらで読めます。
• PremSQLライブラリ | GitHub
• 開発元: Prem AI
• ライセンス: [MIT]

🚀 クイックスタート

モデルの評価結果

私たちは、2つの人気のあるベンチマークデータセット、BirdBenchとSpiderでモデルを評価しました。BirdBenchは、公開検証データセット（1534のデータポイント）と非公開テストデータセットで構成されています。Spiderは、公開検証データセットのみを提供します。以下はその結果です。

BirdBenchデータセットは、さまざまな難易度レベルに分けられています。以下は、異なる難易度レベルでの非公開結果の詳細です。

以下は、人気のあるクローズドソースモデルとオープンソースモデルのより詳細な比較です。

📦 インストール

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を使用してモデルを実行することもできます。これは、あるデータセットに対して大量の生成を行いたい場合に便利です。以下はその例です。

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をDBに対して実行し、失敗した場合はエラーメッセージを使用して修正し、有効な結果が得られるか、リトライが上限に達するまで繰り返します。

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)
)

📚 ドキュメント

Prem-1B-SQLは、以下のデータセットを使用して学習されています。

1. BirdBench学習データセット | HFのPremSQLデータセットにアップロード
2. Spiderデータセット | HFのPremSQLデータセットにアップロード
3. ドメイン特化データセット、収集されてPremSQLデータセットにアップロード
4. Gretel AI合成データセット

さらに、これらのデータセットを基にエラーハンドリングデータセットを作成し、モデルがエラーから学習して自己修正できるようにしています。

Prem-1B-SQLのいくつかの公開ベンチマークでの結果は、近日中に公開されます。

📄 ライセンス

このプロジェクトは、[MIT]ライセンスの下で公開されています。