diff --git a/README.md b/README.md index 3a1820b..29d1585 100644 --- a/README.md +++ b/README.md @@ -83,6 +83,14 @@ Targeted allocation is experimental and remains opt-in because it conditions lat To avoid wasting judge calls on uninformative pairs, the judge skips comparisons where **both** outputs are shorter than `--min-chars` (default 20 — neither model produced meaningful text), and scores **identical** outputs as an automatic tie without calling the judge. Pass `--min-chars 0` to disable the length filter. +> [!IMPORTANT] +> If an existing results repository was judged with an older `ocr-bench` +> checkout from output configs containing multiple appended `inference_info` +> entries, run the first upgraded evaluation with `--full-rejudge`. Older +> versions could pair the latest model identity with an earlier OCR text +> column, and their incremental skip data does not record enough provenance to +> distinguish those verdicts safely. + **`ocr-bench view`** serves a local web viewer with a leaderboard, comparison browser, and human validation. Vote on comparisons to cross-check the automated judge with human judgement. ## Available models diff --git a/src/ocr_bench/dataset.py b/src/ocr_bench/dataset.py index 68238d5..08f76bb 100644 --- a/src/ocr_bench/dataset.py +++ b/src/ocr_bench/dataset.py @@ -327,12 +327,13 @@ def _load_configs( if revision: kwargs["revision"] = revision ds = load_dataset(**kwargs) + model_id, text_col = _resolve_config_output(ds, config) loaded.append( LoadedConfig( config=config, - model_id=_extract_model_id(ds, config), + model_id=model_id, ds=ds, - text_col=_find_text_column(ds), + text_col=text_col, ) ) return loaded @@ -547,48 +548,66 @@ def load_config_dataset( return unified, ocr_columns -def _extract_model_id(ds: Dataset, config: str) -> str: - """Extract model_id from inference_info in first row, falling back to config name. +def _select_inference_info_entry(ds: Dataset) -> dict | None: + """Choose the metadata entry that identifies this config's usable output. - Takes the *last* entry in the inference_info list, since OCR scripts append - new entries — the last one is the model that actually produced this config. + OCR scripts append entries. For list-valued metadata, choose the last entry + whose declared output column still exists; this keeps model identity and + output text coupled even when the newest entry is incomplete. Scalar legacy + metadata remains usable even when it relies on text-column heuristics. """ if "inference_info" not in ds.column_names: - return config + return None try: info_raw = ds["inference_info"][0] # column access avoids image decode - if info_raw: - info = json.loads(info_raw) - if isinstance(info, list): - info = info[-1] - return info.get("model_id", info.get("model_name", config)) + if not info_raw: + return None + info = json.loads(info_raw) + if isinstance(info, list): + for entry in reversed(info): + if not isinstance(entry, dict): + continue + column_name = entry.get("column_name", "") + if column_name and column_name in ds.column_names: + return entry + return None + if isinstance(info, dict): + return info except (json.JSONDecodeError, TypeError, KeyError, IndexError): pass - return config + return None + + +def _extract_model_id(ds: Dataset, config: str) -> str: + """Extract the model from the same inference entry as the output column.""" + info = _select_inference_info_entry(ds) + if info is None: + return config + return info.get("model_id", info.get("model_name", config)) def _find_text_column(ds: Dataset) -> str | None: """Find the likely OCR text column in a dataset. Priority: - 1. ``inference_info[0]["column_name"]`` if present and exists in dataset. + 1. The last ``inference_info`` entry whose ``column_name`` exists in the + dataset. OCR scripts append entries, and using one shared entry keeps + the selected model identity and output column in sync. 2. First column matching ``markdown`` (case-insensitive). 3. First column matching ``ocr`` (case-insensitive). 4. Column named exactly ``text``. """ - # Try inference_info first (column access avoids image decoding) - if "inference_info" in ds.column_names: - try: - info_raw = ds["inference_info"][0] - if info_raw: - info = json.loads(info_raw) - if isinstance(info, list): - info = info[0] - col_name = info.get("column_name", "") - if col_name and col_name in ds.column_names: - return col_name - except (json.JSONDecodeError, TypeError, KeyError, IndexError): - pass + info = _select_inference_info_entry(ds) + if info is not None: + col_name = info.get("column_name", "") + if col_name and col_name in ds.column_names: + return col_name + + return _find_text_column_by_heuristic(ds) + + +def _find_text_column_by_heuristic(ds: Dataset) -> str | None: + """Find a likely output column when metadata cannot identify one.""" # Prioritized heuristic: markdown > ocr > text for pattern in ["markdown", "ocr"]: @@ -600,6 +619,23 @@ def _find_text_column(ds: Dataset) -> str | None: return None +def _resolve_config_output(ds: Dataset, config: str) -> tuple[str, str | None]: + """Resolve a config's model label and text column as one atomic choice.""" + info = _select_inference_info_entry(ds) + if info is not None: + column_name = info.get("column_name", "") + if column_name and column_name in ds.column_names: + model_id = info.get("model_id", info.get("model_name", config)) + return model_id, column_name + + # Scalar legacy metadata often records only a model id. Preserve that + # behavior while using the established heuristic for its sole output. + model_id = info.get("model_id", info.get("model_name", config)) + return model_id, _find_text_column_by_heuristic(ds) + + return config, _find_text_column_by_heuristic(ds) + + # --------------------------------------------------------------------------- # Flat dataset loading # --------------------------------------------------------------------------- diff --git a/tests/test_dataset.py b/tests/test_dataset.py index de3ffd8..5bfd04e 100644 --- a/tests/test_dataset.py +++ b/tests/test_dataset.py @@ -197,6 +197,39 @@ def test_inference_info_list_format(self): ) assert _find_text_column(ds) == "markdown" + def test_inference_info_list_uses_final_appended_entry(self): + info = json.dumps( + [ + {"column_name": "old_ocr", "model_id": "org/old-model"}, + {"column_name": "new_ocr", "model_id": "org/new-model"}, + ] + ) + ds = Dataset.from_dict( + { + "image": [None], + "old_ocr": ["old output"], + "new_ocr": ["new output"], + "inference_info": [info], + } + ) + assert _find_text_column(ds) == "new_ocr" + + def test_inference_info_list_uses_last_entry_with_an_existing_column(self): + info = json.dumps( + [ + {"column_name": "old_ocr", "model_id": "org/old-model"}, + {"model_id": "org/new-model"}, + ] + ) + ds = Dataset.from_dict( + { + "image": [None], + "old_ocr": ["old output"], + "inference_info": [info], + } + ) + assert _find_text_column(ds) == "old_ocr" + # --------------------------------------------------------------------------- # discover_pr_configs @@ -306,6 +339,49 @@ def test_returns_all_when_no_default(self, mock_get): class TestLoadConfigDataset: + @patch("ocr_bench.dataset.load_dataset") + def test_appended_inference_info_keeps_model_and_output_in_sync(self, mock_load): + info = json.dumps( + [ + {"column_name": "old_ocr", "model_id": "org/old-model"}, + {"column_name": "new_ocr", "model_id": "org/new-model"}, + ] + ) + mock_load.return_value = Dataset.from_dict( + { + "image": [None, None], + "old_ocr": ["old one", "old two"], + "new_ocr": ["new one", "new two"], + "inference_info": [info, info], + } + ) + + ds, ocr_cols = load_config_dataset("repo/id", ["cfg"]) + + assert ocr_cols == {"cfg": "org/new-model"} + assert ds["cfg"] == ["new one", "new two"] + + @patch("ocr_bench.dataset.load_dataset") + def test_incomplete_final_entry_keeps_model_and_output_in_sync(self, mock_load): + info = json.dumps( + [ + {"column_name": "old_ocr", "model_id": "org/old-model"}, + {"model_id": "org/new-model"}, + ] + ) + mock_load.return_value = Dataset.from_dict( + { + "image": [None, None], + "old_ocr": ["old one", "old two"], + "inference_info": [info, info], + } + ) + + ds, ocr_cols = load_config_dataset("repo/id", ["cfg"]) + + assert ocr_cols == {"cfg": "org/old-model"} + assert ds["cfg"] == ["old one", "old two"] + @patch("ocr_bench.dataset.load_dataset") def test_merges_two_configs(self, mock_load): ds_a = Dataset.from_dict(