【Databricks】LLM構造化抽出の評価・改善サイクルをMLflowで構築してみる

{

　　"company": "Example Company A",

　　"events": [

　　　　{

　　　　　　"product": "Example Service",

　　　　　　"action": "adopt",

　　　　　　"status": "completed",

　　　　　　"purpose": "データ分析基盤として利用"

　　　　}

　　]

}

%pip install "mlflow-skinny[databricks]==3.16.1" "databricks-agents==1.12.0" "databricks-sdk[openai]"

dbutils.library.restartPython()

{

　　"company": "Example Logistics A",

　　"events": [

　　　　{

　　　　　　"product": "Databricks",

　　　　　　"action": "adopt",

　　　　　　"status": "completed",

　　　　　　"purpose": "データ分析基盤"

　　　　}

　　]

}

from enum import Enum

from pydantic import BaseModel

class Action(str, Enum):

　　ADOPT = "adopt"

　　USE = "use"

　　MIGRATE = "migrate"

　　DISCONTINUE = "discontinue"

class Status(str, Enum):

　　PLANNED = "planned"

　　IN_PROGRESS = "in_progress"

　　COMPLETED = "completed"

class ITEvent(BaseModel):

　　product: str

　　action: Action

　　status: Status

　　purpose: str | None

class ExtractionResult(BaseModel):

　　company: str

　　events: list[ITEvent]

EXTRACTION_RESPONSE_SCHEMA = {

　　"type": "object",

　　"properties": {

　　　　"company": {

　　　　　　"type": "string",

　　　　},

　　　　"events": {

　　　　　　"type": "array",

　　　　　　"items": {

　　　　　　　　"type": "object",

　　　　　　　　"properties": {

　　　　　　　　　　"product": {

　　　　　　　　　　　　"type": "string",

　　　　　　　　　　},

　　　　　　　　　　"action": {

　　　　　　　　　　　　"type": "string",

　　　　　　　　　　　　"enum": [action.value for action in Action],

　　　　　　　　　　},

　　　　　　　　　　"status": {

　　　　　　　　　　　　"type": "string",

　　　　　　　　　　　　"enum": [status.value for status in Status],

　　　　　　　　　　},

　　　　　　　　　　"purpose": {

　　　　　　　　　　　　"type": ["string", "null"],

　　　　　　　　　　},

　　　　　　　　},

　　　　　　　　"required": ["product", "action", "status", "purpose"],

　　　　　　　　"additionalProperties": False,

　　　　　　},

　　　　},

　　},

　　"required": ["company", "events"],

　　"additionalProperties": False,

}

records = []

for i in range(1, 11):

　　case_id = f"case_{i:03d}"

　　case_dir = DATASET_ROOT / case_id

　　input_text = (case_dir / "input.txt").read_text(encoding="utf-8")

　　with open(

　　　　case_dir / "ground_truth.json",

　　　　encoding="utf-8"

　　) as f:

　　　　ground_truth = json.load(f)

　　records.append(

　　　　{

　　　　　　"inputs": {

　　　　　　　　"text": input_text

　　　　　　},

　　　　　　"expectations": {

　　　　　　　　"expected_response": ground_truth

　　　　　　},

　　　　　　"tags": {

　　　　　　　　"case_id": case_id

　　　　　　},

　　　　}

　　)

import mlflow.genai.datasets

DATASET_NAME = "workspace.default.it_product_event_eval"

eval_dataset = mlflow.genai.datasets.get_dataset(

　　name=DATASET_NAME

)

eval_dataset = eval_dataset.merge_records(records)

from mlflow.genai import register_prompt

USER_PROMPT = """

以下の文章から、IT製品・サービスに関するイベントを抽出してください。

文章:

{{text}}

"""

prompt = register_prompt(

　　name="workspace.default.it_product_event_extraction",

　　template=[

　　　　{

　　　　　　"role": "system",

　　　　　　"content": SYSTEM_PROMPT,

　　　　},

　　　　{

　　　　　　"role": "user",

　　　　　　"content": USER_PROMPT,

　　　　},

　　],

　　response_format=EXTRACTION_RESPONSE_SCHEMA,

)

prompt = mlflow.genai.load_prompt(

　　"prompts:/workspace.default.it_product_event_extraction/1"

)

from databricks.sdk import WorkspaceClient

workspace_client = WorkspaceClient()

