Switch the NER reference parser to UltraBERT-150M (v5.1) and take its encoder from the checkpoint

CHANGELOG.md

@@ -28,6 +28,21 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/).
 
 ### Changed
 
+- The default reference parser is v5.1,
+  `scienceverse/bibr-parser-v5-1-ultrabert` (UltraBERT-150M encoder + CRF),
+  trained on the reference markup of CC BY / CC0 JATS papers, LinkedBooks and
+  the v4.5 gold rows. Against v4.5 it raises reference exact match from 0.42 to
+  0.75 on the refsets sealed test and from 0.56 to 0.73 on RenoBench, and edges
+  it on the v4.5 gold validation split (entity F1 0.975 vs 0.966). The parser
+  repo carries UltraBERT's config, modeling code and tokenizer, so the PyTorch
+  runtime loads them from the same pinned commit (`trust_remote_code` only at a
+  commit SHA or from a local directory, so `NER_PARSER_REVISION=main` does not
+  load it under PyTorch); the ONNX bundle runs no custom code. The encoder now
+  comes from the checkpoint (ModernBERT-base for a repo without a
+  `config.json`), and `NER_PARSER_ENCODER` / `NER_PARSER_ENCODER_REVISION` (new)
+  override it. `NER_PARSER_CKPT=scienceverse/bibr-parser-v4-5-gold` with
+  `NER_PARSER_REVISION=ff50a83e7f5b73dcf6f8f973a1a6e3847ec429e6` keeps v4.5.
+
 - Refreshed `uv.lock` to releases at least a week old (as of 2026-10-10):
   transformers 5.17.0, tokenizers 0.23.3, onnxruntime 1.31.0, anthropic
   0.125.0, google-genai 2.29.0, instructor 1.17.0, pydantic 2.14.0, pyarrow

bibr/config.py

@@ -2503,7 +2503,7 @@ def _normalize_serve_log_level(cls, value):
     # Reference extraction strategy — DECOUPLED segmentation + parsing axes
     # (resolved in bibr.extract.ref_extractor._resolve_ref_strategies; per-run
     # overrides ride RunConfig via chew(refs=...) / CLI --refs/--ref-seg).
-    # The local ModernBERT-CRF parser ("ner") is the default: reference parsing
+    # The local encoder-CRF parser ("ner") is the default: reference parsing
     # is the bulk of bibr's LLM tokens, so doing it locally cuts most of the
     # per-paper cost. To use LLM field parsing instead, set parse=llm (CLI
     # `--refs llm`). Segmentation defaults to the local geometry GBM ("geom"),
@@ -2521,7 +2521,7 @@ def _normalize_serve_log_level(cls, value):
     )
     REF_PARSE_STRATEGY: Literal["llm", "ner", "llm-chunked", "ref-llm", "off"] | None = Field(
         "ner",
-        description='Reference-parsing strategy: "ner" (default; local ModernBERT-CRF parser, no '
+        description='Reference-parsing strategy: "ner" (default; local encoder-CRF parser, no '
         'per-reference LLM cost), "ref-llm" (dedicated fine-tuned reference model), "llm" (batched LLM), "llm-chunked" (chunk-tolerant '
         'LLM parse over region-aligned chunks), or "off" (disable reference extraction entirely).',
     )
@@ -2701,13 +2701,15 @@ def _normalize_serve_log_level(cls, value):
         description="Local path or HF Hub repo id for the NER reference-segmenter checkpoint "
         "(optionally suffixed :filename.pt).",
     )
-    # Parser default is the v4.5 gold-trained feature-gated checkpoint (0.9656
-    # strict micro entity F1 on the gold val split vs 0.0544 for giant-v4, whose
-    # span conventions are OOD). The parser path (RefParser →
-    # FeatureGatedEncoderCRF, add_special_tokens=False) requires a v4-format
-    # checkpoint; the legacy v1 parser is no longer loadable here.
+    # Parser default is v5.1: UltraBERT-150M + feature-gated CRF trained on the
+    # reference markup of CC BY / CC0 JATS papers, LinkedBooks and the v4.5 gold
+    # rows (refsets sealed test ref_exact 0.75 vs 0.42 for v4.5, RenoBench 0.73
+    # vs 0.56, v4.5 gold val entity F1 0.975 vs 0.966). The repo carries the encoder's config, modeling code and
+    # tokenizer, so they load from it at NER_PARSER_REVISION. The parser path
+    # (RefParser → FeatureGatedEncoderCRF, add_special_tokens=False) requires a
+    # v4-format checkpoint; the legacy v1 parser is no longer loadable here.
     NER_PARSER_CKPT: str = Field(
-        "scienceverse/bibr-parser-v4-5-gold",
+        "scienceverse/bibr-parser-v5-1-ultrabert",
         description="Local path or HF Hub repo id for the NER reference-parser checkpoint "
         "(optionally suffixed :filename.pt). Requires a v4-format checkpoint.",
     )
