LLMのためのコーディングテスト作成 〜Spark SQLに焦点を当てて〜

はじめに

コード生成のための大規模言語モデル（LLM）の活用はますます一般的になっており、より速く、よりスマートにコーディングできる利点があります。しかし、LLMが生成するコードの正確性が主な懸念点です。多くのオープンソースのコーディングベンチマークは一般的なコーディングスキルの評価を目的としていますが、企業環境では、LLMは一般的なプログラミング能力だけでなく、MLflowやSpark SQLといった特定のライブラリやツールの利用にも対応する必要があります。このため、LLMが特定のコーディングライブラリにおける能力を体系的に評価する方法が求められています。

本ブログ記事では、この課題に対処するため、LLM向けのライブラリ特化型コードテストを生成する手法をご紹介します。これらの生成されたテストケースは、モデルを評価するための構造化された方法を提供し、特定のライブラリに適したモデルを選定する助けとなります。また、ドメイン固有のファインチューニングにより、特定のライブラリへの熟練度向上も測定可能です。

この記事では、Databricks Assistant Autocompleteのモデル評価に使用する内部ベンチマークに組み込まれたSpark SQL用のコードテストの生成方法を示します。関数名、定義、サンプルコードを含むコードドキュメントを活用することで、非常に特化したコードテストを汎用的に生成するプロセスを開発しました。

図1：array_except関数の生成されたコードテスト。左側には、Spark SQL APIに記載されている関数のソース情報が表示されており、右側には2つの生成されたコードテストが示されています。評価時には、モデルに右側のコンテキストが提示され、<here>プレースホルダーで適切なコードを生成するよう指示されます。この生成されたコード指示はテストの要となっており、上の例はコードの目的と必要な入力データが明確に表現されているため理想的です。一方、下の例はコメントの意味が曖昧であるため、問題があります。

アプローチ

コードドキュメントに基づき、私たちのテストケース生成パイプラインは以下の主要ステップで構成されています。

1. シード関数フィルタリング：提供されたコードドキュメントから、パイプラインでの自動テストの基準を満たすシード関数を選定します。

2. コード指示生成：最新のモデル（SOTAモデル）を用いて、各関数に関する情報をもとに詳細なコード指示（コメント）を生成します。これらの指示には、機能の説明と入力データの要件が明確に記述されている必要があります。

3. コード指示の検証：生成されたコード指示の信頼性を確保するため、SOTAモデルを用いて指示を解釈し、可能なソリューションを生成します。この際、モデルの限界を補うために関連するメタ情報が提供されます。生成されたソリューションを実行し、結果を元のコードスニペットの結果と比較することで、指示が正しいコード生成に導いているかを確認します。異なる、または予期しない結果が出た場合は、その質を手動で検証し、逸脱があっても高品質であれば保持し、そうでない場合はテストプロセスの整合性を保つために除外します。

シード関数のフィルタリング

コードドキュメントに記載されている各関数には、通常、高品質な例が付属しており、使用方法を理解しやすくしています。しかし、すべての関数が自動テストの候補に適しているわけではありません。テストケース生成のシードとして有効であるためには、例のコードが以下の2つの条件を満たしている必要があります。

1. 決定論的な出力：コードの実行結果が決定論的である必要があり、これは後続の検証ステップで重要です。例えば、rand()やcurrent_date()のようなランダムや時間依存の結果を生成する関数は、その性質上予測不可能であるため適していません。

2. 実行環境との互換性：コードは指定されたコーディング環境内で実行可能である必要があります。例えば、DatabricksのUnity Catalogで実行する必要がある場合、UC共有モードでサポートされていない関数は避ける必要があります。

これを確認するため、ターゲット環境で各サンプルコードを実行し、その結果を記録します。結果が参照APIドキュメントの内容と一致した場合、その関数とコードは決定性があると判断されて保持されます。逆に、エラーが発生する場合は、実行環境との互換性がないとみなし、自動テストの候補から除外されます。このフィルタリングを終えることで、自動テストが可能であり、指定の環境で実行できる関数のセットが確保されます。

コード指示生成

ここで、自動テストケース生成の核心となるステップに進みます：指示の生成です。この指示を基にコードを生成すると、シード関数の例と同じ実行結果が得られることを目指します。最先端（SOTA）のコードモデルに対し、各シード関数に対応するコーディング指示を生成するよう促します。このモデルへの入力は、関数名、その定義、および1つの例コードから構成されます。生成されるコード指示は、例コードの内容を説明する簡潔なコメントです。

モデルの出力を効果的に導くため、プロンプト内で特定の要件を設定することが重要です。これにより、指示がモデルの知識を信頼性をもってテストするものとなります。プロンプトでは、SOTAモデルに次の指示を出します：

• コメントには関数名を記載せず、例コードで入力データが与えられている場合はそのデータを明示すること。
• コメントは十分な詳細を含み、コメントの情報だけで対応するコードが識別できるようにすること。

これにより、コメント内で解答を明かすことなく、同時にコメントから動作する例が生成できるために必要な情報が提供されるようになります。

コード指示の検証

生成されたコード指示は、テストケースにおいて重要な役割を果たします。ターゲットモデルを効果的に評価するために、これらの指示はプロンプトとして機能し、関数の目的や関連する入力データを明確に表現する必要があります。曖昧さはモデルの出力の正確性を損なうため、正確なコード生成には明確な指示が不可欠です。以下に、不適切とされるコード指示の例を示します。

コード指示が基準を満たしているかを確認するために、次の検証プロセスを行います。これらの指示を用いて最先端（SOTA）コードモデルにプロンプトを出し、対応する解答を生成させます。その解答を実行し、シードコードスニペットの結果と一致する場合、指示は保持され、正確なコード生成を可能にする十分な詳細が含まれていることが確認されます。

