RadGraph-IT 1.0.0__py3-none-any.whl

radgraphit/core/errors.py

@@ -0,0 +1,46 @@
+"""Domain exception hierarchy for predictable public loading and inference failures."""
+
+from __future__ import annotations
+
+
+class RadGraphITError(Exception):
+    """Base class for every error raised by radgraphit."""
+
+
+class InvalidReportError(RadGraphITError, ValueError):
+    """Input is not a non-empty ``str`` or non-empty sequence of non-empty ``str``.
+
+    Raised deterministically instead of the upstream package's behaviour of silently turning an
+    empty report into the literal string ``"None"``.
+    """
+
+
+class ModelConfigurationError(RadGraphITError):
+    """The model artifact could not be located or is misconfigured."""
+
+
+class ModelNotConfiguredError(ModelConfigurationError):
+    """The caller supplied an empty model identifier or revision."""
+
+
+class ModelNotFoundError(ModelConfigurationError):
+    """A configured model id / revision could not be downloaded from the Hugging Face Hub, or a
+    local model directory does not exist or is missing required files."""
+
+
+class ManifestError(ModelConfigurationError):
+    """``model_manifest.json`` is missing, malformed, incompatible, or fails checksum
+    verification."""
+
+
+class DeviceError(RadGraphITError):
+    """An invalid device string was requested, or CUDA was requested but is unavailable."""
+
+
+class CheckpointError(ModelConfigurationError):
+    """The checkpoint could not be loaded onto the reconstructed model topology (e.g. a
+    ``strict=True`` state-dict mismatch)."""
+
+
+class BackendError(RadGraphITError):
+    """The inference backend named by the manifest is unknown or unsupported."""

radgraphit/core/models.py

@@ -0,0 +1,116 @@
+"""Immutable, typed value objects that flow through the inference pipeline.
+
+These are plain ``@dataclass(frozen=True)`` objects (no torch, no I/O) so the pure boundary —
+tokenizer -> backend -> decoder -> serializer — can be reasoned about and unit-tested without a
+model. Token indices everywhere are **token-level, zero-based, and inclusive** (``end_ix`` is the
+last token *in* the span), exactly like RadGraph-XL; there are no character offsets.
+
+The typed protocols (:class:`InferenceBackend`) live here too, kept intentionally small: the only
+thing the service needs from a backend is "given tokens, give me raw scored predictions".
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Protocol, runtime_checkable
+
+
+@dataclass(frozen=True, slots=True)
+class TokenizedReport:
+    """A report after RadGraph-XL-compatible normalization + tokenization.
+
+    ``tokens`` are the normalized word tokens handed verbatim to the backend; ``text`` is their
+    space-joined form (the ``"text"`` field of the RadGraph-XL output).
+    """
+
+    tokens: tuple[str, ...]
+
+    @property
+    def text(self) -> str:
+        return " ".join(self.tokens)
+
+    def __len__(self) -> int:
+        return len(self.tokens)
+
+
+@dataclass(frozen=True, slots=True)
+class ScoredEntity:
+    """One raw NER prediction: an inclusive ``[start_ix, end_ix]`` span with a label and scores."""
+
+    start_ix: int
+    end_ix: int
+    label: str
+    raw_score: float
+    softmax_score: float
+
+    @property
+    def span(self) -> tuple[int, int]:
+        return (self.start_ix, self.end_ix)
+
+
+@dataclass(frozen=True, slots=True)
+class ScoredRelation:
+    """One raw relation prediction, **directed** source -> target.
+
+    Endpoints are referenced by span (not entity id); the serializer resolves them to entity ids.
+    """
+
+    label: str
+    source_start: int
+    source_end: int
+    target_start: int
+    target_end: int
+    raw_score: float
+    softmax_score: float
+
+    @property
+    def source_span(self) -> tuple[int, int]:
+        return (self.source_start, self.source_end)
+
+    @property
+    def target_span(self) -> tuple[int, int]:
+        return (self.target_start, self.target_end)
+
+
+@dataclass(frozen=True, slots=True)
+class ReportPrediction:
+    """The full raw graph predicted for a single report, before RadGraph-XL serialization.
+
+    This is the internal, typed result. The public API serializes it to the RadGraph-XL dict;
+    advanced callers can request these objects directly (see ``RadGraphIT.predict_graphs``).
+    """
+
+    doc_key: str
+    tokens: tuple[str, ...]
+    entities: tuple[ScoredEntity, ...]
+    relations: tuple[ScoredRelation, ...]
+
+    @property
+    def text(self) -> str:
+        return " ".join(self.tokens)
+
+
+@dataclass(frozen=True, slots=True)
+class RawBackendOutput:
+    """The minimal, torch-free contract a backend returns for one report.
+
+    NER tuples are ``(start, end, label, raw_score, softmax_score)`` and relation tuples are
+    ``(s1, e1, s2, e2, label, raw_score, softmax_score)`` — the same shape upstream DyGIE decode
+    produced, kept as plain Python so nothing torch-shaped leaks past the backend boundary.
+    """
+
+    ner: tuple[tuple[int, int, str, float, float], ...]
+    relations: tuple[tuple[int, int, int, int, str, float, float], ...]
+
+
+@runtime_checkable
+class InferenceBackend(Protocol):
+    """What :class:`~radgraphit.inference.service.InferenceService` needs from any backend.
+
+    Deliberately tiny: one report's tokens in, one :class:`RawBackendOutput` out. This is the seam
+    a fake backend is injected at in tests, and the boundary that keeps torch out of the service.
+    """
+
+    def infer(self, tokens: tuple[str, ...]) -> RawBackendOutput:
+        """Run inference on a single report's tokens and return raw scored predictions."""
+        ...

