MT-Bench-101 (2/2): Add MT-Bench-101 golden-context evaluation runner

judgearena/benchmarks/mt_bench_101/evaluate.py

@@ -0,0 +1,251 @@
+"""Per-turn 1-10 judging and min-per-dialogue aggregation for MT-Bench-101."""
+
+from __future__ import annotations
+
+import re
+from functools import lru_cache
+from importlib.resources import files
+
+import pandas as pd
+from langchain_core.prompts import ChatPromptTemplate
+
+from judgearena.datasets.mt_bench_101 import (
+    MT_BENCH_101_REFERENCE_TASKS,
+    MT_BENCH_101_TASK_TO_ABILITY,
+)
+from judgearena.models import do_inference
+from judgearena.prompts.parsing import PairScore
+from judgearena.utils import safe_text, strip_thinking_tags
+
+DOUBLE_BRACKET_PATTERN = re.compile(r"\[\[(\d+)\]\]")
+
+TASK_PROMPT_FILES = {
+    "CM": "CM.txt",
+    "AR": "AR.txt",
+    "SI": "SI.txt",
+    "TS": "TS.txt",
+    "CC": "CC.txt",
+    "CR": "rephrasing.txt",
+    "FR": "rephrasing.txt",
+    "SC": "SC.txt",
+    "SA": "SA.txt",
+    "MR": "MR.txt",
+    "GR": "GR.txt",
+    "IC": "IC.txt",
+    "PI": "PI.txt",
+}
+
+
+def _prompt_text(name: str) -> str:
+    return (
+        files("judgearena.prompts")
+        .joinpath("templates", "mt_bench_101", name)
+        .read_text(encoding="utf-8")
+    )
+
+
+@lru_cache(maxsize=1)
+def load_mt_bench_101_prompts() -> dict[str, object]:
+    return {
+        "global_system": _prompt_text("global_system.txt"),
+        "scoring_format": _prompt_text("scoring_format.txt"),
+        "task_prompts": {
+            task: _prompt_text(prompt_file)
+            for task, prompt_file in TASK_PROMPT_FILES.items()
+        },
+    }
+
+
+def parse_mt_bench_101_rating(judge_completion: str) -> float | None:
+    for match in DOUBLE_BRACKET_PATTERN.finditer(judge_completion):
+        score = int(match.group(1))
+        if 1 <= score <= 10:
+            return float(score)
+    return None
+
+
+def format_mt_bench_101_dialogue(
+    *,
+    golden_context: list[dict[str, str]],
+    user_message: str,
+    assistant_message: str,
+) -> str:
+    chunks: list[str] = []
+    for turn in golden_context:
+        chunks.append(
+            f"\n\n Human: {turn.get('user', '')}\n\nAssistant: {turn.get('bot', '')}"
+        )
+    chunks.append(f"\n\n Human: {user_message}\n\nAssistant: {assistant_message}")
+    return "".join(chunks)
+
+
+def judge_mt_bench_101_single(
+    *,
+    judge_chat_model,
+    eval_items: pd.DataFrame,
+    completions: pd.DataFrame,
+    truncate_input_chars: int | None = 8192,
+    use_tqdm: bool = False,
+    strip_thinking_before_judging: bool = False,
+) -> pd.DataFrame:
+    prompts = load_mt_bench_101_prompts()
+    task_prompts = prompts["task_prompts"]
+    completion_by_idx = (
+        completions
+        if "instruction_index" not in completions.columns
+        else completions.set_index("instruction_index")
+    )
+    rows: list[dict[str, object]] = []
+    for idx in eval_items.index:
+        eval_row = eval_items.loc[idx]
+        completion_row = completion_by_idx.loc[idx]
+        task = str(eval_row["task"])
+        model_response = safe_text(completion_row.get("completion", ""), None)
+        if strip_thinking_before_judging:
+            model_response = strip_thinking_tags(model_response)
+        model_response = safe_text(model_response, truncate_input_chars)
+        dialogue = format_mt_bench_101_dialogue(
+            golden_context=list(eval_row.get("golden_context") or []),
+            user_message=safe_text(
+                eval_row.get("user_message", ""), truncate_input_chars
+            ),
+            assistant_message=model_response,
+        )
+        user_prompt = f"The dialogue need to be judged is: \n *** \n {dialogue} \n ***"
+        if task in MT_BENCH_101_REFERENCE_TASKS:
+            user_prompt += (
+                "\n\nThe reference solution is: \n ### \n "
+                f"{safe_text(eval_row.get('reference_answer'), truncate_input_chars)}"
+                " \n ###\n\n"
+            )
+        system_prompt = (
+            f"{prompts['global_system']}\n\n"
+            f"{task_prompts[task]}\n\n"
+            f"{prompts['scoring_format']}"
+        ).strip()
+        rows.append(
+            {
+                "instruction_index": idx,
+                "dialogue_id": eval_row["dialogue_id"],
+                "dialogue_uid": eval_row["dialogue_uid"],
+                "task": task,
+                "ability": eval_row.get("ability", MT_BENCH_101_TASK_TO_ABILITY[task]),
+                "domain": eval_row["domain"],
+                "turn_index": eval_row["turn_index"],
+                "model_completion": model_response,
+                "system_prompt": system_prompt,
+                "user_prompt": user_prompt,
+            }
+        )
+    prompt_template = ChatPromptTemplate.from_messages(
+        [("system", "{system_prompt}"), ("user", "{user_prompt}")]
+    )
+    inputs = prompt_template.batch(
+        [
+            {"system_prompt": row["system_prompt"], "user_prompt": row["user_prompt"]}
+            for row in rows
+        ]
+    )
+    judge_completions = do_inference(
+        chat_model=judge_chat_model,
+        inputs=inputs,
+        use_tqdm=use_tqdm,
+        stage="judging",
+    )
+    for row, judge_completion in zip(rows, judge_completions, strict=True):
+        row["judge_completion"] = judge_completion
+        row["score"] = parse_mt_bench_101_rating(judge_completion)
+    return pd.DataFrame(rows)
+
+
+def compute_mt_bench_101_dialogue_scores(scored_turns: pd.DataFrame) -> pd.DataFrame:
+    grouped = scored_turns.groupby(
+        ["dialogue_uid", "dialogue_id", "task", "ability", "domain"], as_index=False
+    )["score"].min()
+    return grouped.rename(columns={"score": "dialogue_score"})
+
+
+def summarize_mt_bench_101_absolute_scores(
+    scored_turns: pd.DataFrame,
+) -> dict[str, object]:
+    dialogue_scores = compute_mt_bench_101_dialogue_scores(scored_turns)
+    per_task_series = (
+        dialogue_scores.groupby("task")["dialogue_score"].mean().sort_index()
+    )
+    per_ability_series = (
+        dialogue_scores.groupby("ability")["dialogue_score"].mean().sort_index()
+    )
+    per_domain_series = (
+        dialogue_scores.groupby("domain")["dialogue_score"].mean().sort_index()
+    )
+    overall = per_task_series.mean() if len(per_task_series) else float("nan")
+    return {
+        "num_turns": int(len(scored_turns)),
+        "num_scored_turns": int(scored_turns["score"].notna().sum()),
+        "per_task": {
+            task: float(score)
+            for task, score in per_task_series.items()
+            if pd.notna(score)
+        },
+        "per_ability": {
+            ability: float(score)
+            for ability, score in per_ability_series.items()
+            if pd.notna(score)
+        },
+        "per_domain": {
+            domain: float(score)
+            for domain, score in per_domain_series.items()
+            if pd.notna(score)
+        },
+        "overall": float(overall) if pd.notna(overall) else None,
+    }
+
+
+def derive_mt_bench_101_pairwise_preferences(
+    scored_a: pd.DataFrame,
+    scored_b: pd.DataFrame,
+) -> pd.DataFrame:
+    cols = [
+        "instruction_index",
+        "dialogue_uid",
+        "dialogue_id",
+        "task",
+        "ability",
+        "domain",
+        "turn_index",
+    ]
+    merged = (
+        scored_a.loc[:, cols + ["score"]]
+        .rename(columns={"score": "score_A"})
+        .merge(
+            scored_b.loc[:, cols + ["score"]].rename(columns={"score": "score_B"}),
+            on=cols,
+            how="inner",
+        )
+    )
+    scorer = PairScore()
+    merged["preference"] = [
+        None
+        if pd.isna(score_a) or pd.isna(score_b)
+        else float(scorer.preference_from_scores(score_a, score_b))
+        for score_a, score_b in zip(merged["score_A"], merged["score_B"], strict=True)
+    ]
+    return merged
+
+
+def aggregate_mt_bench_101_dialogues(pairwise_turns: pd.DataFrame) -> pd.DataFrame:
+    """Apply the benchmark's minimum-turn rule before deriving preferences."""
+    group_columns = ["dialogue_uid", "dialogue_id", "task", "ability", "domain"]
+    dialogues = pairwise_turns.groupby(group_columns, as_index=False)[
+        ["score_A", "score_B"]
+    ].min()
+    scorer = PairScore()
+    dialogues["preference"] = [
+        None
+        if pd.isna(score_a) or pd.isna(score_b)
+        else float(scorer.preference_from_scores(score_a, score_b))
+        for score_a, score_b in zip(
+            dialogues["score_A"], dialogues["score_B"], strict=True
+        )
+    ]
+    return dialogues

