Databricks で Text2SQL のパフォーマンスを簡単に向上

テキストから SQL へのタスクで広く使用されているベンチマークである Spider のトップ 10 に LLM をランクアップしたいですか? Spider は、LLM がテキスト クエリを SQL コードに変換できるかどうかを評価します。

text-to-SQL に馴染みのない方のために説明すると、その重要性は、企業がデータと対話する方法を変革することにあります。 クエリの作成について SQL の専門家に頼る代わりに、ユーザーは平易な英語でデータについて質問するだけで、正確な回答を得ることができます。 これにより、データへのアクセスが民主化され、ビジネスインテリジェンスが強化され、より情報に基づいた意思決定が可能になります。

Spider ベンチマークは、text-to-SQL システムのパフォーマンスを評価するための広く認識されている標準です。LLM は自然言語クエリを正確な SQL ステートメントに変換することが求められ、データベース スキーマに関する深い理解と、構文的および意味的に正しい SQL コードを生成する能力が求められます。

この投稿では、オープンソースの Llama3 8B Instruct モデルを使用して、1 日未満の作業で Spider 開発データセットで 79.9% 、テスト データセットで 78.9% のスコアを達成した方法について詳しく説明します。これは、ベースラインに対して 19 ポイントという驚異的な改善です。 このパフォーマンスは、 Databricksでの戦略的なプロンプトとファインチューニングのおかげで、現在は凍結されている Spider リーダーボードでトップ 10 にランクインすることになります。

ベースライン性能のためのゼロショットプロンプト

まず、テーブルを作成した CREATE TABLE ステートメントと、それらのテーブルを使用して回答したい質問で構成される非常にシンプルなプロンプト形式を使用して、Spider dev データセットでの Meta Llama 3 8B Instruct のパフォーマンスを評価してみましょう。

{create_table_queries}

-- {question}
SELECT

このタイプのプロンプトは、プロンプトに他の例がないため、しばしば「ゼロショット」と呼ばれます。 Spider dev データセットの最初の質問では、このプロンプト形式によって次が生成されます。

CREATE TABLE stadium (
Stadium_ID int,
Location text,
Name text,
Capacity int,
Highest int,
Lowest int,
Average int,
PRIMARY KEY (Stadium_ID)
)
…<omitted the singer, concert, and singer_in_concert tables for brevity>

-- How many singers do we have?
SELECT

この形式を使用して開発データセットで Spider ベンチマークを実行すると、実行精度と貪欲なデコードを使用して測定した場合、全体のスコアは 60.9 になります。 これは、モデルが 60.9% の確率で SQL を生成し、それを実行すると、正しいソリューションを表す「ゴールド」クエリと同じ結果が生成されることを意味します。

ベースライン スコアが確立されたので、ファインチューニングに入る前に、Spider 開発ベンチマーク データセットでベース モデルのスコアを上げるために、さまざまなプロンプト戦略を試してみましょう。

サンプル行によるプロンプト

最初に使用したプロンプトの欠点の 1 つは、データ型以外の列のデータに関する情報が含まれていないことです。 Spiderを使用したモデルのテキストからSQLへの機能の評価に関する論文では、サンプリングされた行をプロンプトに追加するとスコアが高くなることがわかったので、それを試してみましょう。

上記のプロンプト形式を更新して、テーブル作成クエリに各テーブルの最初の数行も含まれるようにすることができます。 以前の同じ質問に対して、更新されたプロンプトはありません。

CREATE TABLE stadium (
Stadium_ID int,
Location text,
Name text,
Capacity int,
Highest int,
Lowest int,
Average int,
PRIMARY KEY (Stadium_ID)
)
/*
Stadium_ID    Location    Name    Capacity    Highest    Lowest    
Average
1    Raith Rovers    Stark's Park    10104    4812    1294    2106
2    Ayr United    Somerset Park    11998    2363    1057    1477
3    East Fife    Bayview Stadium    2000    1980    533    864
*/
…<omitted the singer, concert, and singer_in_concert tables for 
brevity>

-- How many singers do we have?
SELECT

各テーブルのサンプル行を含めると、全体のスコアが約 6 パーセント ポイント上昇して 67.0 になります。

少数のプロンプト（Few-shot Prompting）