radgraphit/inference/__init__.py

@@ -0,0 +1,5 @@
+"""Inference orchestration: tokenizer, backend boundary, decoder, and the service.
+
+Importing this package is torch-free; the torch-dependent DyGIE-v2 backend under
+``backends/dygie_v2`` is imported lazily by the loader, never at package import time.
+"""

radgraphit/inference/backends/__init__.py

@@ -0,0 +1 @@
+"""Inference backends. The base protocol is torch-free; concrete backends may pull in torch."""

radgraphit/inference/backends/base.py

@@ -0,0 +1,13 @@
+"""The backend-facing surface of the inference protocol.
+
+A backend is anything that can turn one report's tokens into raw scored predictions. The protocol
+itself is defined in :mod:`radgraphit.core.models` (so it stays torch-free and importable
+everywhere); it is re-exported here as the canonical import for backend authors and for the
+service. Concrete backends (e.g. ``dygie_v2``) live in sibling packages and may depend on torch.
+"""
+
+from __future__ import annotations
+
+from ...core.models import InferenceBackend, RawBackendOutput
+
+__all__ = ["InferenceBackend", "RawBackendOutput"]

radgraphit/inference/backends/dygie_v2/__init__.py

@@ -0,0 +1,11 @@
+"""DyGIE-v2 backend: an inference-only PyTorch port of ``training_v2/src/dygie``.
+
+Private to the package — none of these classes are part of the public API. The parameter-bearing
+modules are faithful ports of the training code so a v2 checkpoint loads with ``strict=True``.
+Importing this package pulls in torch + transformers; it is imported lazily by the API facade.
+"""
+
+from .backend import DyGIEv2Backend
+from .loader import load_backend
+
+__all__ = ["DyGIEv2Backend", "load_backend"]

radgraphit/inference/backends/dygie_v2/backend.py

@@ -0,0 +1,32 @@
+"""The ``dygie_v2`` inference backend: a thin, torch-free-facing wrapper over the ported model.
+
+Implements the :class:`~radgraphit.core.models.InferenceBackend` protocol: tokens in, raw scored
+tuples out. All torch lives inside; nothing torch-shaped crosses ``infer``'s boundary.
+"""
+
+from __future__ import annotations
+
+import torch
+
+from ....core.errors import BackendError
+from ....core.models import RawBackendOutput
+from .model import DyGIEModel
+
+
+class DyGIEv2Backend:
+    """Runs a loaded :class:`DyGIEModel` on one report at a time."""
+
+    def __init__(self, model: DyGIEModel, device: torch.device):
+        self._model = model
+        self._device = device
+
+    @property
+    def device(self) -> torch.device:
+        return self._device
+
+    def infer(self, tokens: tuple[str, ...]) -> RawBackendOutput:
+        try:
+            ner, relations = self._model.predict(tokens)
+        except Exception as exc:
+            raise BackendError(f"dygie_v2 inference failed: {exc}") from exc
+        return RawBackendOutput(ner=tuple(ner), relations=tuple(relations))

radgraphit/inference/backends/dygie_v2/device.py

