duckdb-nsql:7b im LLM-Testsuite-Report

DuckDB-NSQL ist ein Text-to-SQL-Modell mit 7 Milliarden Parametern, das speziell für Aufgaben der SQL-Generierung entwickelt wurde.

Dieses Modell basiert auf dem ursprünglichen Llama-2-7B-Modell von Meta. Es wurde anschließend auf einem Datensatz allgemeiner SQL-Abfragen weiter vortrainiert und danach auf einem Datensatz aus DuckDB-Text-to-SQL-Paaren feinabgestimmt.

Verwendung

Beispielprompt

Gegeben sei dieses Schema: CREATE TABLE taxi ( VendorID bigint, tpep_pickup_datetime timestamp, tpep_dropoff_datetime timestamp, passenger_count double, trip_distance double, fare_amount double, extra double, tip_amount double, tolls_amount double, improvement_surcharge double, total_amount double, ); Gib mir Taxis mit mehr als 2 Fahrgästen

Beispielausgabe

SELECT * FROM taxi WHERE passenger_count > 2

Systemprompt festlegen

Dieses Modell erwartet das Schema als Eingabe im Systemprompt:

/set system """Hier ist das Datenbankschema, auf dem die SQL-Abfrage ausgeführt wird: CREATE TABLE taxi ( VendorID bigint, tpep_pickup_datetime timestamp, tpep_dropoff_datetime timestamp, passenger_count double, trip_distance double, fare_amount double, extra double, tip_amount double, tolls_amount double, improvement_surcharge double, total_amount double, );"""

Sobald das Schema im Systemprompt bereitgestellt wurde, verwendet das Modell es in den folgenden Antworten.

Für den folgenden Prompt:

gib alle Spalten aus der Tabelle taxi zurück, die auf _amount enden

Das Modell gibt etwa Folgendes aus:

SELECT COLUMNS('.*_amount') FROM taxi;

API-Beispiel

$ curl http://localhost:11434/api/generate -d '{ "model": "duckdb-nsql:7b-q4_0", "system": "Hier ist das Datenbankschema, auf dem die SQL-Abfrage ausgeführt wird: CREATE TABLE taxi (VendorID bigint, tpep_pickup_datetime timestamp, tpep_dropoff_datetime timestamp, passenger_count double, trip_distance double, fare_amount double, extra double, tip_amount double, tolls_amount double, improvement_surcharge double, total_amount double,);", "prompt": "gib alle Spalten aus der Tabelle taxi zurück, die auf _amount enden" }'

Beispiel für die Python-Bibliothek

pip install ollama

import ollama r = ollama.generate( model='duckdb-nsql:7b-q4_0', system='''Hier ist das Datenbankschema, auf dem die SQL-Abfrage ausgeführt wird: CREATE TABLE taxi ( VendorID bigint, tpep_pickup_datetime timestamp, tpep_dropoff_datetime timestamp, passenger_count double, trip_distance double, fare_amount double, extra double, tip_amount double, tolls_amount double, improvement_surcharge double, total_amount double, );''', prompt='gib alle Spalten aus der Tabelle taxi zurück, die auf _amount enden', ) print(r['response'])

Trainingsdaten

200.000 DuckDB-Text-to-SQL-Paare, die unter Anleitung der Dokumentation zu DuckDB v0.9.2 synthetisch mit Mixtral-8x7B-Instruct-v0.1 erzeugt wurden. Hinzu kommen Text-to-SQL-Paare aus NSText2SQL, die mithilfe von sqlglot in DuckDB-SQL übertragen wurden.

Trainingsverfahren

DuckDB-NSQL wurde mit Kreuzentropieverlust trainiert, um die Wahrscheinlichkeit sequenzieller Eingaben zu maximieren. Beim Fine-Tuning mit Text-to-SQL-Paaren berechnen wir den Verlust nur für den SQL-Teil des Paares. Das Modell wurde auf A100-GPUs mit 80 GB unter Nutzung von Daten- und Modellparallelität trainiert. Das Fine-Tuning umfasste 10 Epochen.

Vorgesehene Verwendung und Einschränkungen

Das Modell wurde für die Text-to-SQL-Generierung aus vorgegebenen Tabellenschemata und natürlichsprachlichen Prompts entwickelt. Am besten funktioniert es mit dem nachfolgend definierten Promptformat und den entsprechenden Ausgaben. Anders als bestehende Text-to-SQL-Modelle ist die SQL-Generierung nicht auf SELECT-Anweisungen beschränkt. Das Modell kann jede gültige DuckDB-SQL-Anweisung erzeugen, einschließlich Anweisungen für offizielle DuckDB-Erweiterungen.

Referenzen

Hugging Face