ここで、一つの懸念が生じる可能性があります。それは、SOTAモデルが指示を解釈できるだけの能力を持たない場合です。もしモデルが指示をうまく解釈できなければ、それは指示の質ではなく、モデルの限界が原因かもしれません。この点を補うため、関数名や定義など必要な事前知識をすべてプロンプトに含め、SOTAモデルが提供された包括的な情報に基づいて決定論的な解を生成できるようにしています。また、モデル生成の解が失敗したテストについては手動で見直し、失敗していても高品質なものは保持します。

コードモデル評価

実験設定

この実験では、特定のカーソル位置において、モデルが中央部分を埋めるモード（FIM）を用いて評価を行います。カーソル前のコードを「プレフィックス」、カーソル後のコードを「サフィックス」と呼びます。通常、これら2つのセグメントを示すセンチネルトークンが使用され、その後、中央を埋めるコードの生成を依頼するセンチネルが続きます。モデルへのプロンプトは次のようにフォーマットされます：「<fim_prefix>プレフィックスコード<fim_suffix>サフィックスコード<fim_middle>」。異なるモデルによって異なるセンチネルトークンが使用される場合があり、インフィリング形式も異なることがある点に注意が必要です。

私たちのSpark SQLテスト生成パイプラインでは、286のテストケースが生成されました！これらのテストケースを、上記の手法を用いて生成し、評価ベンチマークで実行するためにYAML形式に変換します。各YAMLファイルには以下の重要な要素が含まれます：

• Name: テストする関数名。特定の関数に対するモデルの性能を示すために使用されます。
• Context: このコンテキストは、必要なセンチネルトークンとともにFIM形式に変換されます。「<here>」はプレースホルダーで、後の評価で生成されたコードと置き換えます。この表現により、異なるFIM形式を使用するモデルにテストケースを柔軟に適応させることが可能です。
• Canonical solution: 参照解答であり、テストケースが適切に定義されているかを確認するために使用します。この基準解を用いてベンチマークを実行した場合、スコアは100％になるべきです。
• Test: ここにはアサーションチェックが含まれます。生成されたコードをコンテキストで実行し、その結果が参照結果と一致するかを検証します。

評価結果

パフォーマンスは、モデルが最初の試行で正しい解を生成した問題の割合を示す指標であるpass@1メトリック（Chen et al., 2021）を使用して報告します。これは、モデルが1回の推測でコーディング問題をどの程度成功裏に解けるかを示します。サンプリングには、top_pを0.95、温度を0.2に設定したヌクレアスサンプリングを使用しています。また、パラメータが70億程度のモデルを複数評価し、このベンチマークにおけるSOTA性能を把握するために、GPT-4oの貪欲デコーディングも評価に含めています。

表1：各LLMのSparkSQLベンチマークにおけるPass@k結果。モデルごとの独自のFIM形式および特別なトークンに従って評価を行いました。

モデル評価の際、「# Databricks notebook source」を冒頭に含めることで結果が向上することが分かりました。この行は常にDatabricksノートブックの最初に表示され、通常のPythonモジュールやスクリプトと区別されます。この効果は特にStarCoder2-7Bモデルで顕著で、この行がないとPass@1スコアが0.125まで大幅に低下します。この最初の行がヒントとして作用し、モデルがDatabricksノートブックの文脈で得たSpark SQLに関する重要な知識を推論時に参照できると仮定しています。

モデルが頻繁に失敗するテストを分析すると、多くの失敗が適切な組み込み関数を正しく識別して使用できないことに起因していることがわかります。たとえば、Spark SQLの「find_in_set」関数は、カンマ区切りのリスト内で特定の文字列のインデックスを返すように設計されていますが、モデルはしばしばこれを、ターゲット文字列内で部分文字列のインデックスを見つける「position」関数と混同します。また、モデルが複雑なネストされたサブクエリで指示を過剰に実装し、エラーを引き起こしやすくしていることもあります。一方で、単純な組み込み関数で解決できる場合もあるため、モデルの実装が不要に複雑化する傾向も見られます。

まとめ

私たちは、任意のコードライブラリに対して、提供されたドキュメントからコードテストを生成する手法を提案します。このテストケース生成パイプラインは、ドキュメントからシード関数をフィルタリングし、詳細なコード指示を生成し、それらを検証する一連のステップで構成されています。検証では、関数情報をヒントとしてコード指示と共に使用し、対応するコードソリューションを生成して実行し、その正確性を確認します。これにより、コード指示がモデルのコーディング能力を評価するのに有効であることが保証されます。最終的に、これらのテストケースを使用して、さまざまなモデルのインフィリングモードでの性能を評価します。

この投稿では、ドキュメントのサンプルコードをコードテストに最も直接的に変換する方法を紹介しています。この手法は、より複雑なテストケースにも対応するよう拡張できます。例えば、異なる入力データが必要な場合、シード関数のフィルタリング後にサンプルコードを適宜変更する追加ステップを導入できます。また、さまざまな条件のアサーションを追加することも可能です。現在のシナリオではターゲットコードが1行ですが、複数行のコードに対しては、簡潔なコードコメントではなく、より詳細なドックストリングが必要になるでしょう。また、前のコードをコンテキストとして使用し、モデルに特定の関数行のみを生成するよう指示することもできます。テストケースを特定の要件に合わせるために、さまざまな修正を実施できます。

次回の投稿では、このSpark SQLベンチマークでのパフォーマンス向上を目的にモデルをファインチューニングする方法について解説します。乞うご期待！