@@ -0,0 +1,46 @@
+"""Explicit device selection with clear errors and a working CPU default.
+
+Accepts ``"auto"`` (default), ``"cpu"``, ``"cuda"``, or ``"cuda:<index>"``. ``"auto"`` picks CUDA
+when available and otherwise falls back to CPU. An explicit CUDA request when CUDA is unavailable
+is an error (no silent fallback) — but CPU always works.
+"""
+
+from __future__ import annotations
+
+import torch
+
+from ....core.errors import DeviceError
+
+
+def resolve_device(device: str | None) -> torch.device:
+    """Resolve a device string to a concrete :class:`torch.device`.
+
+    Raises :class:`DeviceError` for an unknown string or a CUDA request on a machine without CUDA.
+    """
+    if device is None or device == "auto":
+        return torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    if not isinstance(device, str):
+        raise DeviceError(f"device must be a string or None, got {type(device).__name__}.")
+    if device == "cpu":
+        return torch.device("cpu")
+    if device == "cuda" or device.startswith("cuda:"):
+        if not torch.cuda.is_available():
+            raise DeviceError(
+                f"Device {device!r} was requested but CUDA is not available. Use device='cpu' or "
+                "device='auto' (which falls back to CPU)."
+            )
+        index = 0
+        if ":" in device:
+            suffix = device.split(":", 1)[1]
+            if not suffix.isdigit():
+                raise DeviceError(f"Invalid CUDA device index in {device!r}.")
+            index = int(suffix)
+            if index >= torch.cuda.device_count():
+                raise DeviceError(
+                    f"CUDA device index {index} out of range: only {torch.cuda.device_count()} "
+                    "CUDA device(s) visible."
+                )
+        return torch.device(f"cuda:{index}")
+    raise DeviceError(
+        f"Unknown device {device!r}. Expected one of: 'auto', 'cpu', 'cuda', 'cuda:<index>'."
+    )

radgraphit/inference/backends/dygie_v2/feedforward.py

@@ -0,0 +1,25 @@
+"""Shared MLP scorer body — faithful port of ``training_v2/src/dygie/feedforward.py``.
+
+Structure (the ``.net`` Sequential and its layer indices) is preserved exactly so the parameter
+names in the state dict match the checkpoint.
+"""
+
+from __future__ import annotations
+
+import torch
+from torch import nn
+
+
+class FeedForward(nn.Module):
+    def __init__(self, input_dim: int, hidden_dims: list[int], dropout: float):
+        super().__init__()
+        layers: list[nn.Module] = []
+        prev = input_dim
+        for dim in hidden_dims:
+            layers += [nn.Linear(prev, dim), nn.ReLU(), nn.Dropout(dropout)]
+            prev = dim
+        self.net = nn.Sequential(*layers)
+        self.output_dim = prev
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.net(x)

radgraphit/inference/backends/dygie_v2/loader.py

