From ef060ca6fbf7046148413c335aa0183f82f98736 Mon Sep 17 00:00:00 2001 From: Justin Xu Date: Mon, 12 Oct 2026 00:52:36 -0700 Subject: [PATCH] Forward annotation identity hints to array-record scoring --- README.md | 2 + .../evaluation/evaluators/extract.py | 2 + .../test_cases/identity_hints.py | 36 ++++++ src/extract_bench/test_cases/loader.py | 1 + src/extract_bench/test_cases/schema.py | 17 +++ .../test_cases/test_identity_hints.py | 106 ++++++++++++++++++ 6 files changed, 164 insertions(+) create mode 100644 src/extract_bench/test_cases/identity_hints.py create mode 100644 tests/extract_bench/test_cases/test_identity_hints.py diff --git a/README.md b/README.md index 62808be9..5d984aff 100644 --- a/README.md +++ b/README.md @@ -277,6 +277,8 @@ uv run extract-bench evaluation run output/ --test_cases_dir data uv run extract-bench analysis generate_report --evaluation_dir ./output/ ``` +For `.test.json` annotations, `_identity_hints` can declare row identity beside `data_schema`, for example `"_identity_hints": {"holdings": ["cusip"]}`. Array-record metrics use these hints to partition top-level arrays only when the residual assignment exceeds 100 million cells after exact-row matching. Each bucket keeps the same cap; smaller assignments remain keyless. Hints are evaluation-only and are not added to the schema sent to providers. +
diff --git a/src/extract_bench/evaluation/evaluators/extract.py b/src/extract_bench/evaluation/evaluators/extract.py index 85369354..333b0daa 100644 --- a/src/extract_bench/evaluation/evaluators/extract.py +++ b/src/extract_bench/evaluation/evaluators/extract.py @@ -37,6 +37,7 @@ from extract_bench.schemas.pipeline_io import InferenceResult from extract_bench.schemas.product import ProductType from extract_bench.test_cases.extract_field_paths import parse_field_path +from extract_bench.test_cases.identity_hints import top_level_identity_fields from extract_bench.test_cases.schema import ExtractFieldTestRule, ExtractTestCase, TestCase @@ -304,6 +305,7 @@ def compute_metrics(self, inference_result: InferenceResult, test_case: TestCase expected=expected_output, actual=extracted_data, data_schema=test_case.data_schema, + identity_keys_by_field=top_level_identity_fields(getattr(test_case, "identity_hints", None)), ) ) diff --git a/src/extract_bench/test_cases/identity_hints.py b/src/extract_bench/test_cases/identity_hints.py new file mode 100644 index 00000000..ebb4492e --- /dev/null +++ b/src/extract_bench/test_cases/identity_hints.py @@ -0,0 +1,36 @@ +"""Read evaluation-only row identity declarations from annotation sidecars.""" + +from collections.abc import Mapping +from typing import Any + + +def top_level_identity_fields(block: Mapping[str, Any] | None) -> Mapping[str, tuple[str, ...]]: + """Validate hints and select the top-level arrays scored by array-record F1. + + Nested array paths are retained by the test case for downstream scorers, + but array-record F1 only scores top-level arrays. Identity bucketing is + activated by the matcher only above its residual assignment cap. + """ + if block is None: + return {} + if not isinstance(block, Mapping): + raise ValueError("_identity_hints must be an object") + result: dict[str, tuple[str, ...]] = {} + for path, spec in block.items(): + if not isinstance(path, str) or any(len(part.strip()) == 0 for part in path.split(".")): + raise ValueError("_identity_hints paths must be non-empty dotted array names") + where = f"_identity_hints[{path!r}]" + if isinstance(spec, Mapping): + if set(spec) != {"fields"}: + raise ValueError(f"{where} must contain only 'fields'") + spec = spec["fields"] + fields = [spec] if isinstance(spec, str) else spec + if ( + not isinstance(fields, list) + or len(fields) == 0 + or any(not isinstance(field, str) or len(field.strip()) == 0 for field in fields) + ): + raise ValueError(f"{where} must be a field name or a non-empty list of field names") + if "." not in path: + result[path] = tuple(fields) + return result diff --git a/src/extract_bench/test_cases/loader.py b/src/extract_bench/test_cases/loader.py index 7738d14e..ab9e7a43 100644 --- a/src/extract_bench/test_cases/loader.py +++ b/src/extract_bench/test_cases/loader.py @@ -406,6 +406,7 @@ def load_test_case( test_rules=test_rules, **({"_schema_version": test_config["_schema_version"]} if "_schema_version" in test_config else {}), **_row_identity_kwargs(test_config, test_json_path), + identity_hints=test_config.get("_identity_hints"), ) # Drift check applies only to legacy rules (no v0.2 evidence list): diff --git a/src/extract_bench/test_cases/schema.py b/src/extract_bench/test_cases/schema.py index b8c826d4..72594672 100644 --- a/src/extract_bench/test_cases/schema.py +++ b/src/extract_bench/test_cases/schema.py @@ -8,6 +8,7 @@ from pydantic import BaseModel, ConfigDict, Discriminator, Field, Tag, field_validator from extract_bench.schemas.pipeline_io import InputFile +from extract_bench.test_cases.identity_hints import top_level_identity_fields from extract_bench.test_cases.parse_rule_schemas import ( ParseRule, coerce_parse_rule, @@ -439,6 +440,22 @@ class ExtractTestCase(BaseTestCase): ), ) + identity_hints: dict[str, Any] | None = Field( + default=None, + alias="_identity_hints", + description=( + "Evaluation-only identity fields by array path. Array-record F1 uses " + "top-level hints to partition assignments above its residual cell cap. " + "Keep these annotations outside the schema sent to inference providers." + ), + ) + + @field_validator("identity_hints", mode="before") + @classmethod + def validate_identity_hints(cls, value: Any) -> Any: + top_level_identity_fields(value) + return value + @property def eval_data_schema(self) -> dict[str, Any]: """``data_schema`` with row identity re-attached, for scoring only. diff --git a/tests/extract_bench/test_cases/test_identity_hints.py b/tests/extract_bench/test_cases/test_identity_hints.py new file mode 100644 index 00000000..5b38c508 --- /dev/null +++ b/tests/extract_bench/test_cases/test_identity_hints.py @@ -0,0 +1,106 @@ +"""Annotation loading and native evaluation of row identity hints.""" + +import json +from datetime import datetime +from pathlib import Path + +import pytest + +from extract_bench.evaluation.evaluators.extract import ExtractEvaluator +from extract_bench.evaluation.metrics.extract import array_record_match_metric +from extract_bench.schemas.extract_output import ExtractOutput +from extract_bench.schemas.pipeline_io import InferenceRequest, InferenceResult +from extract_bench.schemas.product import ProductType +from extract_bench.test_cases.identity_hints import top_level_identity_fields +from extract_bench.test_cases.loader import load_test_cases + +SCHEMA = { + "type": "object", + "properties": { + "rows": { + "type": "array", + "items": { + "type": "object", + "properties": {"id": {"type": "string"}, "amount": {"type": "number"}}, + }, + }, + }, +} +EXPECTED = {"rows": [{"id": "A", "amount": 1}, {"id": "B", "amount": 2}]} +ACTUAL = {"rows": [{"id": "B", "amount": 3}, {"id": "A", "amount": 4}]} + + +def write_case(directory: Path, hints: object) -> None: + (directory / "doc.pdf").touch() + (directory / "doc.test.json").write_text( + json.dumps({"data_schema": SCHEMA, "expected_output": EXPECTED, "_identity_hints": hints}) + ) + + +@pytest.mark.parametrize("hint", ["id", ["id"], {"fields": ["id"]}]) +@pytest.mark.parametrize("cap", [1, 100_000_000]) +def test_loaded_hints_reach_array_matcher( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch, hint: object, cap: int +) -> None: + write_case(tmp_path, {"rows": hint}) + (case,) = load_test_cases(str(tmp_path), product_type="extract") + assert case.identity_hints == {"rows": hint} + assert case.data_schema == SCHEMA + assert case.eval_data_schema == SCHEMA + monkeypatch.setattr(array_record_match_metric, "_MAX_RESIDUAL_ASSIGNMENT_CELLS", cap) + now = datetime.now() + result = InferenceResult( + request=InferenceRequest( + example_id="doc", source_file_path=str(tmp_path / "doc.pdf"), product_type=ProductType.EXTRACT + ), + pipeline_name="fixture", + product_type=ProductType.EXTRACT, + raw_output={}, + output=ExtractOutput(example_id="doc", pipeline_name="fixture", extracted_data=ACTUAL), + started_at=now, + completed_at=now, + latency_in_ms=1, + ) + evaluator = ExtractEvaluator() + hinted = {m.metric_name: m for m in evaluator.compute_metrics(result, case).metrics} + plain = { + m.metric_name: m + for m in evaluator.compute_metrics(result, case.model_copy(update={"identity_hints": None})).metrics + } + assert hinted["array_record_f1"].metadata["correct"] == 2 + assert hinted["array_record_f1"].value == 0.5 + assert plain["array_record_f1"].value == (0.0 if cap == 1 else 0.5) + + +@pytest.mark.parametrize( + "hints", + [ + [], + "rows", + {"rows": []}, + {"rows": [1]}, + {"rows": ""}, + {"rows": {"fields": ["id"], "typo": True}}, + {"": ["id"]}, + {"rows..nested": ["id"]}, + ], +) +def test_loader_rejects_malformed_hints(tmp_path: Path, hints: object) -> None: + write_case(tmp_path, hints) + with pytest.raises(ValueError, match="_identity_hints"): + load_test_cases(str(tmp_path), product_type="extract") + + +def test_nested_hints_are_preserved_but_not_forwarded_as_top_level(tmp_path: Path) -> None: + hints = {"rows": ["id"], "rows.children": {"fields": ["code", "date"]}} + write_case(tmp_path, hints) + (case,) = load_test_cases(str(tmp_path), product_type="extract") + assert case.identity_hints == hints + assert top_level_identity_fields(case.identity_hints) == {"rows": ("id",)} + + +@pytest.mark.parametrize("hints", [None, {}]) +def test_optional_hints(tmp_path: Path, hints: object) -> None: + write_case(tmp_path, hints) + (case,) = load_test_cases(str(tmp_path), product_type="extract") + assert top_level_identity_fields(case.identity_hints) == {}