少数のプロンプトは、LLM で使用されるよく知られた戦略であり、実行するタスクを示すいくつかの例を含めることで、正しい SQL を生成するなどのタスクのパフォーマンスを向上させることができます。 ゼロショットのプロンプトで、スキーマを提供し、質問をしました。 数回のプロンプトで、いくつかのスキーマ、質問、その質問に答える SQL を提供し、そのシーケンスを数回繰り返してから、実際に尋ねたい質問に進みます。 これにより、通常、ゼロショットプロンプトよりもパフォーマンスが向上します。

SQL生成タスクを示す例の良いソースは、実際には Spider トレーニング データセットそのものです。 このデータセットから、対応するテーブルを含むいくつかの質問をランダムにサンプルとして取得し、これらの質問のそれぞれに回答できる SQL を示す数回のプロンプトを作成できます。 前のプロンプトの時点でサンプル行を使用しているため、これらの例の 1 つにサンプル行も含まれていることを確認して、使用法を示す必要があります。

以前のゼロショットプロンプトから改善できるもう1つの改善点は、最初に「システムプロンプト」を含めることです。 システム プロンプトは通常、実行されるタスクの概要を示す詳細なガイダンスをモデルに提供するために使用されます。 ユーザーはモデルとのチャットの過程で複数の質問をすることができますが、システムプロンプトはユーザーが質問する前に一度だけ提供されるため、基本的にチャット中に「システム」がどのように機能すべきかについての期待が確立されます。

これらの戦略を念頭に置いて、上部に大きな SQL コメント ブロックとして表されるシステム メッセージで始まり、その後に 3 つの例が続く、数回のプロンプトを作成できます。

/*
You are a helpful assistant who answers questions about database tables 
by responding with SQL queries.  Users will provide you with a set of 
tables represented as CREATE TABLE statements.  Each CREATE TABLE 
statement may optionally be followed by the first few rows from the 
table in order to help write the correct SQL to answer questions. After 
the CREATE TABLE statements users will ask a question using a SQL 
comment starting with two dashes. You should answer the user's question 
by writing a SQL statement starting with SELECT and ending with a 
semicolon.
*/

CREATE TABLE "Campuses" (
	"Id" INTEGER PRIMARY KEY,
	"Campus" TEXT,
	"Location" TEXT,
	"County" TEXT,
	"Year" INTEGER
);
/*
Id    Campus    Location    County    Year
1    California State University-Bakersfield    Bakersfield    Kern    
1965
2    California State University-Channel Islands    Camarillo    
Ventura    2002
3    California State University-Chico    Chico    Butte    1887
*/

… <more tables omitted>

-- Please answer the following question using the tables above.
-- Find the name of the campuses that is in Northridge, Los Angeles or 
-- in San Francisco, San Francisco.
SELECT Campus FROM Campuses WHERE Location="Northridge" AND County="Los 
Angeles" 
UNION SELECT Campus FROM Campuses WHERE Location="San Francisco" AND 
County="San Francisco";

… <two more examples omitted>

CREATE TABLE stadium (
Stadium_ID int,
Location text,
Name text,
Capacity int,
Highest int,
Lowest int,
Average int,
PRIMARY KEY (Stadium_ID)
)
/*
Stadium_ID    Location    Name    Capacity    Highest    Lowest    
Average
1    Raith Rovers    Stark's Park    10104    4812    1294    2106
2    Ayr United    Somerset Park    11998    2363    1057    1477
3    East Fife    Bayview Stadium    2000    1980    533    864
*/
…<omitted the singer, concert, and singer_in_concert tables for 
brevity>

-- How many singers do we have?
SELECT

この新しいプロンプトの結果、スコアは 70.8 となり、前回のスコアからさらに 3.8 ポイント改善されました。 私たちは、単純なプロンプト戦略だけで、開始したところからスコアを 10%近く 引き上げました。

おそらく、プロンプトを微調整することによる収穫逓減のポイントに達しているのでしょう。 モデルを微調整して、さらにどのようなメリットが得られるかを確認しましょう。

LoRAとの連携