@@ -0,0 +1,165 @@
+"""Load a v2 checkpoint into a runnable backend.
+
+Reconstructs the model topology from ``config.json`` + ``vocab.json``, then loads the weights with
+``strict=True`` so any mismatch between the checkpoint and the reconstructed graph fails loudly and
+comprehensibly. For ``.pt`` checkpoints the safe ``torch.load(..., weights_only=True)`` path is
+used; ``.safetensors`` is supported (and preferred for a future release) when present.
+
+This is the only module that reads the config schema. It accepts every option v2 configs can carry
+— ``encoder``, ``max_span_width``, ``feature_size``, ``feedforward_params``, ``loss_weights`` (read
+but not needed at inference), ``relation_spans_per_word``, and, when present, ``span_pooling``,
+``transformer_params``, ``relation_context``, ``relation_feedforward_params``.
+"""
+
+from __future__ import annotations
+
+import json
+import logging
+from pathlib import Path
+from typing import Any
+
+import torch
+
+from ....core.errors import CheckpointError
+from .backend import DyGIEv2Backend
+from .device import resolve_device
+from .model import DyGIEModel
+from .tokenizer_embedder import MismatchedEmbedder
+from .vocab import Vocabulary, derive_dataset
+
+logger = logging.getLogger("radgraphit")
+
+
+def build_model(
+    config: dict[str, Any],
+    vocab: Vocabulary,
+    dataset: str,
+    *,
+    encoder_revision: str | None = None,
+    embedder: MismatchedEmbedder | None = None,
+) -> DyGIEModel:
+    """Reconstruct the model topology from a v2 config dict + loaded vocab.
+
+    ``embedder`` is injectable for testing (to avoid downloading an encoder); production leaves it
+    ``None`` so the default :class:`MismatchedEmbedder` is built from ``config["encoder"]``.
+    """
+    try:
+        encoder = config["encoder"]
+        model = DyGIEModel(
+            vocab,
+            dataset,
+            encoder_name=encoder["model_name"],
+            encoder_revision=encoder_revision,
+            max_length=encoder["max_length"],
+            max_span_width=config["max_span_width"],
+            feature_size=config["feature_size"],
+            feedforward_params=config["feedforward_params"],
+            relation_spans_per_word=config["relation_spans_per_word"],
+            span_pooling=config.get("span_pooling", False),
+            transformer_params=config.get("transformer_params"),
+            relation_context=config.get("relation_context", False),
+            relation_feedforward_params=config.get("relation_feedforward_params"),
+            train_encoder=encoder.get("train_parameters", True),
+            embedder=embedder,
+        )
+    except (KeyError, TypeError, ValueError) as exc:
+        raise CheckpointError(
+            f"config.json is missing or has an invalid required field: {exc}. A valid v2 config "
+            "needs at least "
+            "'encoder', 'max_span_width', 'feature_size', 'feedforward_params', and "
+            "'relation_spans_per_word'."
+        ) from exc
+    except Exception as exc:
+        raise CheckpointError(
+            f"Could not construct the encoder/model topology described by config.json: {exc}"
+        ) from exc
+    return model
+
+
+def _read_state_dict(weights_path: Path) -> dict[str, torch.Tensor]:
+    """Read a state dict from a ``.pt`` (safe, weights-only) or ``.safetensors`` checkpoint."""
+    if weights_path.suffix == ".safetensors":
+        try:
+            from safetensors.torch import load_file
+        except ImportError as exc:  # pragma: no cover - exercised only when safetensors is absent
+            raise CheckpointError(
+                f"{weights_path.name} is a safetensors checkpoint but the 'safetensors' package is "
+                "not installed. Install radgraphit[safetensors]."
+            ) from exc
+        try:
+            state_dict = load_file(str(weights_path))
+        except Exception as exc:
+            raise CheckpointError(
+                f"Could not read safetensors checkpoint {weights_path}: {exc}"
+            ) from exc
+        return state_dict
+
+    # .pt path: weights_only=True refuses to execute arbitrary pickled code (safe for untrusted
+    # bundles). The v2 trainer saves {"epoch": int, "model_state_dict": OrderedDict[str, Tensor]}.
+    try:
+        obj = torch.load(str(weights_path), map_location="cpu", weights_only=True)
+    except Exception as exc:
+        raise CheckpointError(f"Could not safely read checkpoint {weights_path}: {exc}") from exc
+    if isinstance(obj, dict) and "model_state_dict" in obj:
+        obj = obj["model_state_dict"]
+    if not isinstance(obj, dict) or not all(
+        isinstance(key, str) and isinstance(value, torch.Tensor) for key, value in obj.items()
+    ):
+        raise CheckpointError(
+            f"Unexpected checkpoint structure in {weights_path.name}: expected a mapping from "
+            "parameter names to tensors, optionally under 'model_state_dict'."
+        )
+    return obj
+
+
+def load_weights(model: DyGIEModel, weights_path: Path) -> None:
+    """Load weights into ``model`` with ``strict=True``, raising a clear error on mismatch."""
+    state_dict = _read_state_dict(weights_path)
+    try:
+        model.load_state_dict(state_dict, strict=True)
+    except Exception as exc:
+        raise CheckpointError(
+            f"Failed to load {weights_path.name} into the reconstructed model (strict=True): "
+            f"{exc}. This usually means config.json / vocab.json do not describe the same "
+            "architecture the checkpoint was trained with."
+        ) from exc
+
+
+def load_backend(
+    *,
+    config_path: Path,
+    vocab_path: Path,
+    weights_path: Path,
+    device: str = "auto",
+    encoder_revision: str | None = None,
+    embedder: MismatchedEmbedder | None = None,
+) -> DyGIEv2Backend:
+    """Build and return a ready-to-run :class:`DyGIEv2Backend` from a verified bundle's files."""
+    try:
+        config = json.loads(config_path.read_text(encoding="utf-8"))
+    except (OSError, UnicodeError, json.JSONDecodeError) as exc:
+        raise CheckpointError(
+            f"Could not read a valid JSON config at {config_path}: {exc}"
+        ) from exc
+    if not isinstance(config, dict):
+        raise CheckpointError(f"config.json at {config_path} must contain a JSON object.")
+    vocab = Vocabulary.load(vocab_path)
+    dataset = derive_dataset(vocab)
+    torch_device = resolve_device(device)
+
+    model = build_model(
+        config,
+        vocab,
+        dataset,
+        encoder_revision=encoder_revision,
+        embedder=embedder,
+    )
+    load_weights(model, weights_path)
+    try:
+        model.to(torch_device).eval()
+    except Exception as exc:
+        raise CheckpointError(
+            f"Could not move the reconstructed model to {torch_device}: {exc}"
+        ) from exc
+    logger.debug("Loaded dygie_v2 backend (dataset=%s) onto %s", dataset, torch_device)
+    return DyGIEv2Backend(model, torch_device)