client = workspace_client.serving_endpoints.get_open_ai_client()

def generate_structured(

　　messages: list[dict[str, Any]],

　　response_schema: dict[str, Any],

) -> dict[str, Any]:

　　response_format = {

　　　　"type": "json_schema",

　　　　"json_schema": {

　　　　　　"name": "extraction_result",

　　　　　　"schema": response_schema,

　　　　　　"strict": True,

　　　　}

　　}

　　response = client.chat.completions.create(

　　　　model="databricks-gpt-oss-120b",

　　　　messages=messages,

　　　　response_format=response_format,

　　　　temperature=0.0,

　　)

　　content = response.choices[0].message.content

　　if not content:

　　　　raise ValueError("LLM returned an empty response.")

　　if isinstance(content, str):

　　　　text = content

　　elif isinstance(content, list):

　　　　text_blocks = [

　　　　　　block["text"] for block in content

　　　　　　if block.get("type") == "text"

　　　　]

　　　　if len(text_blocks) != 1:

　　　　　　raise ValueError(

　　　　　　　　f"Expected exactly one text block, got {len(text_blocks)}."

　　　　　　)

　　　　text = text_blocks[0]

　　else:

　　　　raise ValueError(f"Unexpected response format: {content}")

　　try:

　　　　return json.loads(text)

　　except json.JSONDecodeError as e:

　　　　raise ValueError(

　　　　　　f"Failed to parse structured output as JSON: {content}"

　　　　) from e

@mlflow.trace

def extract(text: str) -> ExtractionResult:

　　raw = generate_structured(

　　　　messages= prompt.format(text=text),

　　　　response_schema=prompt.response_format,

　　)

　　try:

　　　　return ExtractionResult.model_validate(raw)

　　except ValidationError as exc:

　　　　raise ValueError(

　　　　　　f"LLM output does not match the extraction schema.\nRaw output: {raw}"

　　　　) from exc

from dataclasses import dataclass

@dataclass

class InputCase:

　　case_id: str

　　text: str

@dataclass

class DatasetCase(InputCase):

　　ground_truth: dict

def load_cases(

　　dataset_name: str,

) -> list[DatasetCase]:

　　dataset = mlflow.genai.datasets.get_dataset(

　　　　name=dataset_name,

　　)

　　df = dataset.to_df()

　　cases = []

　　for _, row in df.iterrows():

　　　　cases.append(

　　　　　　DatasetCase(

　　　　　　　　case_id=row["tags"]["case_id"],

　　　　　　　　text=row["inputs"]["text"],

　　　　　　　　ground_truth=row["expectations"]["expected_response"],

　　　　　　)

　　　　)

　　return cases

@mlflow.trace

def run_case(

　　case: InputCase,

　　run_type: str,

) -> dict[str, Any]:

　　mlflow.update_current_trace(

　　　　tags={

　　　　　　"case_id": case.case_id,

　　　　　　"run_type": run_type,

　　　　}

　　)

　　try:

　　　　result = extract(case.text)

　　　　return {

　　　　　　"case_id": case.case_id,

　　　　　　"success": True,

　　　　　　"prediction": result.model_dump(),

　　　　　　"error": None,

　　　　}

　　except Exception as exc:

　　　　return {

　　　　　　"case_id": case.case_id,

　　　　　　"success": False,

　　　　　　"prediction": None,

　　　　　　"error": str(exc),

　　　　}

def run_cases(

　　cases: list[InputCase],

　　run_type: str,

) -> list[dict[str, Any]]:

　　results = []

　　for case in cases:

　　　　print(f"Running: {case.case_id}")

　　　　result = run_case(

　　　　　　case=case,

　　　　　　run_type=run_type,

　　　　)

　　　　results.append(result)

　　return results

EXPERIMENT_NAME = f"/Users/{_USER}/llm-structured-extraction-evaluation"

mlflow.set_experiment(EXPERIMENT_NAME)

DATASET_NAME = "workspace.default.it_product_event_eval"

cases = load_cases(DATASET_NAME)

RUN_TYPE = "baseline_v1"

baseline_results = run_cases(

　　cases=cases,

　　run_type=RUN_TYPE,

)

"events": [

　　{

　　　　"product": "product A",

　　　　"action": "discontinue",

　　　　"status": "completed",

　　　　"purpose": "..."

　　},

　　{

　　　　"product": "product B",

　　　　"action": "migrate",

　　　　"status": "completed",

　　　　"purpose": "..."

　　}

]