モデルをファインチューニングする場合、最初の質問はどのトレーニング データを使用するかということです。 Spider にはトレーニング データセットが含まれているため、ここから始めるのが良いと思われます。 モデルを微調整するには、 Databricks の Standard_NC24ads_A100_v4 などの単一の A100 80GB Databricks GPU クラスターでモデルを効率的にトレーニングできるように QLoRA を使用します。 これは、Spider トレーニング データセットの 7,000 件のレコードを使用して約 4 時間で完了できます。 以前のブログ記事で、LoRA によるファインチューニングについて説明しました。 興味のある読者は、その投稿で詳細を参照してください。 trl、peft、bitsandbytes ライブラリを使用して、標準的なトレーニング レシピに従うことができます。

Spider からトレーニング レコードを取得していますが、モデルが学習できるような形式でフォーマットする必要があります。 目標は、スキーマ (サンプル行を含む)、質問、SQL で構成される各レコードを 1 つのテキスト文字列にマップすることです。 まず、生の Spider データセットに対していくつかの処理を実行します。 生データから、各レコードがschema_with_rows、question、query の 3 つのフィールドで構成されるデータセットを生成します。 schema_with_rows フィールドは、以前の数回のプロンプトで使用されたCREATE TABLEステートメントと行のフォーマットに従って、質問に対応するテーブルから取得されます。

次に、トークナイザーをロードします。

tokenizer = 
AutoTokenizer.from_pretrained("meta-llama/Meta-Llama-3-8B-Instruct")

処理済みの Spider トレーニング データセットの各レコードをテキスト文字列に変換するマッピング関数を定義します。 トークナイザーから apply_chat_template を使用して、テキストを Instruct モデルが期待するチャット形式に便利にフォーマットできます。 これは、few-shot プロンプトに使用している形式とまったく同じではありませんが、プロンプトの定型文形式がわずかに異なっていても、モデルは十分に一般化されています。

def _mapper(rec):
    schema = rec["schema_with_rows"].strip()
    question = rec["question"].strip()
    query = rec["query"].strip()

    user_message = USER_MESSAGE_FORMAT.format(schema=schema, 
    question=question)

    messages = [
     {
       "role": "system",
       "content": SYSTEM_PROMPT,
     },
     {"role": "user", "content": user_message},
     {"role": "assistant", "content": query},
    ]
    prompt = tokenizer.apply_chat_template(messages, tokenize=False, 
    add_generation_prompt=False)
    return {"text": prompt}

SYSTEM_PROMPTには、前のfew-shotプロンプトで使用したのと同じシステムプロンプトを使用します。 USER_MESSAGE_FORMAT の場合も同様に使用します。

{schema}

Please answer the following question using the tables above.
{question}

この関数を定義したら、あとは処理された Spider データセットをそれを使って変換し、JSONL ファイルとして保存するだけです。

dataset.map(_mapper)

トレーニングする準備が整いました。 数時間後、微調整されたLlama3 8B Instructが手に入りました。 この新しいモデルで数ショットのプロンプトを再実行すると、スコアは 79.9になり、以前のスコアからさらに 9ポイント 改善されました。 これで、合計スコアが単純なゼロショットのベースラインよりも~19パーセントポイント 上昇しました。

Llama3 8B Instructモデルと微調整バージョンが、Llama3 70B Instructなどの大型モデルとどのように比較されるのか疑問に思われるかもしれません。 8 つの A100 40 GB GPU を搭載した開発データセットで既製の 70B モデルを使用して評価プロセスを繰り返し、以下の結果を記録しました。

予想通り、既製のモデルを比較すると、同じプロンプト形式で測定すると、70Bモデルが8Bモデルを上回っています。 しかし、驚くべきことに、微調整されたLlama3 8B Instructモデルは、Llama3 70B Instructモデルよりも 3%高いスコアを獲得しています。 テキストから SQL への変換などの特定のタスクに重点を置くと、ファインチューニングによって、サイズがはるかに大きいモデルとパフォーマンスが同等の小さなモデルを作成できます。

モデルサービングエンドポイントにデプロイする

Llama3 は Mosaic AI Model Serving でサポートされているため、微調整された Llama3 モデルをエンドポイントにデプロイし、アプリケーションを強化するために使用することもできます。 必要なのは、微調整されたモデルをUnity Catalogに記録し、UI を使用してエンドポイントを作成することだけです。 デプロイされると、共通ライブラリを使用してクエリを実行できます。

まとめ