radgraphit/inference/backends/dygie_v2/model.py

@@ -0,0 +1,101 @@
+"""The joint model — inference-only port of ``training_v2/src/dygie/model.py``.
+
+Wires ``embedder -> span_extractor -> ner + relation``. The four submodule names (``embedder``,
+``span_extractor``, ``ner``, ``relation``) and everything beneath them are preserved exactly, so
+``state_dict()`` keys line up with the training checkpoint and ``load_state_dict(strict=True)``
+succeeds. The training-only pieces (loss weighting, xavier re-init, metrics) are gone — every
+parameter is overwritten by the checkpoint at load time anyway.
+
+``embedder`` is injectable purely so the loader can be unit-tested without downloading an encoder;
+production always uses the default :class:`MismatchedEmbedder`.
+"""
+
+from __future__ import annotations
+
+import torch
+from torch import nn
+
+from .ner_head import NERHead, NerTuple
+from .relation_head import RelationHead, RelationTuple
+from .span_extractor import EndpointSpanExtractor
+from .spans import enumerate_spans
+from .tokenizer_embedder import MismatchedEmbedder
+from .vocab import Vocabulary
+
+
+class DyGIEModel(nn.Module):
+    def __init__(
+        self,
+        vocab: Vocabulary,
+        dataset: str,
+        *,
+        encoder_name: str,
+        encoder_revision: str | None,
+        max_length: int,
+        max_span_width: int,
+        feature_size: int,
+        feedforward_params: dict,
+        relation_spans_per_word: float,
+        span_pooling: bool = False,
+        transformer_params: dict | None = None,
+        relation_context: bool = False,
+        relation_feedforward_params: dict | None = None,
+        train_encoder: bool = True,
+        embedder: MismatchedEmbedder | None = None,
+    ):
+        super().__init__()
+        self.vocab = vocab
+        self.dataset = dataset
+        self.max_span_width = max_span_width
+
+        self.embedder = (
+            embedder
+            if embedder is not None
+            else MismatchedEmbedder(
+                encoder_name,
+                max_length,
+                train_encoder,
+                revision=encoder_revision,
+            )
+        )
+        self.span_extractor = EndpointSpanExtractor(
+            self.embedder.get_output_dim(),
+            num_width_embeddings=max_span_width,
+            span_width_embedding_dim=feature_size,
+            mean_pool=span_pooling,
+        )
+        span_emb_dim = self.span_extractor.get_output_dim()
+
+        self.ner = NERHead(vocab, span_emb_dim, feedforward_params, transformer_params)
+        self.relation = RelationHead(
+            vocab,
+            span_emb_dim,
+            self.embedder.get_output_dim(),
+            feedforward_params,
+            relation_spans_per_word,
+            transformer_params,
+            relation_context,
+            relation_feedforward_params,
+        )
+
+    @torch.no_grad()
+    def predict(self, tokens: tuple[str, ...]) -> tuple[list[NerTuple], list[RelationTuple]]:
+        """Run NER + relation inference on one report's tokens.
+
+        Returns ``(ner_predictions, relation_predictions)`` as plain tuples. Assumes a non-empty
+        token list (the service guarantees this).
+        """
+        self.eval()
+        words = list(tokens)
+        spans = enumerate_spans(words, self.max_span_width)
+        device = next(self.parameters()).device
+
+        word_embeddings = self.embedder(words)
+        spans_tensor = torch.tensor(spans, dtype=torch.long, device=device)
+        span_embeddings = self.span_extractor(word_embeddings, spans_tensor)
+
+        ner = self.ner.predict(self.dataset, spans, span_embeddings)
+        relations = self.relation.predict(
+            self.dataset, spans, span_embeddings, len(words), word_embeddings
+        )
+        return ner, relations

radgraphit/inference/backends/dygie_v2/ner_head.py