judgearena/benchmarks/mt_bench_101/generate.py

@@ -0,0 +1,95 @@
+"""Golden-context completion generation for MT-Bench-101."""
+
+from __future__ import annotations
+
+from typing import Any
+
+import pandas as pd
+from langchain_core.prompts import ChatPromptTemplate
+
+from judgearena.models import do_inference, make_model
+from judgearena.utils import truncate
+
+DEFAULT_SYSTEM_PROMPT = "You are a helpful assistant."
+
+
+def _escape_template_braces(text: str) -> str:
+    return text.replace("{", "{{").replace("}", "}}")
+
+
+def _build_golden_context_input(
+    *,
+    system_prompt: str,
+    golden_context: list[dict[str, str]],
+    user_message: str,
+    truncate_input_chars: int | None,
+):
+    messages: list[tuple[str, str]] = [
+        ("system", _escape_template_braces(system_prompt))
+    ]
+    for turn in golden_context:
+        messages.append(
+            (
+                "user",
+                _escape_template_braces(
+                    truncate(str(turn.get("user") or ""), max_len=truncate_input_chars)
+                ),
+            )
+        )
+        messages.append(
+            (
+                "assistant",
+                _escape_template_braces(
+                    truncate(str(turn.get("bot") or ""), max_len=truncate_input_chars)
+                ),
+            )
+        )
+    messages.append(
+        (
+            "user",
+            _escape_template_braces(
+                truncate(user_message, max_len=truncate_input_chars)
+            ),
+        )
+    )
+    return ChatPromptTemplate.from_messages(messages).invoke({})
+
+
+def generate_mt_bench_101_completions(
+    eval_items: pd.DataFrame,
+    model: str,
+    truncate_input_chars: int | None = 8192,
+    max_tokens: int | None = 8192,
+    use_tqdm: bool = True,
+    system_prompt: str = DEFAULT_SYSTEM_PROMPT,
+    **model_kwargs: Any,
+) -> pd.DataFrame:
+    """Generate MT-Bench-101 responses from golden-context eval items."""
+    chat_model = make_model(model, max_tokens=max_tokens, **model_kwargs)
+    inputs = [
+        _build_golden_context_input(
+            system_prompt=system_prompt,
+            golden_context=list(row.get("golden_context") or []),
+            user_message=str(row.get("user_message") or ""),
+            truncate_input_chars=truncate_input_chars,
+        )
+        for _, row in eval_items.iterrows()
+    ]
+    completions = do_inference(
+        chat_model=chat_model,
+        inputs=inputs,
+        use_tqdm=use_tqdm,
+        stage="generation",
+    )
+    idxs = list(eval_items.index)
+    return pd.DataFrame(
+        {
+            "instruction_index": idxs,
+            "dialogue_id": [eval_items.loc[idx, "dialogue_id"] for idx in idxs],
+            "dialogue_uid": [eval_items.loc[idx, "dialogue_uid"] for idx in idxs],
+            "task": [eval_items.loc[idx, "task"] for idx in idxs],
+            "ability": [eval_items.loc[idx, "ability"] for idx in idxs],
+            "turn_index": [eval_items.loc[idx, "turn_index"] for idx in idxs],
+            "completion": completions,
+        }
+    )