私たちは、ゼロショットプロンプトを使用して Spider 開発データセットで Llama3 8B Instruct の取り組みを開始し、60.9 という控えめなスコアを達成しました。 システムメッセージ、複数の例、サンプル行を含む数発のプロンプトでこれを強化することで、スコアを70.8に向上させました。 Spider トレーニング データセットでモデルをファインチューニングすることで、Spider dev で79.9 、Spider test で78.9 という素晴らしいスコアを達成し、さらなる向上が実現しました。 スタート地点から19ポイントの大幅な上昇と、ベースのLlama3 70B Instructに対する3ポイントのリードは、私たちのモデルの実力を示すだけでなく、Spiderのトップ10リザルトで切望されていた場所を確保することにもなります。

オープンソース LLM とデータ インテリジェンス プラットフォームのパワーを活用する方法の詳細については、Data+ AI Summitに登録してください。

Appendix

評価セットアップ

生成は、vLLM、貪欲デコード（温度 0）、2 つの A100 80 GB GPU、および最大 1024 個の新しいトークンを使用して実行されました。 世代を評価するために、Github のtaoyds/test-suite-sql-evalリポジトリのテスト スイートを使用しました。

トレーニングのセットアップ

ファインチューニング設定に関する具体的な詳細は次のとおりです。

ゼロショット プロンプトの例

これは、テーブル スキーマを含むゼロ ショット プロンプトとしてフォーマットされた、評価に使用した開発データセットの最初のレコードです。 質問の対象となっているテーブルは、そのテーブルを作成したCREATE TABLEステートメントを使用して表されます。

CREATE TABLE stadium (
Stadium_ID int,
Location text,
Name text,
Capacity int,
Highest int,
Lowest int,
Average int,
PRIMARY KEY (Stadium_ID)
)

CREATE TABLE singer (
Singer_ID int,
Name text,
Country text,
Song_Name text,
Song_release_year text,
Age int,
Is_male bool,
PRIMARY KEY (Singer_ID)
)

CREATE TABLE concert (
concert_ID int,
concert_Name text,
Theme text,
Stadium_ID text,
Year text,
PRIMARY KEY (concert_ID),
FOREIGN KEY (Stadium_ID) REFERENCES stadium(Stadium_ID)
)

CREATE TABLE singer_in_concert (
concert_ID int,
Singer_ID text,
PRIMARY KEY (concert_ID,Singer_ID),
FOREIGN KEY (concert_ID) REFERENCES concert(concert_ID),
FOREIGN KEY (Singer_ID) REFERENCES singer(Singer_ID)
)

-- How many singers do we have?
SELECT

サンプル行を使用したゼロショットプロンプトの例

これは、テーブル スキーマとサンプル行を含むゼロ ショット プロンプトとしてフォーマットされた、評価に使用した開発データセットの最初のレコードです。 質問の対象となっているテーブルは、そのテーブルを作成した CREATE TABLE ステートメントを使用して表されます。 行は、各テーブルから "SELECT * {table_name} LIMIT 3" を使用して選択され、列名がヘッダーとして表示されます。

CREATE TABLE stadium (
Stadium_ID int,
Location text,
Name text,
Capacity int,
Highest int,
Lowest int,
Average int,
PRIMARY KEY (Stadium_ID)
)
/*
Stadium_ID   Location   Name   Capacity   Highest   Lowest   Average
1   Raith Rovers   Stark's Park    10104    4812    1294    2106
2   Ayr United   Somerset Park    11998    2363    1057    1477
3   East Fife   Bayview Stadium    2000    1980    533    864
*/

CREATE TABLE singer (
Singer_ID int,
Name text,
Country text,
Song_Name text,
Song_release_year text,
Age int,
Is_male bool,
PRIMARY KEY (Singer_ID)
)
/*
Singer_ID    Name    Country    Song_Name   Song_release_year   Age   Is_male
1    Joe Sharp    Netherlands    You    1992    52    F
2    Timbaland    United States    Dangerous    2008    32    T
3    Justin Brown    France    Hey Oh    2013    29    T
*/

CREATE TABLE concert (
concert_ID int,
concert_Name text,
Theme text,
Stadium_ID text,
Year text,
PRIMARY KEY (concert_ID),
FOREIGN KEY (Stadium_ID) REFERENCES stadium(Stadium_ID)
)
/*
concert_ID    concert_Name    Theme    Stadium_ID    Year
1    Auditions    Free choice    1    2014
2    Super bootcamp    Free choice 2    2    2014
3    Home Visits    Bleeding Love    2    2015
*/