def event_key(event: dict) -> tuple[str, str]:

　　return (

　　　　normalize_product(event["product"]),

　　　　event["action"],

　　)

def normalize_product(value: str) -> str:

　　return value.strip().lower()

from mlflow.genai.judges import make_judge

purpose_semantic_judge = make_judge(

　　name="purpose_semantic_match",

　　instructions="""

{{ outputs }} はLLMが抽出したIT製品・サービスの利用目的です。

{{ expectations }} はGround Truthの利用目的です。

2つの利用目的が意味的に同じ内容を表しているか判定してください。

表現や言い回し、詳細度が多少異なっていても、

本質的な利用目的が同じであれば "match" としてください。

利用目的の意味が異なる場合は "mismatch" としてください。

""".strip(),

　　feedback_value_type=Literal["match", "mismatch"],

　　model="databricks:/databricks-gpt-oss-120b",

)

def purpose_match(

　　prediction: str | None,

　　expected: str | None,

) -> bool:

　　if prediction is None and expected is None:

　　　　return True

　　if prediction is None or expected is None:

　　　　return False

　　feedback = purpose_semantic_judge(

　　　　outputs=prediction,

　　　　expectations={"expected_purpose": expected},

　　)

　　return feedback.value == "match"

from mlflow.entities import Feedback

from mlflow.genai.scorers import scorer

def evaluate_event(

　　prediction: dict,

　　expected: dict,

) -> EventEvaluation:

　　scores = FieldScores(

　　　　～省略（product, action, statusの評価処理）～

　　　　purpose=purpose_match(

　　　　　　prediction["purpose"],

　　　　　　expected["purpose"],

　　　　),

　　)

　　return EventEvaluation(

　　　　event_key=event_key(expected),

　　　　prediction=prediction,

　　　　expected=expected,

　　　　field_scores=scores,

　　)

def evaluate_case(

　　case_id: str,

　　prediction: dict,

　　ground_truth: dict,

) -> CaseEvaluation:

　　company_correct = exact_match(

　　　　prediction["company"],

　　　　ground_truth["company"],

　　)

　　～省略～

　　predictions_by_key = {}

　　for event in predicted_events:

　　　　key = event_key(event)

　　　　predictions_by_key.setdefault(key, []).append(event)

　　for expected in expected_events:

　　　　key = event_key(expected)

　　　　candidates = predictions_by_key.pop(key, [])

　　　　if not candidates:

　　　　　　false_negatives += 1

　　　　　　continue

　　　　# シンプルなマッチングルール:

　　　　# 1つ目の予測をマッチさせる、それ以外はFP

　　　　matched_prediction = candidates[0]

　　　　evaluation = evaluate_event(

　　　　　　prediction=matched_prediction,

　　　　　　expected=expected,

　　　　)

　　　　matched_events.append(evaluation)

　　　　false_positives += len(candidates) - 1

　　# 予測に含まれるが、期待値に含まれないEventはFP

　　for remaining in predictions_by_key.values():

　　　　false_positives += len(remaining)

　　return CaseEvaluation(

　　　　case_id=case_id,

　　　　company_correct=company_correct,

　　　　matched_events=matched_events,

　　　　false_positives=false_positives,

　　　　false_negatives=false_negatives,

　　)

def summarize_case(

　　evaluation: CaseEvaluation,

) -> dict[str, int]:

　　tp = 0

　　fp = evaluation.false_positives

　　fn = evaluation.false_negatives

　　～省略（正解数を0で初期化）～

　　for event in evaluation.matched_events:

　　　　product_correct += int(event.field_scores.product)

　　　　action_correct += int(event.field_scores.action)

　　　　status_correct += int(event.field_scores.status)

　　　　purpose_correct += int(event.field_scores.purpose)

　　　　if event.correct:

　　　　　　tp += 1

　　　　else:

　　　　　　# 対応するEventは存在したが、完全には正しくなかったケース

　　　　　　fp += 1

　　　　　　fn += 1

　　return {

　　　　"company_correct": int(evaluation.company_correct),

　　　　"event_tp": tp,

　　　　"event_fp": fp,

　　　　～省略（マッチしたイベント数や各フィールドの正解数を同様に返す）～

　　}

@scorer(name="structured_extraction")

def structured_extraction_scorer(

　　outputs: dict,

　　expectations: dict,

) -> list[Feedback]:

　　～省略～

　　evaluation = evaluate_case(

　　　　case_id=case_id,

　　　　prediction=prediction,

　　　　ground_truth=expectations[

　　　　　　"expected_response"

　　　　],

　　)

　　summary = summarize_case(evaluation)

　　return [

　　　　Feedback(

　　　　　　name="company_correct",

　　　　　　value=summary["company_correct"],

　　　　),

　　　　Feedback(

　　　　　　name="event_tp",

　　　　　　value=summary["event_tp"],

　　　　),

　　　　Feedback(

　　　　　　name="event_fp",

　　　　　　value=summary["event_fp"],

　　　　),

　　　　～省略（マッチしたイベント数や各フィールドの正解数を同様に返す）～

　　]

os.environ["MLFLOW_TRACING_SQL_WAREHOUSE_ID"] = "<SQL_WAREHOUSE_ID>"

experiment = mlflow.set_experiment(EXPERIMENT_NAME)

traces = mlflow.search_traces(

　　locations=[experiment.experiment_id],

　　filter_string=(

　　　　f"tag.run_type = '{run_type}'"

　　),

　　return_type="pandas",

)

cases_by_id = {

　　case.case_id: case for case in cases

}

traces["expectations"] = traces.apply(

　　lambda row: {

　　　　"expected_response":

　　　　　　cases_by_id[

　　　　　　　　row["tags"]["case_id"]

　　　　　　].ground_truth

　　},

　　axis=1,

)

with mlflow.start_run(

　　run_name="baseline_v1_evaluation"

):

　　evaluation_result = mlflow.genai.evaluate(

　　　　data=traces,

　　　　scorers=[structured_extraction_scorer],

　　)

result_df = evaluation_result.result_df

def aggregate_summaries(

　　summaries: list[dict[str, int]],

) -> dict[str, float | int]:

　　total_company_correct = sum(s["company_correct"] for s in summaries)

　　total_tp = sum(s["event_tp"] for s in summaries)

　　total_fp = sum(s["event_fp"] for s in summaries)

　　total_fn = sum(s["event_fn"] for s in summaries)

　　～省略（マッチしたイベント数や各フィールドの正解数に対し同様の処理を行う）～

　　precision = (

　　　　total_tp / (total_tp + total_fp)

　　　　if total_tp + total_fp

　　　　else 0.0

　　)

　　recall = (

　　　　total_tp / (total_tp + total_fn)

　　　　if total_tp + total_fn

　　　　else 0.0

　　)

　　f1 = (

　　　　2 * precision * recall / (precision + recall)

　　　　if precision + recall

　　　　else 0.0

　　)

　　def field_accuracy(

　　　　correct: int,

　　) -> float:

　　　　return (

　　　　　　correct / total_matched

　　　　　　if total_matched

　　　　　　else 0.0

　　　　)

　　return （集計結果等を返す）

metrics = aggregate_summaries(summaries)

metrics = aggregate_summaries(summaries)

mlflow.log_metrics(

　　{

　　　　"company_accuracy": metrics["company_accuracy"],

　　　　"event_tp": metrics["event_tp"],

　　　　"event_fp": metrics["event_fp"],

　　　　"event_fn": metrics["event_fn"],

　　　　"event_precision": metrics["event_precision"],

　　　　"event_recall": metrics["event_recall"],

　　　　"event_f1": metrics["event_f1"],

　　　　"product_accuracy": metrics["product_accuracy"],

　　　　"action_accuracy": metrics["action_accuracy"],

　　　　"status_accuracy": metrics["status_accuracy"],

　　　　"purpose_accuracy": metrics["purpose_accuracy"],

　　},

　　run_id=evaluation_result.run_id,

)

additional_cases = []

for i in range(11, 21):

　　case_id = f"case_{i:03d}"

　　case_dir = ADDITIONAL_CASES_PATH / case_id

　　input_text = (case_dir / "input.txt").read_text(

　　　　encoding="utf-8"

　　)

　　additional_cases.append(

　　　　InputCase(

　　　　　　case_id=case_id,

　　　　　　text=input_text,

　　　　)

　　)

RUN_TYPE = "additional_v1"

additional_results = run_cases(

　　cases=additional_cases,

　　run_type=RUN_TYPE,

)