@@ -2721,9 +2723,23 @@ def _normalize_serve_log_level(cls, value):
         "track the latest revision instead.",
     )
     NER_PARSER_REVISION: str = Field(
-        "ff50a83e7f5b73dcf6f8f973a1a6e3847ec429e6",
+        "97b06cc7dd6511bdd53f8855c93a29020231d215",
         description='Pinned commit revision for the NER parser checkpoint. Set to "main" to track '
-        "the latest revision instead.",
+        "the latest revision instead; a parser repo that ships its own encoder code (v5) loads "
+        "only from a commit SHA or a local directory.",
+    )
+    # Unset, the encoder comes from the checkpoint (bibr.ner.parser.resolve_encoder):
+    # the parser repo itself when it ships a config.json, else ModernBERT-base.
+    NER_PARSER_ENCODER: str | None = Field(
+        None,
+        description="Encoder (HF Hub repo id or local directory) whose config, modeling code and "
+        "tokenizer the torch NER parser loads. Unset: the checkpoint's own (a parser repo with a "
+        "config.json, at NER_PARSER_REVISION), else answerdotai/ModernBERT-base.",
+    )
+    NER_PARSER_ENCODER_REVISION: str | None = Field(
+        None,
+        description="Revision of NER_PARSER_ENCODER. An encoder with its own modeling code runs it "
+        "only from a local directory or a commit SHA revision.",
     )
     NER_DEVICE: str | None = Field(
         None, description="Device for NER checkpoints: cpu, cuda, or mps (auto-detected if unset)."

bibr/demo/server.py

@@ -59,7 +59,7 @@ def build_parser() -> argparse.ArgumentParser:
         default=None,
         help=(
             "Reference extraction strategy: 'ner' (default) parses each reference "
-            "with the local ModernBERT-CRF model (no per-reference LLM cost); 'llm' "
+            "with the local encoder-CRF model (no per-reference LLM cost); 'llm' "
             "parses the bibliography with the configured LLM. "
             "Overrides REF_PARSE_STRATEGY for the demo."
         ),

bibr/extract/ref_extractor.py

@@ -3,7 +3,7 @@
 ``ReferenceExtractor`` turns a located references DataFrame into
 ``PaperReference`` objects: segmentation (geom GBM → region anchors → LLM
 anchor-emit → CRF cascade, plus marker-split last resort) followed by parsing
-(batched LLM or local ModernBERT-CRF NER, dispatched through
+(batched LLM or local encoder-CRF NER, dispatched through
 ``REF_PARSE_STRATEGIES``).
 Module-level helpers cover the strategy-independent finalize steps (DOI/author
 rescue, vol/issue normalization, bib_type inference, bib_id sequencing).
@@ -583,7 +583,13 @@ def _get_ner_parser(settings: GlobalSettings | None = None, memory_mode: str | N
     global _NER_PARSER, _NER_PARSER_KEY
     settings = settings if settings is not None else snapshot_settings()
     device = resolve_ner_device(settings, memory_mode)
-    key = (settings.NER_PARSER_CKPT, device, settings.NER_PARSER_REVISION)
+    key = (
+        settings.NER_PARSER_CKPT,
+        device,
+        settings.NER_PARSER_REVISION,
+        settings.NER_PARSER_ENCODER,
+        settings.NER_PARSER_ENCODER_REVISION,
+    )
     if _NER_PARSER is None or key != _NER_PARSER_KEY:
         with _NER_LOCK:
             if _NER_PARSER is None or key != _NER_PARSER_KEY:
@@ -970,8 +976,8 @@ def _backfill_vancouver_numbers(fields: dict[str, Any], segment: str) -> None:
 
     Fill-only: a field the parser already populated survives untouched, except
     when it holds the undecomposed span (see ``_is_lumped_span``). Interim
-    measure for a ref-parse training-distribution gap — the v4.5-gold parser is
-    APA-heavy and tags this tail ``O``.
+    measure for a ref-parse training-distribution gap — the v4.5-gold parser was
+    APA-heavy and tagged this tail ``O``.
     """
     match = _match_vancouver_tail(fields, segment)
     if match is None:
@@ -1664,7 +1670,7 @@ async def _parse_refs_via_ner(
     ref_text: str,  # noqa: ARG001 — uniform strategy signature
     ref_strings: list[str],
 ) -> list[PaperReference]:
-    # Offload the synchronous ModernBERT-CRF forward pass to a worker thread so
+    # Offload the synchronous encoder-CRF forward pass to a worker thread so
     # it does not block the shared serve event loop (BibrPipelineAPI runs
     # enable_async=True, one loop for many concurrent requests).
     if (
@@ -2554,7 +2560,7 @@ def _crf_segment_or_recover(self, ref_text: str) -> list[str]:
         """CRF last-resort segmentation, guarded so a populated references region
         never silently yields 0.
 
-        Synchronous: a ModernBERT-CRF load plus N sliding-window forward passes
+        Synchronous: an encoder-CRF load plus N sliding-window forward passes
         holding the process-wide inference lock. Every caller offloads it with
         ``asyncio.to_thread``, matching ``_segment_geom`` and
         ``_parse_refs_via_ner`` — ``bibr serve`` runs one async worker, so this

bibr/extract/ref_field_repair.py

@@ -1,6 +1,6 @@
 """Deterministic repairs of NER reference fields (post-parse, before finalize).
 
-The NER reference parser (ModernBERT emissions + CRF) segments a reference
+The NER reference parser (encoder emissions + CRF) segments a reference
 correctly far more often than it places the field boundaries correctly. Its
 recurring boundary errors are systematic enough to repair from the reference
 text itself, without retraining:

bibr/local/cli/dry_run.py

@@ -57,7 +57,9 @@
     "scienceverse/bibr-geom-segmenter-v1": "1 MB",
     # answerdotai/ModernBERT-base (~150M params) encoder + CRF head.
     "scienceverse/bibr-segmenter-v1": "300 MB",
-    "scienceverse/bibr-parser-v4-5-gold": "300 MB",
+    # dmis-lab/UltraBERT-150M encoder + CRF head: best.pt (torch) or onnx/model.onnx
+    # (ONNX Runtime), 605 MB either way.
+    "scienceverse/bibr-parser-v5-1-ultrabert": "605 MB",
     # GLM-OCR weight repos — same figures as _OCR_BACKEND_DOWNLOAD_NOTES in
     # bibr/setup_wizard.py (bf16 ~2.7 GB, llama.cpp Q8_0 GGUF quant smaller).
     # Bare repo id — the ":Q8_0" quant suffix is stripped before lookup (see

bibr/local/cli/parser.py

@@ -235,7 +235,7 @@ def _add_pipeline_options(parser: argparse.ArgumentParser) -> None:
         "--refs",
         choices=["llm", "ner", "llm-chunked", "ref-llm", "off"],
         help=(
-            "Reference parser: 'ner' (default) uses the local ModernBERT-CRF model "
+            "Reference parser: 'ner' (default) uses the local encoder-CRF model "
             "through ONNX Runtime or the optional PyTorch runtime. "
             "'llm' parses references with the configured LLM in batches. "
             "'llm-chunked' parses region-aligned chunks of the reference list. "

bibr/ner/model.py

@@ -43,17 +43,26 @@ def __init__(
         num_tags: int,
         dropout: float = 0.0,
         bio_tags: list[str] | None = None,
+        revision: str | None = None,
+        trust_remote_code: bool = False,
     ) -> None:
         super().__init__()
         # Build the encoder from config only (no weight download). The base
         # pretrained weights would be immediately overwritten by the trained
         # checkpoint in load_state_dict, so loading them is wasted I/O and emits
         # a spurious transformers "UNEXPECTED keys" load report for the base
-        # MLM head we never use.
+        # MLM head we never use. An encoder with its own modeling code (UltraBERT)
+        # needs ``trust_remote_code``; ``code_revision`` keeps that code at the
+        # same pinned ``revision`` as its config (``from_config`` would otherwise
+        # fetch it from the repo's default branch).
         from bibr.utils.hf_cache import disable_hf_cache_symlinks_on_windows
 
         disable_hf_cache_symlinks_on_windows()
-        self.encoder = AutoModel.from_config(AutoConfig.from_pretrained(encoder_name))
+        config = AutoConfig.from_pretrained(
+            encoder_name, revision=revision, trust_remote_code=trust_remote_code
+        )
+        remote = {"trust_remote_code": True, "code_revision": revision} if trust_remote_code else {}
+        self.encoder = AutoModel.from_config(config, **remote)
         hidden_size: int = self.encoder.config.hidden_size
         self.dropout = nn.Dropout(dropout)
         self.linear = nn.Linear(hidden_size, num_tags)
@@ -102,12 +111,16 @@ def __init__(
         dropout: float = 0.0,
         bio_tags: list[str] | None = None,
         feature_dim: int = _FEATURE_DIM,
+        revision: str | None = None,
+        trust_remote_code: bool = False,
     ) -> None:
         super().__init__(
             encoder_name=encoder_name,
             num_tags=num_tags,
             dropout=dropout,
             bio_tags=bio_tags,
+            revision=revision,
+            trust_remote_code=trust_remote_code,
         )
         self.feature_dim = feature_dim
         self.feature_proj = nn.Linear(feature_dim, self.encoder.config.hidden_size)

bibr/ner/parser.py

@@ -1,8 +1,13 @@
-"""Trained ref parser (ModernBERT + feature-gated CRF), v4 checkpoint.
+"""Trained ref parser (encoder + feature-gated CRF), v4-format checkpoint.
 
 Takes a single reference string and emits BIO field tags, then maps the 39-tag
 parser output to a flat dict suitable for ``PaperReference.model_validate``.
 
+The encoder comes from the checkpoint (see :func:`resolve_encoder`): a parser
+repo that carries its encoder's ``config.json``, modeling code and tokenizer
+(v5, UltraBERT-150M) loads them from that repo at the parser's pinned revision;
+one that does not (v4.5) uses ModernBERT-base.
+
 The v4 checkpoint (``best.pt``) is a
 :class:`~bibr.ner.model.FeatureGatedEncoderCRF` trained with
 ``add_special_tokens=False`` and no tag shift. Its ``feature_proj.weight`` is
@@ -17,6 +22,7 @@
 from __future__ import annotations
 
 import logging
+import re
 from pathlib import Path
 
 from bibr.utils.ml_extra import ml_import_error
@@ -27,6 +33,7 @@
 except ImportError as e:  # pragma: no cover
     raise ml_import_error("NER reference parsing (bibr.ner)") from e
 
+from bibr.utils.ml_runtime import is_hub_repo_id
 from bibr.utils.text import strip_lone_surrogates
 
 from .checkpoint import resolve_checkpoint
@@ -38,6 +45,55 @@
 logger = logging.getLogger(__name__)
 
 DEFAULT_ENCODER = "answerdotai/ModernBERT-base"
+ENCODER_CONFIG = "config.json"
+_COMMIT_SHA = re.compile(r"[0-9a-f]{40}")
+
+
+def resolve_encoder(ckpt: str | Path, revision: str | None = None) -> tuple[str, str | None]:
+    """The encoder source (Hub id or local directory) for ``ckpt``, and its revision.
+
+    A checkpoint that ships its encoder's ``config.json`` next to the weights is
+    its own encoder source: a Hub repo at the parser's ``revision``, or the
+    checkpoint file's directory. Otherwise the encoder is ModernBERT-base, which
+    every checkpoint before v5 was trained on. A wrong guess cannot load
+    silently: the state dict is loaded strictly.
+    """
+    ckpt_str = str(ckpt)
+    repo_id, _, filename = ckpt_str.partition(":")
+    if isinstance(ckpt, Path) or not is_hub_repo_id(repo_id, filename=filename or None, warn=False):
+        local = Path(ckpt_str).expanduser()
+        folder = local if local.is_dir() else local.parent
+        if (folder / ENCODER_CONFIG).is_file():
+            return str(folder.resolve()), None
+        return DEFAULT_ENCODER, None
+    from bibr.utils.hf_cache import hf_download_or_cached
+
+    try:
+        hf_download_or_cached(repo_id, ENCODER_CONFIG, revision)
+    except Exception as exc:  # noqa: BLE001 — no encoder config: a pre-v5 checkpoint
+        logger.debug(
+            "%s@%s has no %s (%s); encoder %s",
+            repo_id,
+            revision or "main",
+            ENCODER_CONFIG,
+            type(exc).__name__,
+            DEFAULT_ENCODER,
+        )
+        return DEFAULT_ENCODER, None
+    return repo_id, revision
+
+
+def trusts_remote_code(encoder: str, revision: str | None) -> bool:
+    """Whether the encoder's own modeling code may run (``trust_remote_code``).
+
+    Only code that cannot change under bibr: a local directory, or a Hub repo
+    pinned to a commit SHA. A branch or tag revision (or none) is never
+    trusted, so an encoder that needs custom code then fails to load instead of
+    running whatever the repo holds today.
+    """
+    if not is_hub_repo_id(encoder, warn=False):
+        return True
+    return revision is not None and _COMMIT_SHA.fullmatch(revision) is not None
 
 
 class RefParser:
@@ -47,10 +103,11 @@ class RefParser:
     def __init__(
         self,
         ckpt_path: str | Path,
-        encoder_name: str = DEFAULT_ENCODER,
+        encoder_name: str | None = None,
         device: str | None = None,
         max_seq_len: int = 256,
         revision: str | None = None,
+        encoder_revision: str | None = None,
     ) -> None:
         if device is None:
             from bibr.utils.device import detect_torch_device
@@ -61,14 +118,22 @@ def __init__(
         from bibr.utils.hf_cache import disable_hf_cache_symlinks_on_windows
 
         disable_hf_cache_symlinks_on_windows()
-        self.tokenizer = AutoTokenizer.from_pretrained(encoder_name)
+        local_path = resolve_checkpoint(ckpt_path, revision=revision)
+        if encoder_name is None:
+            encoder_name, encoder_revision = resolve_encoder(ckpt_path, revision)
+        self.encoder_name = encoder_name
+        trust = trusts_remote_code(encoder_name, encoder_revision)
+        self.tokenizer = AutoTokenizer.from_pretrained(
+            encoder_name, revision=encoder_revision, trust_remote_code=trust
+        )
         self.model = FeatureGatedEncoderCRF(
             encoder_name=encoder_name,
             num_tags=len(BIO_TAGS),
             dropout=0.0,
             bio_tags=BIO_TAGS,
+            revision=encoder_revision,
+            trust_remote_code=trust,
         ).to(device)
-        local_path = resolve_checkpoint(ckpt_path, revision=revision)
         state = torch.load(local_path, map_location=device, weights_only=True)
         self.model.load_state_dict(state)
         self.model.eval()

bibr/ner/parser_onnx.py

@@ -1,4 +1,4 @@
-"""ONNX Runtime reference parser (ModernBERT emissions + numpy CRF decode).
+"""ONNX Runtime reference parser (encoder emissions + numpy CRF decode).
 
 Behaviourally identical to :class:`bibr.ner.parser.RefParser`: the same
 tokenizer settings (``add_special_tokens=False``, truncation to

bibr/ner/runtime.py

@@ -42,4 +42,10 @@ def load_ref_parser(
     import bibr.ner.parser as parser_mod
 
     logger.info("NER parser runtime: torch (%s@%s)", ckpt_str, revision or "main")
-    return parser_mod.RefParser(ckpt, device=device, revision=revision)
+    return parser_mod.RefParser(
+        ckpt,
+        encoder_name=getattr(settings, "NER_PARSER_ENCODER", None),
+        device=device,
+        revision=revision,
+        encoder_revision=getattr(settings, "NER_PARSER_ENCODER_REVISION", None),
+    )

bibr/setup_wizard.py

@@ -1670,7 +1670,7 @@ def _step_external_services(self) -> None:
         # Reference parsing strategy (mirrors the ``--refs`` chew flag)
         self.console.print("\n[dim]Reference parsing (bibliography → structured fields):[/dim]")
         self.console.print(
-            "  [cyan]ner[/cyan] — local ModernBERT-CRF parser (default): cuts most of\n"
+            "  [cyan]ner[/cyan] — local encoder-CRF parser (default): cuts most of\n"
             "        bibr's LLM token use (references are the bulk of it); reference\n"
             "        field precision is lower (requires the ml extra)\n"
             "  [cyan]llm[/cyan] — batched LLM parsing: full reference precision, at the\n"