CREATE TABLE singer_in_concert (
concert_ID int,
Singer_ID text,
PRIMARY KEY (concert_ID,Singer_ID),
FOREIGN KEY (concert_ID) REFERENCES concert(concert_ID),
FOREIGN KEY (Singer_ID) REFERENCES singer(Singer_ID)
)
/*
concert_ID    Singer_ID
1    2
1    3
1    5
*/

-- How many singers do we have?
SELECT

サンプル行を使用した数ショットのプロンプトの例

これは、テーブル スキーマとサンプル行を含む、数回のプロンプトとしてフォーマットされた、評価に使用した開発データセットの最初のレコードです。 質問の対象となっているテーブルは、そのテーブルを作成した CREATE TABLE ステートメントを使用して表されます。 行は、各テーブルから "SELECT * {table_name} LIMIT 3" を使用して選択され、列名がヘッダーとして表示されます。

/*
You are a helpful assistant who answers questions about database tables by 
responding with SQL
queries.  Users will provide you with a set of tables represented as CREATE 
TABLE statements.  Each CREATE TABLE statement may optionally be followed by 
the first few rows from the table in order to help write the correct SQL to 
answer questions. After the CREATE TABLE statements users will ask a 
question using a SQL comment starting with two dashes. You should answer the 
user's question by writing a SQL statement starting with SELECT and ending 
with a semicolon.
*/

CREATE TABLE "Campuses" (
	"Id" INTEGER PRIMARY KEY,
	"Campus" TEXT,
	"Location" TEXT,
	"County" TEXT,
	"Year" INTEGER
);
/*
Id    Campus    Location    County    Year
1    California State University-Bakersfield    Bakersfield    Kern    1965
2    California State University-Channel Islands    Camarillo    Ventura    
2002
3    California State University-Chico    Chico    Butte    1887
*/

CREATE TABLE "csu_fees" (
	"Campus" INTEGER PRIMARY KEY,
	"Year" INTEGER,
	"CampusFee" INTEGER,
	FOREIGN KEY (Campus) REFERENCES Campuses(Id)
);
/*
Campus    Year    CampusFee
1    1996    1951
2    2003    1868
3    1996    2042
*/

CREATE TABLE "degrees" (
	"Year" INTEGER,
	"Campus" INTEGER,
	"Degrees" INTEGER,
	PRIMARY KEY (Year, Campus),
	FOREIGN KEY (Campus) REFERENCES Campuses(Id)
);
/*
Year    Campus    Degrees
1990    1    701
1991    1    681
1992    1    791
*/

CREATE TABLE "discipline_enrollments" (
	"Campus" INTEGER,
	"Discipline" INTEGER,
	"Year" INTEGER,
	"Undergraduate" INTEGER,
	"Graduate" INTEGER,
	PRIMARY KEY (Campus, Discipline),
	FOREIGN KEY (Campus) REFERENCES Campuses(Id)
);
/*
Campus    Discipline    Year    Undergraduate    Graduate
1    4    2004    248    0
1    5    2004    811    73
1    6    2004    199    0
*/

CREATE TABLE "enrollments" (
	"Campus" INTEGER,
	"Year" INTEGER,
	"TotalEnrollment_AY" INTEGER,
	"FTE_AY" INTEGER,
	PRIMARY KEY(Campus, Year),
	FOREIGN KEY (Campus) REFERENCES Campuses(Id)
);
/*
Campus    Year    TotalEnrollment_AY    FTE_AY
1    1956    384    123
1    1957    432    151
1    1958    422    178
*/

CREATE TABLE "faculty" (
	"Campus" INTEGER,
	"Year" INTEGER,
	"Faculty" REAL,
	FOREIGN KEY (Campus) REFERENCES Campuses(Id)
);
/*
Campus    Year    Faculty
1    2002    357.1
2    2002    48.4
3    2002    742.8
*/

-- Please answer the following question using the tables above.
-- Find the name of the campuses that is in Northridge, Los Angeles or in 
San Francisco, San Francisco.
SELECT Campus FROM Campuses WHERE Location="Northridge" AND County="Los 
Angeles" UNION SELECT Campus \
FROM Campuses WHERE Location="San Francisco" AND County="San Francisco";