@@ -0,0 +1,75 @@
+"""NER head — inference-only port of ``training_v2/src/dygie/ner_head.py``.
+
+The parameter-bearing structure (``scorers`` + ``classifiers`` ModuleDicts, and the null-column
+trick: score ``n_labels - 1`` real classes, prepend a fixed-0 null column) is preserved exactly so
+the state dict matches the checkpoint. Loss and metrics are dropped — inference only needs the
+decode step, which returns plain ``(start, end, label, raw_score, softmax_score)`` tuples.
+"""
+
+from __future__ import annotations
+
+import torch
+from torch import nn
+
+from .feedforward import FeedForward
+from .transformer_block import TransformerBlock
+from .vocab import Vocabulary
+
+NerTuple = tuple[int, int, str, float, float]
+
+
+class NERHead(nn.Module):
+    def __init__(
+        self,
+        vocab: Vocabulary,
+        span_emb_dim: int,
+        feedforward_params: dict,
+        transformer_params: dict | None = None,
+    ):
+        super().__init__()
+        self.vocab = vocab
+        self.namespaces = [ns for ns in vocab.get_namespaces() if ns.endswith("ner_labels")]
+        self.n_labels = {ns: vocab.get_vocab_size(ns) for ns in self.namespaces}
+        self.use_transformer = transformer_params is not None
+
+        self.scorers = nn.ModuleDict()
+        self.classifiers = nn.ModuleDict()
+        for ns in self.namespaces:
+            if self.use_transformer:
+                assert transformer_params is not None
+                block = TransformerBlock(span_emb_dim, **transformer_params)
+                self.scorers[ns] = block
+                body_out = block.output_dim
+            else:
+                ff = FeedForward(
+                    span_emb_dim, feedforward_params["hidden_dims"], feedforward_params["dropout"]
+                )
+                self.scorers[ns] = ff
+                body_out = ff.output_dim
+            self.classifiers[ns] = nn.Linear(body_out, self.n_labels[ns] - 1)
+
+    def predict(
+        self, dataset: str, spans: list[tuple[int, int]], span_embeddings: torch.Tensor
+    ) -> list[NerTuple]:
+        ns = f"{dataset}__ner_labels"
+        if self.use_transformer:
+            starts = torch.tensor(
+                [s[0] for s in spans], dtype=torch.long, device=span_embeddings.device
+            )
+            hidden = self.scorers[ns](span_embeddings, starts)
+        else:
+            hidden = self.scorers[ns](span_embeddings)
+        scores = self.classifiers[ns](hidden)  # (num_spans, n_labels - 1)
+        dummy = scores.new_zeros(scores.size(0), 1)  # null-label score, fixed at 0
+        scores = torch.cat([dummy, scores], dim=-1)  # (num_spans, n_labels)
+
+        softmax_scores = torch.softmax(scores, dim=-1)
+        raw, predicted = scores.max(dim=-1)
+        soft, _ = softmax_scores.max(dim=-1)
+
+        predictions: list[NerTuple] = []
+        for i in (predicted != 0).nonzero(as_tuple=True)[0].tolist():
+            label = self.vocab.get_token_from_index(int(predicted[i]), ns)
+            start, end = spans[i]
+            predictions.append((int(start), int(end), label, float(raw[i]), float(soft[i])))
+        return predictions

radgraphit/inference/backends/dygie_v2/pruner.py

@@ -0,0 +1,27 @@
+"""Mention pruner — faithful port of ``training_v2/src/dygie/pruner.py``.
+
+Score every span, keep the top-k in original span order. The ``scorer`` submodule name is preserved
+for state-dict compatibility. batch_size=1 path: no masking, no batch dimension.
+"""
+
+from __future__ import annotations
+
+import torch
+from torch import nn
+
+
+class Pruner(nn.Module):
+    def __init__(self, scorer: nn.Module):
+        super().__init__()
+        self.scorer = scorer  # (num_spans, dim) -> (num_spans, 1)
+
+    def forward(
+        self, span_embeddings: torch.Tensor, num_items_to_keep: int
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        scores = self.scorer(span_embeddings).squeeze(-1)  # (num_spans,)
+        k = max(1, min(num_items_to_keep, span_embeddings.size(0)))
+        _, top_indices = scores.topk(k)
+        top_indices, _ = torch.sort(top_indices)
+        top_scores = scores[top_indices]
+        top_embeddings = span_embeddings[top_indices]
+        return top_embeddings, top_indices, top_scores