confidence_judge = make_judge(

　　name="extraction_confidence",

　　instructions="""

{{ inputs }} は、IT製品・サービスのイベント抽出対象となった入力文章です。

{{ outputs }} は、その文章からLLMが抽出した構造化結果です。

入力文章と抽出結果を比較し、

以下の抽出ルールに基づいて、

抽出結果全体が正確かつ完全である確信度を0.0~1.0で評価してください。

【抽出ルール】

抽出対象となるイベントは、企業におけるIT製品・サービスの

採用、利用、移行、利用終了です。

各イベントには以下の情報が含まれます。

～Prompt v1とほぼ同じのため省略～

以下の場合は抽出対象に含めません。

- 検討、評価、比較、候補として言及されただけの製品

- 導入や移行の計画が取り下げられ、実際には採用・利用・移行されなかった製品

- 入力文章から確認できない、推測によるイベント

【評価観点】

以下を総合的に評価してください。

- 抽出対象となるイベントを漏れなく抽出できているか

- 抽出対象ではないイベントを抽出していないか

- productが適切か

- actionが適切か

- statusが適切か

- purposeが入力文章の内容と整合しているか

0.0は、

「抽出結果全体が正しい可能性が非常に低い」

ことを意味します。

1.0は、

「抽出結果全体が正しいと強く確信できる」

ことを意味します。

Ground Truthは与えられていません。

入力文章、抽出結果、および上記の抽出ルールだけに基づいて

Confidenceを評価してください。

評価理由（rationale）は日本語で記述してください。

""".strip(),

　　feedback_value_type=float,

　　model="databricks:/databricks-gpt-oss-120b",

)

confidence_results = []

for _, trace in traces.iterrows():

　　trace_id = trace["trace_id"]

　　case_id = trace["tags"]["case_id"]

　　request =json.loads(trace["request"])

　　response = json.loads(trace["response"])

　　judge_inputs = {

　　　　"text": request["case"]["text"],

　　}

　　judge_outputs = response["prediction"]

　　print(f"Scoring: {case_id}")

　　feedback = confidence_judge(

　　　　inputs=judge_inputs,

　　　　outputs=judge_outputs,

　　)

　　confidence = feedback.value

　　mlflow.log_assessment(

　　　　trace_id=trace_id,

　　　　assessment=feedback,

　　)

　　confidence_results.append(

　　　　{

　　　　　　"case_id": case_id,

　　　　　　"trace_id": trace_id,

　　　　　　"confidence": confidence,

　　　　　　"rationale": feedback.rationale,

　　　　}

　　)

REVIEWED_CASE_IDS = {

　　"case_014",

　　"case_016",

　　"case_018",

　　"case_020",

}

reviewed_traces = traces[

　　traces["tags"].apply(

　　　　lambda tags: tags["case_id"] in REVIEWED_CASE_IDS

　　)

].copy()

reviewed_traces["case_id"] = (

　　reviewed_traces["tags"].apply(

　　　　lambda tags: tags["case_id"]

　　)

)

reviewed_records = []

for _, row in reviewed_traces.iterrows():

　　case_id = row["case_id"]

　　trace_id = row["trace_id"]

　　trace = mlflow.get_trace(trace_id=trace_id)

　　expectations = [

　　　　assessment

　　　　for assessment in trace.info.assessments

　　　　if (

　　　　　　assessment.source.source_type == "HUMAN"

　　　　　　and assessment.expectation is not None

　　　　)

　　]

　　if len(expectations) != 1:

　　　　raise ValueError(

　　　　　　f"{case_id}: expected exactly one human expectation, got {len(expectations)}"

　　　　)

　　expected_response = json.loads(

　　　　expectations[0].expectation.value

　　)

　　validated = ExtractionResult.model_validate(expected_response)

　　request = json.loads(trace.data.request)

　　if isinstance(request, str):

　　　　request = json.loads(request)

　　reviewed_records.append(

　　　　{

　　　　　　"inputs": {

　　　　　　　　"text": request["case"]["text"],

　　　　　　},

　　　　　　"expectations": {

　　　　　　　　"expected_response": (

　　　　　　　　　　validated.model_dump(mode="json")

　　　　　　　　),

　　　　　　},

　　　　　　"tags": {

　　　　　　　　"case_id": case_id,

　　　　　　},

　　　　}

　　)

dataset = mlflow.genai.datasets.get_dataset(

　　name=DATASET_NAME

)

dataset.merge_records(reviewed_records)