CREATE TABLE Allergy_Type (
       Allergy 		  VARCHAR(20) PRIMARY KEY,
       AllergyType 	  VARCHAR(20)
);

CREATE TABLE Has_Allergy (
       StuID 		 INTEGER,
       Allergy 		 VARCHAR(20),
       FOREIGN KEY(StuID) REFERENCES Student(StuID),
       FOREIGN KEY(Allergy) REFERENCES Allergy_Type(Allergy)
);

CREATE TABLE Student (
        StuID        INTEGER PRIMARY KEY,
        LName        VARCHAR(12),
        Fname        VARCHAR(12),
        Age      INTEGER,
        Sex      VARCHAR(1),
        Major        INTEGER,
        Advisor      INTEGER,
        city_code    VARCHAR(3)
 );

-- Please answer the following question using the tables above.
-- Which allergy type has most number of allergies?
SELECT AllergyType FROM Allergy_Type GROUP BY AllergyType ORDER BY count(*) 
DESC LIMIT 1;


CREATE TABLE "building" (
"building_id" text,
"Name" text,
"Street_address" text,
"Years_as_tallest" text,
"Height_feet" int,
"Floors" int,
PRIMARY KEY("building_id")
);

CREATE TABLE "Institution" (
"Institution_id"  text,
"Institution" text,
"Location" text,
"Founded" real,
"Type" text,
"Enrollment" int,
"Team" text,
"Primary_Conference" text,
"building_id" text,
PRIMARY KEY("Institution_id"),
FOREIGN  KEY ("building_id") REFERENCES "building"("building_id")
);

CREATE TABLE "protein" (
"common_name" text,
"protein_name" text,
"divergence_from_human_lineage" real,
"accession_number" text,
"sequence_length" real,
"sequence_identity_to_human_protein" text,
"Institution_id" text,
PRIMARY KEY("common_name"),
FOREIGN KEY("Institution_id") REFERENCES "Institution"("Institution_id")
);


-- Please answer the following question using the tables above.
-- For each building, show the name of the building and the number of 
institutions in it.
SELECT T1.name, count(*) FROM building AS T1 JOIN Institution AS T2 ON 
T1.building_id=\
T2.building_id GROUP BY T1.building_id;


CREATE TABLE stadium (
Stadium_ID int,
Location text,
Name text,
Capacity int,
Highest int,
Lowest int,
Average int,
PRIMARY KEY (Stadium_ID)
)
/*
Stadium_ID   Location   Name   Capacity   Highest   Lowest   Average
1   Raith Rovers   Stark's Park   10104   4812   1294   2106
2   Ayr United   Somerset Park   11998   2363   1057   1477
3   East Fife   Bayview Stadium   2000   1980   533   864
*/

CREATE TABLE singer (
Singer_ID int,
Name text,
Country text,
Song_Name text,
Song_release_year text,
Age int,
Is_male bool,
PRIMARY KEY (Singer_ID)
)
/*
Singer_ID    Name    Country    Song_Name    Song_release_year    Age    
Is_male
1    Joe Sharp    Netherlands    You    1992    52    F
2    Timbaland    United States    Dangerous    2008    32    T
3    Justin Brown    France    Hey Oh    2013    29    T
*/

CREATE TABLE concert (
concert_ID int,
concert_Name text,
Theme text,
Stadium_ID text,
Year text,
PRIMARY KEY (concert_ID),
FOREIGN KEY (Stadium_ID) REFERENCES stadium(Stadium_ID)
)
/*
concert_ID    concert_Name    Theme    Stadium_ID    Year
1    Auditions    Free choice    1    2014
2    Super bootcamp    Free choice 2    2    2014
3    Home Visits    Bleeding Love    2    2015
*/

CREATE TABLE singer_in_concert (
concert_ID int,
Singer_ID text,
PRIMARY KEY (concert_ID,Singer_ID),
FOREIGN KEY (concert_ID) REFERENCES concert(concert_ID),
FOREIGN KEY (Singer_ID) REFERENCES singer(Singer_ID)
)
/*
concert_ID    Singer_ID
1    2
1    3
1    5
*/

-- How many singers do we have?
SELECT