microsoft/XL-DocBench

1#!/usr/bin/env python32"""Evaluate XL-DocBench predictions.3 4This script is intentionally self-contained for public release. It computes the5deterministic metrics used in the benchmark tables: relaxed Accuracy,6token-level F1, and ANLS. It does not call any model or require private files.7 8Prediction JSONL format:9    {"question_id": "adubench_single_000001", "prediction": "..."}10 11The prediction field may also be named ``model_answer``, ``answer``,12``response``, or ``output``. JSON files are also accepted, including mappings13from question_id to answer or internal-style ``{"items": {...}}`` files.14"""15 16from __future__ import annotations17 18import argparse19import csv20import json21import re22import sys23from collections import defaultdict24from dataclasses import dataclass, field25from pathlib import Path26from typing import Any27 28ANSWER_FORMAT_MAP = {29    "Str": "entity",30    "Int": "numeric",31    "Float": "numeric",32    "None": "unanswerable",33    "Bool": "boolean",34    "Boolean": "boolean",35    "Percentage": "percentage",36}37 38PREDICTION_FIELDS = ("prediction", "model_answer", "answer", "response", "output")39QUESTION_ID_FIELDS = ("question_id", "global_qa_id", "global_id", "id")40SUCCESS_STATUSES = {"success", "ok", "completed"}41 42 43@dataclass44class MetricBucket:45    accuracy: list[float] = field(default_factory=list)46    token_f1: list[float] = field(default_factory=list)47    anls: list[float] = field(default_factory=list)48 49    def add(self, accuracy: float, token_f1: float, anls: float) -> None:50        self.accuracy.append(accuracy)51        self.token_f1.append(token_f1)52        self.anls.append(anls)53 54    def summary(self) -> dict[str, float | int]:55        return {56            "count": len(self.accuracy),57            "accuracy": average(self.accuracy),58            "token_f1": average(self.token_f1),59            "anls": average(self.anls),60        }61 62 63def average(values: list[float]) -> float:64    return round(sum(values) / len(values), 6) if values else 0.065 66 67def load_jsonl(path: Path) -> list[dict[str, Any]]:68    rows: list[dict[str, Any]] = []69    with path.open("r", encoding="utf-8") as handle:70        for line_number, line in enumerate(handle, start=1):71            line = line.strip()72            if not line:73                continue74            try:75                value = json.loads(line)76            except json.JSONDecodeError as exc:77                raise ValueError(f"Invalid JSON on {path}:{line_number}") from exc78            if not isinstance(value, dict):79                raise TypeError(f"Expected object on {path}:{line_number}")80            rows.append(value)81    return rows82 83 84def load_json_or_jsonl(path: Path) -> Any:85    if path.suffix.lower() == ".jsonl":86        return load_jsonl(path)87    with path.open("r", encoding="utf-8") as handle:88        return json.load(handle)89 90 91def get_question_id(row: dict[str, Any]) -> str:92    for field_name in QUESTION_ID_FIELDS:93        value = row.get(field_name)94        if value is not None and str(value).strip():95            return str(value).strip()96    return ""97 98 99def string_value(value: Any) -> str:100    if value is None:101        return ""102    if isinstance(value, (str, int, float, bool)):103        return str(value)104    return json.dumps(value, ensure_ascii=False, sort_keys=True)105 106 107def answer_payload(row: dict[str, Any]) -> dict[str, Any]:108    value = row.get("answer", {})109    return value if isinstance(value, dict) else {"value": value}110 111 112def gold_answer(row: dict[str, Any]) -> str:113    return string_value(answer_payload(row).get("value", ""))114 115 116def answer_format(row: dict[str, Any]) -> str:117    payload = answer_payload(row)118    raw_format = string_value(payload.get("format", "Str")) or "Str"119    verification_rule = string_value(payload.get("verification_rule", ""))120    if verification_rule == "choice_exact_match":121        return "single_choice"122    if verification_rule == "percentage_exact":123        return "percentage"124    if raw_format in ANSWER_FORMAT_MAP:125        return ANSWER_FORMAT_MAP[raw_format]126    if "numeric" in verification_rule or "tolerance" in verification_rule:127        return "numeric"128    return raw_format.lower()129 130 131def metadata(row: dict[str, Any]) -> dict[str, Any]:132    value = row.get("metadata", {})133    return value if isinstance(value, dict) else {}134 135 136def load_gold_records(gold_files: list[Path]) -> dict[str, dict[str, Any]]:137    if not gold_files:138        raise ValueError("At least one gold file is required")139 140    records: dict[str, dict[str, Any]] = {}141    for path in gold_files:142        for row in load_jsonl(path):143            question_id = get_question_id(row)144            if not question_id:145                raise ValueError(f"Missing question_id in {path}")146            if question_id in records:147                raise ValueError(f"Duplicate question_id in gold data: {question_id}")148            records[question_id] = row149    if not records:150        raise ValueError("No gold records found")151    return records152 153 154def extract_prediction(row: Any, prediction_field: str = "") -> str:155    if not isinstance(row, dict):156        return string_value(row)157 158    if prediction_field:159        if prediction_field not in row:160            raise ValueError(f"Prediction field not found: {prediction_field}")161        value = row[prediction_field]162        if isinstance(value, dict) and "value" in value:163            return string_value(value["value"])164        return string_value(value)165 166    for field_name in PREDICTION_FIELDS:167        if field_name not in row:168            continue169        value = row[field_name]170        if field_name == "answer" and isinstance(value, dict):171            return string_value(value.get("value", ""))172        return string_value(value)173    return ""174 175 176def load_predictions(177    path: Path, prediction_field: str = ""178) -> tuple[dict[str, str], dict[str, str]]:179    payload = load_json_or_jsonl(path)180    predictions: dict[str, str] = {}181    statuses: dict[str, str] = {}182 183    def add(question_id: str, value: Any) -> None:184        if not question_id:185            raise ValueError(f"Prediction row is missing a question id: {value!r}")186        if question_id in predictions:187            raise ValueError(f"Duplicate question_id in predictions: {question_id}")188        predictions[question_id] = extract_prediction(value, prediction_field)189        if isinstance(value, dict):190            statuses[question_id] = (191                string_value(value.get("status", "success")) or "success"192            )193        else:194            statuses[question_id] = "success"195 196    if isinstance(payload, list):197        for row in payload:198            if not isinstance(row, dict):199                raise TypeError("Prediction JSONL/list rows must be objects")200            add(get_question_id(row), row)201    elif isinstance(payload, dict) and isinstance(payload.get("items"), dict):202        for question_id, row in payload["items"].items():203            add(str(question_id), row)204    elif isinstance(payload, dict) and get_question_id(payload):205        add(get_question_id(payload), payload)206    elif isinstance(payload, dict):207        for question_id, row in payload.items():208            add(str(question_id), row)209    else:210        raise TypeError("Unsupported prediction file format")211 212    return predictions, statuses213 214 215def normalize_answer(text: str) -> str:216    text = text.strip().lower()217    for prefix in ("the answer is", "answer:", "answer is"):218        if text.startswith(prefix):219            text = text[len(prefix) :].strip()220    text = re.sub(r"[^\w\s\.\-\%]", "", text)221    text = re.sub(r"\b(a|an|the)\b", " ", text)222    return re.sub(r"\s+", " ", text).strip()223 224 225def extract_number(text: str) -> float | None:226    compact = text.replace(" ", "")227 228    power_match = re.search(229        r"(?P<base>[-+]?(?:\d+(?:\.\d+)?|\.\d+))\^(?P<exponent>[-+]?\d+)",230        compact,231    )232    if power_match:233        try:234            return float(power_match.group("base")) ** int(235                power_match.group("exponent")236            )237        except (OverflowError, ValueError):238            return None239 240    fraction_match = re.search(241        r"(?P<numerator>[-+]?(?:\d+(?:\.\d+)?|\.\d+))/(?P<denominator>[-+]?(?:\d+(?:\.\d+)?|\.\d+))",242        compact,243    )244    if fraction_match:245        try:246            denominator = float(fraction_match.group("denominator"))247            if denominator == 0:248                return None249            return float(fraction_match.group("numerator")) / denominator250        except ValueError:251            return None252 253    match = re.search(254        r"[-+]?(?:\d{1,3}(?:,\d{3})+(?:\.\d+)?|\d+(?:[.,]\d+)?|\.\d+)(?:[eE][-+]?\d+)?",255        compact,256    )257    if not match:258        return None259    token = match.group()260    if "," in token and "." not in token:261        integer, fractional = token.split(",", maxsplit=1)262        token = (263            f"{integer}.{fractional}"264            if len(fractional) <= 2265            else token.replace(",", "")266        )267    else:268        token = token.replace(",", "")269    try:270        return float(token)271    except ValueError:272        return None273 274 275def parse_boolean(text: str) -> bool | None:276    tokens = set(normalize_answer(text).split())277    positive = bool(tokens & {"yes", "true", "correct"})278    negative = bool(tokens & {"no", "false", "incorrect"})279    if positive == negative:280        return None281    return positive282 283 284def levenshtein_distance(left: str, right: str) -> int:285    if len(left) < len(right):286        return levenshtein_distance(right, left)287    if not right:288        return len(left)289 290    previous_row = list(range(len(right) + 1))291    for left_index, left_char in enumerate(left):292        current_row = [left_index + 1]293        for right_index, right_char in enumerate(right):294            substitution_cost = 0 if left_char == right_char else 1295            current_row.append(296                min(297                    current_row[right_index] + 1,298                    previous_row[right_index + 1] + 1,299                    previous_row[right_index] + substitution_cost,300                )301            )302        previous_row = current_row303    return previous_row[-1]304 305 306def normalized_levenshtein_similarity(prediction: str, gold: str) -> float:307    prediction = prediction.strip().lower()308    gold = gold.strip().lower()309    if not prediction and not gold:310        return 1.0311    if not prediction or not gold:312        return 0.0313    distance = levenshtein_distance(prediction, gold)314    return 1.0 - distance / max(len(prediction), len(gold))315 316 317def anls_score(prediction: str, gold: str, threshold: float = 0.5) -> float:318    similarity = normalized_levenshtein_similarity(prediction, gold)319    return similarity if similarity >= threshold else 0.0320 321 322def accuracy_score(prediction: str, gold: str, answer_type: str) -> float:323    prediction_norm = normalize_answer(prediction)324    gold_norm = normalize_answer(gold)325 326    if answer_type == "unanswerable":327        phrases = (328            "not answerable",329            "unanswerable",330            "cannot be determined",331            "cannot be answered",332            "not enough information",333        )334        return 1.0 if any(phrase in prediction_norm for phrase in phrases) else 0.0335 336    if answer_type == "boolean":337        prediction_bool = parse_boolean(prediction)338        gold_bool = parse_boolean(gold)339        return (340            1.0 if prediction_bool is not None and prediction_bool == gold_bool else 0.0341        )342 343    if answer_type in {"numeric", "percentage"}:344        prediction_number = extract_number(prediction)345        gold_number = extract_number(gold)346        if prediction_number is None or gold_number is None:347            return 0.0348        if gold_number == 0:349            return 1.0 if abs(prediction_number) < 1e-6 else 0.0350        relative_error = abs(prediction_number - gold_number) / abs(gold_number)351        return 1.0 if relative_error <= 0.05 else 0.0352 353    if answer_type == "single_choice":354        prediction_option = re.search(r"\b([A-D])\b", prediction.strip().upper())355        gold_option = re.search(r"\b([A-D])\b", gold.strip().upper())356        return (357            1.0358            if prediction_option359            and gold_option360            and prediction_option.group(1) == gold_option.group(1)361            else 0.0362        )363 364    if gold_norm and gold_norm in prediction_norm:365        return 1.0366    if normalized_levenshtein_similarity(prediction_norm, gold_norm) >= 0.8:367        return 1.0368    return 0.0369 370 371def token_f1_score(prediction: str, gold: str) -> float:372    prediction_tokens = set(normalize_answer(prediction).split())373    gold_tokens = set(normalize_answer(gold).split())374    if not gold_tokens:375        return 1.0 if not prediction_tokens else 0.0376    if not prediction_tokens:377        return 0.0378    overlap = prediction_tokens & gold_tokens379    if not overlap:380        return 0.0381    precision = len(overlap) / len(prediction_tokens)382    recall = len(overlap) / len(gold_tokens)383    return 2 * precision * recall / (precision + recall)384 385 386def add_breakdown(387    breakdowns: dict[str, dict[str, MetricBucket]],388    name: str,389    key: Any,390    accuracy: float,391    token_f1: float,392    anls: float,393) -> None:394    label = string_value(key) or "unknown"395    breakdowns[name][label].add(accuracy, token_f1, anls)396 397 398def evaluate(399    gold_records: dict[str, dict[str, Any]],400    predictions: dict[str, str],401    statuses: dict[str, str],402    ignore_missing: bool = False,403) -> dict[str, Any]:404    overall = MetricBucket()405    breakdowns: dict[str, dict[str, MetricBucket]] = {406        "split": defaultdict(MetricBucket),407        "domain": defaultdict(MetricBucket),408        "reasoning_type": defaultdict(MetricBucket),409        "answer_format": defaultdict(MetricBucket),410        "difficulty": defaultdict(MetricBucket),411        "doc_type": defaultdict(MetricBucket),412        "evidence_source": defaultdict(MetricBucket),413    }414    per_question: list[dict[str, Any]] = []415    missing_count = 0416 417    for question_id, row in gold_records.items():418        if question_id not in predictions:419            missing_count += 1420            if ignore_missing:421                continue422        prediction = predictions.get(question_id, "")423        status = statuses.get(question_id, "missing")424 425        gold = gold_answer(row)426        answer_type = answer_format(row)427        normalized_status = (428            status.strip().casefold().replace("-", "_").replace(" ", "_")429        )430        if normalized_status in SUCCESS_STATUSES:431            accuracy = accuracy_score(prediction, gold, answer_type)432            token_f1 = token_f1_score(prediction, gold)433            anls = anls_score(prediction, gold)434        else:435            accuracy = 0.0436            token_f1 = 0.0437            anls = 0.0438        overall.add(accuracy, token_f1, anls)439 440        row_metadata = metadata(row)441        split = string_value(row.get("task_type", "unknown"))442        add_breakdown(breakdowns, "split", split, accuracy, token_f1, anls)443        add_breakdown(444            breakdowns, "domain", row_metadata.get("domain"), accuracy, token_f1, anls445        )446        add_breakdown(447            breakdowns,448            "reasoning_type",449            row_metadata.get("reasoning_type"),450            accuracy,451            token_f1,452            anls,453        )454        add_breakdown(455            breakdowns,456            "answer_format",457            answer_payload(row).get("format"),458            accuracy,459            token_f1,460            anls,461        )462        add_breakdown(463            breakdowns,464            "difficulty",465            row_metadata.get("difficulty"),466            accuracy,467            token_f1,468            anls,469        )470        add_breakdown(471            breakdowns,472            "doc_type",473            row_metadata.get("doc_type"),474            accuracy,475            token_f1,476            anls,477        )478 479        evidence_sources = row_metadata.get("evidence_sources") or ["unknown"]480        if not isinstance(evidence_sources, list):481            evidence_sources = [evidence_sources]482        for evidence_source in evidence_sources:483            add_breakdown(484                breakdowns, "evidence_source", evidence_source, accuracy, token_f1, anls485            )486 487        per_question.append(488            {489                "question_id": question_id,490                "prediction": prediction,491                "gold_answer": gold,492                "answer_format": answer_payload(row).get("format", "Str"),493                "status": status,494                "accuracy": round(accuracy, 6),495                "token_f1": round(token_f1, 6),496                "anls": round(anls, 6),497                "split": split,498                "domain": row_metadata.get("domain", "unknown"),499                "reasoning_type": row_metadata.get("reasoning_type", "unknown"),500            }501        )502 503    extra_prediction_count = len(set(predictions) - set(gold_records))504    return {505        "gold_count": len(gold_records),506        "prediction_count": len(predictions),507        "evaluated_count": overall.summary()["count"],508        "missing_prediction_count": missing_count,509        "extra_prediction_count": extra_prediction_count,510        "overall": overall.summary(),511        "breakdowns": {512            name: {key: bucket.summary() for key, bucket in sorted(group.items())}513            for name, group in breakdowns.items()514        },515        "per_question": per_question,516    }517 518 519def write_per_question_csv(rows: list[dict[str, Any]], output_path: Path) -> None:520    output_path.parent.mkdir(parents=True, exist_ok=True)521    fieldnames = [522        "question_id",523        "prediction",524        "gold_answer",525        "answer_format",526        "status",527        "accuracy",528        "token_f1",529        "anls",530        "split",531        "domain",532        "reasoning_type",533    ]534    with output_path.open("w", encoding="utf-8", newline="") as handle:535        writer = csv.DictWriter(handle, fieldnames=fieldnames)536        writer.writeheader()537        writer.writerows(rows)538 539 540def default_data_dir() -> Path:541    script_dir = Path(__file__).resolve().parent542    for data_dir in (script_dir / "data", script_dir.parent / "data"):543        if data_dir.exists():544            return data_dir545    return script_dir.parent / "data"546 547 548def parse_args() -> argparse.Namespace:549    parser = argparse.ArgumentParser(description="Evaluate XL-DocBench predictions")550    parser.add_argument(551        "--predictions", required=True, type=Path, help="Prediction JSON/JSONL file"552    )553    parser.add_argument(554        "--data-dir",555        type=Path,556        default=default_data_dir(),557        help="Directory containing QA JSONL files",558    )559    parser.add_argument(560        "--gold-files",561        nargs="+",562        type=Path,563        default=None,564        help="Gold QA JSONL files; defaults to qa_single_doc and qa_cross_doc",565    )566    parser.add_argument(567        "--prediction-field", default="", help="Optional explicit prediction field name"568    )569    parser.add_argument(570        "--ignore-missing",571        action="store_true",572        help="Evaluate only questions present in the prediction file",573    )574    parser.add_argument(575        "--output", type=Path, default=None, help="Write JSON report to this path"576    )577    parser.add_argument(578        "--per-question-csv",579        type=Path,580        default=None,581        help="Optional per-question CSV output",582    )583    parser.add_argument(584        "--no-per-question-json",585        action="store_true",586        help="Omit per-question rows from the JSON report",587    )588    return parser.parse_args()589 590 591def main() -> None:592    args = parse_args()593    gold_files = args.gold_files594    if gold_files is None:595        gold_files = [596            args.data_dir / "qa_single_doc.jsonl",597            args.data_dir / "qa_cross_doc.jsonl",598        ]599 600    missing_gold_files = [str(path) for path in gold_files if not path.exists()]601    if missing_gold_files:602        raise FileNotFoundError(f"Gold file(s) not found: {missing_gold_files}")603 604    gold_records = load_gold_records(gold_files)605    predictions, statuses = load_predictions(args.predictions, args.prediction_field)606    report = evaluate(607        gold_records, predictions, statuses, ignore_missing=args.ignore_missing608    )609 610    if args.per_question_csv:611        write_per_question_csv(report["per_question"], args.per_question_csv)612    if args.no_per_question_json:613        report = {key: value for key, value in report.items() if key != "per_question"}614 615    if args.output:616        args.output.parent.mkdir(parents=True, exist_ok=True)617        args.output.write_text(618            json.dumps(report, ensure_ascii=False, indent=2), encoding="utf-8"619        )620 621    overall = report["overall"]622    print("XL-DocBench evaluation")623    print(f"  gold questions:       {report['gold_count']}")624    print(f"  predictions:          {report['prediction_count']}")625    print(f"  evaluated:            {report['evaluated_count']}")626    print(f"  missing predictions:  {report['missing_prediction_count']}")627    print(f"  extra predictions:    {report['extra_prediction_count']}")628    print(f"  Accuracy:             {overall['accuracy'] * 100:.2f}")629    print(f"  Token F1:             {overall['token_f1'] * 100:.2f}")630    print(f"  ANLS:                 {overall['anls'] * 100:.2f}")631 632 633if __name__ == "__main__":634    try:635        main()636    except (OSError, TypeError, ValueError) as exc:637        print(f"ERROR: {exc}", file=sys.stderr)638        sys.exit(1)639