DatabricksのバッチLLM推論による構造化データの抽出および評価

こちらのサンプルノートブック**バッチ推論と構造化データ抽出ノートブック**をウォークスルーします。

https://docs.databricks.com/aws/ja/large-language-models/ai-query-batch-inference#%E3%83%90%E3%83%83%E3%83%81%E6%8E%A8%E8%AB%96%E3%81%A8%E6%A7%8B%E9%80%A0%E5%8C%96%E3%83%87%E3%83%BC%E3%82%BF%E6%8A%BD%E5%87%BA

翻訳したノートブックはこちらです。

https://github.com/taka-yayoi/llm_batch_inference/blob/main/simple_structured_data_extraction_ai_query.ipynb

# サンプル: 構造化データ抽出、バッチ推論、および評価

このノートブックは、`ai_query` ([AWS](https://docs.databricks.com/aws/ja/sql/language-manual/functions/ai_query) | [Azure](https://learn.microsoft.com/ja-jp/azure/databricks/sql/language-manual/functions/ai_query)) を使用して基本的な構造化データ抽出を実行する方法を示します。

このプロセスは、生の非構造化データを自動抽出技術を通じて効果的に整理された実用的な情報に変換する方法を説明します。

このノートブックは、Mosaic AI Agent Evaluation ([AWS](https://docs.databricks.com/aws/ja/generative-ai/agent-evaluation) | [Azure](https://learn.microsoft.com/ja-jp/azure/databricks/generative-ai/agent-evaluation/)) を活用して、正解データが利用可能な場合の精度を評価する方法も示します。

```py
%pip install -U -qqqq mlflow databricks-agents
dbutils.library.restartPython()
```

## `ai_query` を使用したバッチ推論の実行
このノートブックでは、構造化データ抽出のために `ai_query` を使用する方法を示すために、雇用契約のシミュレートされたデータセットを作成します。このダミーデータセットは、雇用者および従業員の名前などの重要な情報に焦点を当てたエンティティ抽出のテストベッドとして機能します。抽出されるデータの真実の値が含まれており、後で評価に使用されます。

次に、このノートブックでは、このデータセットを使用して `ai_query` ([AWS](https://docs.databricks.com/aws/ja/sql/language-manual/functions/ai_query) | [Azure](https://learn.microsoft.com/ja-jp/azure/databricks/sql/language-manual/functions/ai_query)) を使用したバッチ推論を実施します。

### ダミーの契約データ

```py
# ダミーの雇用契約とその正解データ
employment_contracts = [
    dict(
        contract_text="この雇用契約は2023年5月15日にTechCorp Inc.（以下「雇用主」）とSarah Johnson（以下「従業員」）の間で締結されます。従業員は2023年6月1日からソフトウェアエンジニアとして勤務を開始し、年俸は85,000ドルです。従業員は3か月の試用期間に同意します。",
        ground_truth='{"signature_date": "2023年5月15日", "employer": "TechCorp Inc.", "employee": "Sarah Johnson", "bonuses": ["N/A"]}'
    ),
    dict(
        contract_text="雇用契約：2023年7月1日付で、DataSystems LLC（以下「雇用主」）はMichael Chen（以下「従業員」）をデータアナリストとして雇用することに同意します。従業員の初任給は年俸70,000ドルです。この契約には、終了後12か月間有効な競業避止条項が含まれています。",
        ground_truth='{"signature_date": "2023年7月1日", "employer": "DataSystems LLC", "employee": "Michael Chen", "bonuses": ["N/A"]}'
    ),
    dict(
        contract_text="2023年8月15日、CloudNet Solutions（「雇用主」）とEmma Rodriguez（「従業員」）はこの雇用契約を締結します。従業員はネットワーク管理者として2023年9月1日に勤務を開始します。年俸は78,000ドルで、5,000ドルのサインボーナスがあります。",
        ground_truth='{"signature_date": "2023年8月15日", "employer":"CloudNet Solutions", "employee": "Emma Rodriguez", "bonuses":["5,000ドルのサインボーナス"]}'
    ),
    dict(
        contract_text="この契約は2023年10月1日付で、AI Innovations Corp（「雇用主」）とDr. James Lee（「従業員」）の間で締結されます。Dr. Leeは主任研究科学者として2023年11月1日に就任します。基本給は年俸150,000ドルで、業績に基づくボーナスは付録Aに記載されています。",
        ground_truth='{"signature_date": "2023年10月1日", "employer": "AI Innovations Corp", "employee": "Dr. James Lee", "bonuses":["付録Aに記載された業績に基づくボーナス"]}'
    ),
]

employment_contracts_df = spark.createDataFrame(employment_contracts)
employment_contracts_df.display()
```

|contract_text|ground_truth|
|---|---|
|この雇用契約は2023年5月15日にTechCorp Inc.（以下「雇用主」）とSarah Johnson（以下「従業員」）の間で締結されます。従業員は2023年6月1日からソフトウェアエンジニアとして勤務を開始し、年俸は85,000ドルです。従業員は3か月の試用期間に同意します。|{"signature_date": "2023年5月15日", "employer": "TechCorp Inc.", "employee": "Sarah Johnson", "bonuses": ["N/A"]}|
|雇用契約：2023年7月1日付で、DataSystems LLC（以下「雇用主」）はMichael Chen（以下「従業員」）をデータアナリストとして雇用することに同意します。従業員の初任給は年俸70,000ドルです。この契約には、終了後12か月間有効な競業避止条項が含まれています。|{"signature_date": "2023年7月1日", "employer": "DataSystems LLC", "employee": "Michael Chen", "bonuses": ["N/A"]}|
|2023年8月15日、CloudNet Solutions（「雇用主」）とEmma Rodriguez（「従業員」）はこの雇用契約を締結します。従業員はネットワーク管理者として2023年9月1日に勤務を開始します。年俸は78,000ドルで、5,000ドルのサインボーナスがあります。|{"signature_date": "2023年8月15日", "employer":"CloudNet Solutions", "employee": "Emma Rodriguez", "bonuses":["5,000ドルのサインボーナス"]}|
|この契約は2023年10月1日付で、AI Innovations Corp（「雇用主」）とDr. James Lee（「従業員」）の間で締結されます。Dr. Leeは主任研究科学者として2023年11月1日に就任します。基本給は年俸150,000ドルで、業績に基づくボーナスは付録Aに記載されています。|{"signature_date": "2023年10月1日", "employer": "AI Innovations Corp", "employee": "Dr. James Lee", "bonuses":["付録Aに記載された業績に基づくボーナス"]}|

## `ai_query` を使用した構造化データ抽出

次のセルでは、`ai_query` を使用して構造化データ抽出を実行するために必要な主な入力を定義します:
- LLM エンドポイント名
- データ抽出を実行し、JSON を応答形式として使用するように LLM に指示するプロンプト
- 応答の JSON スキーマ

### プロンプトとレスポンスのフォーマットの定義

```py
import json

LLM_ENDPOINT_NAME = "databricks-meta-llama-3-3-70b-instruct"

PROMPT = """あなたは法的文書の分析に特化したAIアシスタントです。
あなたのタスクは、与えられた契約書から関連情報を抽出することです。
出力は構造化されたJSONオブジェクトでなければなりません。

指示:
1. このプロンプトの最後に提供される契約書全体を注意深く読んでください。
2. 関連情報を抽出してください。
3. 以下に指定された形式のJSONで結果を提示してください。

重要な注意事項:
- 関連情報のみを抽出してください。
- 契約全体の文脈を考慮して関連性を判断してください。
- 冗長にならず、正しい形式と情報のみを返してください。
- 一部の質問には関連する抜粋がない場合があります。その場合は、期待される型に応じて "N/A" または ["N/A"] を返してください。
- ここに記載されているJSONキー以外は含めないでください。
- 同じキーをJSONに複数回含めないでください。

期待されるJSONキーとその説明:
- signature_date: 契約の署名日。
- employer: 雇用主の名前。
- employee: 従業員の名前。
- bonuses: 記載されている特定のボーナスのリスト。

分析する契約書:
"""

response_format = json.dumps({
    "type": "json_schema",
    "json_schema": {
        "name": "employment_contract_extraction",
        "schema": {
            "type": "object",
            "properties": {
                "signature_date": {"type": "string"},
                "employer": {"type": "string"},
                "employee": {"type": "string"},
                "bonuses": {"type": "array", "items": {"type": "string"}},
            },
            "strict": True,
        },
    },
})
```

## バッチ推論

以下では、上記で定義された入力を使用して、`ai_query` を SQL 式として Spark データフレームに適用します。LLM の応答である JSON 文字列を解析して、個々のデータポイントを抽出します。

### バッチ推論でai_queryを活用

```py
from pyspark.sql.functions import col, from_json

# クエリを定義
ai_query_expr = f"""
  ai_query(
    endpoint => '{LLM_ENDPOINT_NAME}',
    request => CONCAT('{PROMPT}', contract_text),
    responseFormat => '{response_format}',
    modelParameters => named_struct('temperature', 0.)
    ) AS response
  """

# 解析したいLLMレスポンス文字列のJSONスキーマ
json_schema = "STRUCT<signature_date STRING, employee STRING, employer STRING, bonuses ARRAY<STRING>>"

# バッチクエリを実行し、レスポンスを解析
employment_contracts_df = employment_contracts_df.selectExpr(
    "*", ai_query_expr
).withColumn("parsed_response", from_json(col("response"), json_schema))

employment_contracts_df.display()
```

|contract_text|ground_truth|response|parsed_response|
|---|---|---|---|
|この雇用契約は2023年5月15日にTechCorp Inc.（以下「雇用主」）とSarah Johnson（以下「従業員」）の間で締結されます。従業員は2023年6月1日からソフトウェアエンジニアとして勤務を開始し、年俸は85,000ドルです。従業員は3か月の試用期間に同意します。|{"signature_date": "2023年5月15日", "employer": "TechCorp Inc.", "employee": "Sarah Johnson", "bonuses": ["N/A"]}|{"signature_date": "2023年5月15日", "employer": "TechCorp Inc.", "employee": "Sarah Johnson", "bonuses": ["N/A"]}|{"signature_date":"2023年5月15日","employee":"Sarah Johnson","employer":"TechCorp Inc.","bonuses":["N/A"]}|
|雇用契約：2023年7月1日付で、DataSystems LLC（以下「雇用主」）はMichael Chen（以下「従業員」）をデータアナリストとして雇用することに同意します。従業員の初任給は年俸70,000ドルです。この契約には、終了後12か月間有効な競業避止条項が含まれています。|{"signature_date": "2023年7月1日", "employer": "DataSystems LLC", "employee": "Michael Chen", "bonuses": ["N/A"]}|{"signature_date": "2023年7月1日", "employer": "DataSystems LLC", "employee": "Michael Chen", "bonuses": ["N/A"]}|{"signature_date":"2023年7月1日","employee":"Michael Chen","employer":"DataSystems LLC","bonuses":["N/A"]}|
|2023年8月15日、CloudNet Solutions（「雇用主」）とEmma Rodriguez（「従業員」）はこの雇用契約を締結します。従業員はネットワーク管理者として2023年9月1日に勤務を開始します。年俸は78,000ドルで、5,000ドルのサインボーナスがあります。|{"signature_date": "2023年8月15日", "employer":"CloudNet Solutions", "employee": "Emma Rodriguez", "bonuses":["5,000ドルのサインボーナス"]}|{"signature_date": "2023年8月15日", "employer": "CloudNet Solutions", "employee": "Emma Rodriguez", "bonuses": ["5,000ドルのサインボーナス"]}|{"signature_date":"2023年8月15日","employee":"Emma Rodriguez","employer":"CloudNet Solutions","bonuses":["5,000ドルのサインボーナス"]}|
|この契約は2023年10月1日付で、AI Innovations Corp（「雇用主」）とDr. James Lee（「従業員」）の間で締結されます。Dr. Leeは主任研究科学者として2023年11月1日に就任します。基本給は年俸150,000ドルで、業績に基づくボーナスは付録Aに記載されています。|{"signature_date": "2023年10月1日", "employer": "AI Innovations Corp", "employee": "Dr. James Lee", "bonuses":["付録Aに記載された業績に基づくボーナス"]}|{"signature_date": "2023年10月1日", "employer": "AI Innovations Corp", "employee": "Dr. James Lee", "bonuses": ["業績に基づくボーナス（付録Aに記載）"]}|{"signature_date":"2023年10月1日","employee":"Dr. James Lee","employer":"AI Innovations Corp","bonuses":["業績に基づくボーナス（付録Aに記載）"]}|

## エージェント評価によるエージェントの評価

エージェントの品質を評価するために、エージェント評価フレームワーク ([AWS](https://docs.databricks.com/aws/ja/generative-ai/agent-evaluation) | [Azure](https://learn.microsoft.com/ja-jp/azure/databricks/generative-ai/agent-evaluation/)) を使用します。このアプローチでは、正確性判定を使用して、期待されるエンティティ（または事実）と実際の応答を比較し、エージェントのパフォーマンスを包括的に評価します。

:::note info
**注意**

別のアプローチとして、個々のエンティティに対して `recall` や `precision` などのメトリクスを計算する方法もありますが、これには追加のデータ変換や [カスタムメトリクス](https://mlflow.org/docs/latest/python_api/mlflow.metrics.html#mlflow-metrics) が必要です。
:::

```py
import mlflow
from pyspark.sql.functions import split

# mlflow.evaluate() が期待する評価用データフレームを準備
eval_pdf = employment_contracts_df.select(
    col("contract_text").alias("request"),
    col("response"),
    split("ground_truth", ",").alias("expected_facts"),
).toPandas()

# 評価を実行し、結果を mlflow 実験に記録
with mlflow.start_run():
    results = mlflow.evaluate(
        data=eval_pdf,
        model_type="databricks-agent",
        evaluator_config={
            "databricks-agent": {
                "metrics": ["correctness"],  # 正確性の評価のみ有効化
            }
        }
    )
```

評価結果が記録されます。

![Screenshot 2025-03-17 at 13.31.09.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1168882/81d630a1-034f-4020-a3b5-3cc514768ece.png)
![Screenshot 2025-03-17 at 13.30.33.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1168882/7183c9ea-1b1b-4ec2-9ab2-9f5886429c1a.png)
![Screenshot 2025-03-17 at 13.30.46.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1168882/24b632f4-8fd2-49f1-94ef-0ed8df695df1.png)

## 次のステップ
Databricksでの構造化データ抽出に関するさらなる洞察と関連する例については、以下の包括的な技術ブログ投稿を参照してください:
- [LLMを使用したエンドツーエンドの構造化抽出 – パート1: バッチエンティティ抽出](https://community.databricks.com/t5/technical-blog/end-to-end-structured-extraction-with-llm-part-1-batch-entity/ba-p/98396)
- [LLMを使用したエンドツーエンドの構造化抽出 – パート2: 合成データを用いたファインチューニング](https://community.databricks.com/t5/technical-blog/end-to-end-structured-extraction-with-llm-part-2-fine-tuning/ba-p/99900)

### はじめてのDatabricks

[はじめてのDatabricks](https://qiita.com/taka_yayoi/items/8dc72d083edb879a5e5d)

### Databricks無料トライアル

[Databricks無料トライアル](https://databricks.com/jp/try-databricks)