RadGraph-IT 1.0.0__py3-none-any.whl

radgraphit/inference/backends/dygie_v2/relation_head.py

@@ -0,0 +1,147 @@
+"""Relation head — inference-only port of ``training_v2/src/dygie/relation_head.py``.
+
+Prune to top-K spans, score every ordered pair, decode. The parameter-bearing structure
+(``pruners`` / ``span_transformers`` / ``relation_feedforwards`` / ``relation_scorers`` ModuleDicts,
+the null-column trick, and the optional ``span_pooling`` / ``transformer_params`` /
+``relation_context`` / ``relation_feedforward_params`` variants) is preserved exactly so the state
+dict matches the checkpoint. Loss and metrics are dropped; decode returns plain
+``(s1, e1, s2, e2, label, raw_score, softmax_score)`` tuples.
+"""
+
+from __future__ import annotations
+
+import math
+
+import torch
+from torch import nn
+
+from .feedforward import FeedForward
+from .pruner import Pruner
+from .transformer_block import TransformerBlock
+from .vocab import Vocabulary
+
+RelationTuple = tuple[int, int, int, int, str, float, float]
+
+
+class RelationHead(nn.Module):
+    def __init__(
+        self,
+        vocab: Vocabulary,
+        span_emb_dim: int,
+        word_emb_dim: int,
+        feedforward_params: dict,
+        spans_per_word: float,
+        transformer_params: dict | None = None,
+        use_between_context: bool = False,
+        relation_feedforward_params: dict | None = None,
+    ):
+        super().__init__()
+        self.vocab = vocab
+        self.spans_per_word = spans_per_word
+        self.namespaces = [ns for ns in vocab.get_namespaces() if ns.endswith("relation_labels")]
+        self.n_labels = {ns: vocab.get_vocab_size(ns) for ns in self.namespaces}
+        self.use_transformer = transformer_params is not None
+        self.use_between_context = use_between_context
+        between_dim = word_emb_dim if use_between_context else 0
+        relation_feedforward_params = relation_feedforward_params or feedforward_params
+
+        self.pruners = nn.ModuleDict()
+        self.span_transformers = nn.ModuleDict()
+        self.relation_feedforwards = nn.ModuleDict()
+        self.relation_scorers = nn.ModuleDict()
+
+        for ns in self.namespaces:
+            mention_ff = FeedForward(
+                span_emb_dim, feedforward_params["hidden_dims"], feedforward_params["dropout"]
+            )
+            mention_scorer = nn.Sequential(mention_ff, nn.Linear(mention_ff.output_dim, 1))
+            self.pruners[ns] = Pruner(mention_scorer)
+
+            if self.use_transformer:
+                assert transformer_params is not None
+                block = TransformerBlock(span_emb_dim, **transformer_params)
+                self.span_transformers[ns] = block
+                self.relation_scorers[ns] = nn.Linear(
+                    3 * block.output_dim + between_dim, self.n_labels[ns] - 1
+                )
+            else:
+                relation_ff = FeedForward(
+                    3 * span_emb_dim + between_dim,
+                    relation_feedforward_params["hidden_dims"],
+                    relation_feedforward_params["dropout"],
+                )
+                self.relation_feedforwards[ns] = relation_ff
+                self.relation_scorers[ns] = nn.Linear(relation_ff.output_dim, self.n_labels[ns] - 1)
+
+    @staticmethod
+    def _pair_embeddings(top_embeddings: torch.Tensor) -> torch.Tensor:
+        k = top_embeddings.size(0)
+        e1 = top_embeddings.unsqueeze(1).expand(k, k, -1)
+        e2 = top_embeddings.unsqueeze(0).expand(k, k, -1)
+        return torch.cat([e1, e2, e1 * e2], dim=-1)
+
+    @staticmethod
+    def _between_context(
+        word_embeddings: torch.Tensor, top_spans: list[tuple[int, int]]
+    ) -> torch.Tensor:
+        device, dtype = word_embeddings.device, word_embeddings.dtype
+        starts = torch.tensor([s[0] for s in top_spans], dtype=torch.long, device=device)
+        ends = torch.tensor([s[1] for s in top_spans], dtype=torch.long, device=device)
+
+        zero_row = word_embeddings.new_zeros(1, word_embeddings.size(-1))
+        prefix = torch.cat([zero_row, word_embeddings.cumsum(dim=0)], dim=0)  # (num_words + 1, dim)
+
+        low = torch.minimum(ends.unsqueeze(1), ends.unsqueeze(0))
+        high = torch.maximum(starts.unsqueeze(1), starts.unsqueeze(0))
+        count = (high - low - 1).clamp(min=0)
+
+        seg_sum = prefix[high] - prefix[low + 1]  # (K, K, dim)
+        mean = seg_sum / count.clamp(min=1).unsqueeze(-1).to(dtype)
+        return torch.where((count > 0).unsqueeze(-1), mean, seg_sum.new_zeros(()))
+
+    def predict(
+        self,
+        dataset: str,
+        spans: list[tuple[int, int]],
+        span_embeddings: torch.Tensor,
+        num_words: int,
+        word_embeddings: torch.Tensor,
+    ) -> list[RelationTuple]:
+        ns = f"{dataset}__relation_labels"
+        num_to_keep = math.ceil(num_words * self.spans_per_word)
+
+        top_embeddings, top_indices, top_scores = self.pruners[ns](span_embeddings, num_to_keep)
+        top_spans = [spans[i] for i in top_indices.tolist()]
+
+        if self.use_transformer:
+            starts = torch.tensor(
+                [s[0] for s in top_spans], dtype=torch.long, device=span_embeddings.device
+            )
+            top_embeddings = self.span_transformers[ns](top_embeddings, starts)
+
+        pair_embeddings = self._pair_embeddings(top_embeddings)
+        if self.use_between_context:
+            between = self._between_context(word_embeddings, top_spans)
+            pair_embeddings = torch.cat([pair_embeddings, between], dim=-1)
+
+        k = top_embeddings.size(0)
+        if self.use_transformer:
+            scores = self.relation_scorers[ns](pair_embeddings.view(k * k, -1)).view(k, k, -1)
+        else:
+            projected = self.relation_feedforwards[ns](pair_embeddings.view(k * k, -1))
+            scores = self.relation_scorers[ns](projected).view(k, k, -1)
+        scores = scores + top_scores.view(k, 1, 1) + top_scores.view(1, k, 1)
+        dummy = scores.new_zeros(k, k, 1)
+        scores = torch.cat([dummy, scores], dim=-1)  # (K, K, n_labels), null at index 0
+
+        raw, predicted = scores.max(dim=-1)
+        soft, _ = torch.softmax(scores, dim=-1).max(dim=-1)
+
+        predictions: list[RelationTuple] = []
+        for i, j in (predicted != 0).nonzero(as_tuple=False).tolist():
+            label = self.vocab.get_token_from_index(int(predicted[i, j]), ns)
+            (s1, e1), (s2, e2) = top_spans[i], top_spans[j]
+            predictions.append(
+                (int(s1), int(e1), int(s2), int(e2), label, float(raw[i, j]), float(soft[i, j]))
+            )
+        return predictions

radgraphit/inference/backends/dygie_v2/span_extractor.py

@@ -0,0 +1,51 @@
+"""Endpoint span extractor — faithful port of ``training_v2/src/dygie/span_extractor.py``.
+
+concat(start emb, end emb, [mean-pooled interior emb if ``mean_pool``], width emb). The
+``width_embedding`` submodule name is preserved for state-dict compatibility. Inference runs the
+batch_size=1 path (no batch dim, no span masking), exactly as upstream.
+"""
+
+from __future__ import annotations
+
+import torch
+from torch import nn
+
+
+class EndpointSpanExtractor(nn.Module):
+    def __init__(
+        self,
+        input_dim: int,
+        num_width_embeddings: int,
+        span_width_embedding_dim: int,
+        mean_pool: bool = False,
+    ):
+        super().__init__()
+        self.width_embedding = nn.Embedding(num_width_embeddings, span_width_embedding_dim)
+        self.mean_pool = mean_pool
+        self.output_dim = (3 if mean_pool else 2) * input_dim + span_width_embedding_dim
+
+    def get_output_dim(self) -> int:
+        return self.output_dim
+
+    def forward(self, word_embeddings: torch.Tensor, spans: torch.LongTensor) -> torch.Tensor:
+        """word_embeddings: (num_words, input_dim). spans: (num_spans, 2) inclusive [start, end]."""
+        starts, ends = spans[:, 0], spans[:, 1]
+        start_emb = word_embeddings[starts]
+        end_emb = word_embeddings[ends]
+        width_emb = self.width_embedding(ends - starts)  # 0-indexed width (widths 1..N -> 0..N-1)
+        if not self.mean_pool:
+            return torch.cat([start_emb, end_emb, width_emb], dim=-1)
+        mean_emb = self._interior_mean(word_embeddings, starts, ends)
+        return torch.cat([start_emb, end_emb, mean_emb, width_emb], dim=-1)
+
+    @staticmethod
+    def _interior_mean(
+        word_embeddings: torch.Tensor, starts: torch.Tensor, ends: torch.Tensor
+    ) -> torch.Tensor:
+        widths = ends - starts + 1  # (num_spans,)
+        max_width = int(widths.max().item())
+        offsets = torch.arange(max_width, device=word_embeddings.device).unsqueeze(0)  # (1, W)
+        mask = offsets < widths.unsqueeze(1)  # (num_spans, W)
+        idx = (starts.unsqueeze(1) + offsets).clamp(max=word_embeddings.size(0) - 1)
+        gathered = word_embeddings[idx] * mask.unsqueeze(-1)  # (num_spans, W, input_dim)
+        return gathered.sum(dim=1) / widths.unsqueeze(1).to(word_embeddings.dtype)

radgraphit/inference/backends/dygie_v2/spans.py

@@ -0,0 +1,20 @@
+"""Candidate span enumeration — the inference-relevant half of ``training_v2/src/dygie/dataset.py``.
+
+One document is a single "sentence" spanning the whole report (v2 invariant), so enumeration is the
+only piece of the dataset reader inference needs: all inclusive ``(start, end)`` spans of width
+1..``max_span_width``, ordered by ``(start, end)``.
+"""
+
+from __future__ import annotations
+
+
+def enumerate_spans(
+    words: tuple[str, ...] | list[str], max_span_width: int
+) -> list[tuple[int, int]]:
+    """All inclusive (start, end) spans of width 1..max_span_width, ordered by (start, end)."""
+    n = len(words)
+    spans: list[tuple[int, int]] = []
+    for start in range(n):
+        for end in range(start, min(start + max_span_width, n)):
+            spans.append((start, end))
+    return spans

radgraphit/inference/backends/dygie_v2/tokenizer_embedder.py

@@ -0,0 +1,123 @@
+"""Word-level embeddings from a subword encoder — faithful port of
+``training_v2/src/dygie/tokenizer_embedder.py``.
+
+Tokenize each word into subwords, run the transformer, mean-pool subwords back to one vector per
+word. Long documents are split into non-overlapping, word-aligned chunks of at most ``max_length``
+transformer positions ("context folding"), so a word's subwords are always encoded together.
+
+The ``encoder`` submodule is built from the pinned encoder configuration without downloading its
+pretrained weights: every encoder parameter is subsequently populated by the verified RadGraphIT
+checkpoint through ``load_state_dict(strict=True)``.
+"""
+
+from __future__ import annotations
+
+import logging
+
+import torch
+from torch import nn
+from transformers import AutoConfig, AutoModel, AutoTokenizer
+
+logger = logging.getLogger("radgraphit")
+
+
+class MismatchedEmbedder(nn.Module):
+    def __init__(
+        self,
+        model_name: str,
+        max_length: int,
+        train_parameters: bool = True,
+        *,
+        revision: str | None = None,
+    ):
+        super().__init__()
+        self.tokenizer = AutoTokenizer.from_pretrained(model_name, revision=revision)
+        encoder_config = AutoConfig.from_pretrained(model_name, revision=revision)
+        self.encoder = AutoModel.from_config(encoder_config)  # type: ignore[no-untyped-call]
+        self.max_length = max_length
+        self.output_dim = self.encoder.config.hidden_size
+        if not train_parameters:
+            for p in self.encoder.parameters():
+                p.requires_grad = False
+
+        num_specials = self.tokenizer.num_special_tokens_to_add(pair=False)
+        budget = max_length - num_specials
+        if budget <= 0:
+            raise ValueError(
+                f"max_length={max_length} too small for {model_name}'s "
+                f"{num_specials} special tokens"
+            )
+        self._budget = budget
+
+    def get_output_dim(self) -> int:
+        return self.output_dim
+
+    def _word_subword_counts(self, words: list[str]) -> list[int]:
+        enc = self.tokenizer(words, is_split_into_words=True, add_special_tokens=False)
+        counts = [0] * len(words)
+        for w in enc.word_ids():
+            if w is not None:
+                counts[w] += 1
+        return counts
+
+    def _pack_chunks(self, counts: list[int]) -> list[list[int]]:
+        """Greedily group word indices into chunks whose total subword count fits the budget."""
+        chunks: list[list[int]] = []
+        current: list[int] = []
+        current_len = 0
+        for w, n in enumerate(counts):
+            n = max(n, 1)
+            if n > self._budget:
+                logger.warning(
+                    "word %d has %d subwords > budget %d; it will be truncated by the tokenizer",
+                    w,
+                    n,
+                    self._budget,
+                )
+            if current and current_len + n > self._budget:
+                chunks.append(current)
+                current, current_len = [], 0
+            current.append(w)
+            current_len += n
+        if current:
+            chunks.append(current)
+        return chunks
+
+    def forward(self, words: list[str]) -> torch.Tensor:
+        """Returns (num_words, hidden_size)."""
+        device = next(self.encoder.parameters()).device
+        counts = self._word_subword_counts(words)
+        chunks = self._pack_chunks(counts)
+
+        word_embeddings: list[torch.Tensor | None] = [None] * len(words)
+        hidden: torch.Tensor | None = None
+        for chunk_word_ixs in chunks:
+            chunk_words = [words[i] for i in chunk_word_ixs]
+            enc = self.tokenizer(
+                chunk_words,
+                is_split_into_words=True,
+                add_special_tokens=True,
+                truncation=True,
+                max_length=self.max_length,
+                return_tensors="pt",
+            )
+            input_ids = enc["input_ids"].to(device)
+            attention_mask = enc["attention_mask"].to(device)
+            hidden = self.encoder(
+                input_ids=input_ids, attention_mask=attention_mask
+            ).last_hidden_state[0]
+
+            word_ids = enc.word_ids(batch_index=0)
+            positions_by_word: dict[int, list[int]] = {}
+            for pos, w in enumerate(word_ids):
+                if w is not None:
+                    positions_by_word.setdefault(w, []).append(pos)
+            for local_w, positions in positions_by_word.items():
+                word_embeddings[chunk_word_ixs[local_w]] = hidden[positions].mean(dim=0)
+
+        for i, emb in enumerate(word_embeddings):
+            if emb is None:  # word tokenized to zero subwords (rare; e.g. an unknown char)
+                assert hidden is not None  # at least one chunk always runs for a non-empty report
+                word_embeddings[i] = hidden.new_zeros(self.output_dim)
+
+        return torch.stack(word_embeddings, dim=0)  # type: ignore[arg-type]

radgraphit/inference/backends/dygie_v2/transformer_block.py

@@ -0,0 +1,62 @@
+"""Span-attention scorer body — faithful port of ``training_v2/src/dygie/transformer_block.py``.
+
+Used by the NER / relation heads in the ``medbit_transformer`` arm in place of the plain
+``FeedForward`` scorer. Submodule names (``input_proj``, ``encoder``) are preserved for state-dict
+compatibility. Spans are the "sequence"; a sinusoidal positional encoding keyed on each span's
+start-word index gives attention a notion of span order/distance.
+"""
+
+from __future__ import annotations
+
+import math
+
+import torch
+from torch import nn
+
+
+def sinusoidal_position_encoding(positions: torch.Tensor, d_model: int) -> torch.Tensor:
+    """positions: (N,) int64 span start indices. Returns (N, d_model)."""
+    div_term = torch.exp(
+        torch.arange(0, d_model, 2, device=positions.device, dtype=torch.float32)
+        * (-math.log(10000.0) / d_model)
+    )
+    angles = positions.unsqueeze(1).to(torch.float32) * div_term.unsqueeze(0)  # (N, d_model/2)
+    pe = torch.zeros(positions.size(0), d_model, device=positions.device)
+    pe[:, 0::2] = torch.sin(angles)
+    pe[:, 1::2] = torch.cos(angles[:, : pe[:, 1::2].size(1)])
+    return pe
+
+
+class TransformerBlock(nn.Module):
+    def __init__(
+        self,
+        input_dim: int,
+        d_model: int,
+        nhead: int,
+        num_layers: int,
+        dim_feedforward: int,
+        dropout: float,
+        activation: str = "gelu",
+    ):
+        super().__init__()
+        self.input_proj = nn.Linear(input_dim, d_model) if input_dim != d_model else nn.Identity()
+        layer = nn.TransformerEncoderLayer(
+            d_model=d_model,
+            nhead=nhead,
+            dim_feedforward=dim_feedforward,
+            dropout=dropout,
+            activation=activation,
+            norm_first=True,
+            batch_first=True,
+        )
+        self.encoder = nn.TransformerEncoder(
+            layer, num_layers=num_layers, enable_nested_tensor=False
+        )
+        self.d_model = d_model
+        self.output_dim = d_model
+
+    def forward(self, span_embeddings: torch.Tensor, start_positions: torch.Tensor) -> torch.Tensor:
+        x = self.input_proj(span_embeddings) + sinusoidal_position_encoding(
+            start_positions, self.d_model
+        )
+        return self.encoder(x.unsqueeze(0)).squeeze(0)

radgraphit/inference/backends/dygie_v2/vocab.py

@@ -0,0 +1,123 @@
+"""Minimal label vocabulary — inference-only port of ``training_v2/src/dygie/vocab.py``.
+
+Per-namespace string<->int maps with the null label ``""`` pinned to index 0 (so "argmax == 0"
+means "no label"). Namespaces are ``"{dataset}__ner_labels"`` / ``"{dataset}__relation_labels"``;
+the dataset prefix is *derived from the saved vocab*, never hard-coded, so the label schema is
+whatever the checkpoint was trained with (the current v2 artifact: 6 NER classes, 3 relations).
+"""
+
+from __future__ import annotations
+
+import json
+from pathlib import Path
+
+from ....core.errors import CheckpointError
+
+NULL_LABEL = ""
+NER_SUFFIX = "__ner_labels"
+RELATION_SUFFIX = "__relation_labels"
+
+
+class Vocabulary:
+    """String<->index maps per namespace, loaded from the checkpoint's ``vocab.json``."""
+
+    def __init__(self) -> None:
+        self._token_to_index: dict[str, dict[str, int]] = {}
+        self._index_to_token: dict[str, dict[int, str]] = {}
+
+    def add_namespace(self, namespace: str) -> None:
+        if namespace not in self._token_to_index:
+            self._token_to_index[namespace] = {NULL_LABEL: 0}
+            self._index_to_token[namespace] = {0: NULL_LABEL}
+
+    def add_token(self, token: str, namespace: str) -> int:
+        self.add_namespace(namespace)
+        t2i = self._token_to_index[namespace]
+        if token not in t2i:
+            idx = len(t2i)
+            t2i[token] = idx
+            self._index_to_token[namespace][idx] = token
+        return t2i[token]
+
+    def get_token_index(self, token: str, namespace: str) -> int:
+        return self._token_to_index[namespace][token]
+
+    def get_token_from_index(self, index: int, namespace: str) -> str:
+        return self._index_to_token[namespace][index]
+
+    def get_vocab_size(self, namespace: str) -> int:
+        return len(self._token_to_index[namespace])
+
+    def get_namespaces(self) -> list[str]:
+        return list(self._token_to_index.keys())
+
+    @classmethod
+    def load(cls, path: str | Path) -> Vocabulary:
+        """Load the exact label<->index mapping saved next to the checkpoint."""
+        vocab_path = Path(path)
+        try:
+            data = json.loads(vocab_path.read_text(encoding="utf-8"))
+        except (OSError, UnicodeError, json.JSONDecodeError) as exc:
+            raise CheckpointError(
+                f"Could not read a valid JSON vocabulary at {vocab_path}: {exc}"
+            ) from exc
+        if not isinstance(data, dict) or not data:
+            raise CheckpointError(f"vocab.json at {vocab_path} must contain a non-empty object.")
+
+        vocab = cls()
+        for namespace, t2i in data.items():
+            if not isinstance(namespace, str) or not namespace:
+                raise CheckpointError("Every vocab namespace must be a non-empty string.")
+            if not isinstance(t2i, dict) or not t2i:
+                raise CheckpointError(f"Vocab namespace {namespace!r} must be a non-empty object.")
+            if t2i.get(NULL_LABEL) != 0:
+                raise CheckpointError(
+                    f"Vocab namespace {namespace!r} must map the null label '' to index 0."
+                )
+            if not all(
+                isinstance(token, str)
+                and isinstance(index, int)
+                and not isinstance(index, bool)
+                and index >= 0
+                for token, index in t2i.items()
+            ):
+                raise CheckpointError(
+                    f"Vocab namespace {namespace!r} must map strings to non-negative integers."
+                )
+            indices = list(t2i.values())
+            if len(indices) != len(set(indices)) or sorted(indices) != list(range(len(indices))):
+                raise CheckpointError(
+                    f"Vocab namespace {namespace!r} indices must be unique and contiguous from 0."
+                )
+            vocab.add_namespace(namespace)
+            for token, _idx in sorted(t2i.items(), key=lambda kv: kv[1]):
+                if token != NULL_LABEL:
+                    vocab.add_token(token, namespace)
+        return vocab
+
+
+def derive_dataset(vocab: Vocabulary) -> str:
+    """Derive the single ``dataset`` prefix from the vocab's namespaces.
+
+    The model was trained with one ``dataset`` constant (``"radgraph-it"`` for the current
+    artifact), which prefixes both label namespaces. We recover it from the vocab rather than
+    hard-coding it, and validate that the NER and relation prefixes agree, so the head lookups
+    ``f"{dataset}__ner_labels"`` / ``f"{dataset}__relation_labels"`` are guaranteed to hit.
+    """
+    ner_prefixes = [
+        ns[: -len(NER_SUFFIX)] for ns in vocab.get_namespaces() if ns.endswith(NER_SUFFIX)
+    ]
+    rel_prefixes = [
+        ns[: -len(RELATION_SUFFIX)] for ns in vocab.get_namespaces() if ns.endswith(RELATION_SUFFIX)
+    ]
+    if len(ner_prefixes) != 1 or len(rel_prefixes) != 1:
+        raise CheckpointError(
+            "vocab.json must contain exactly one '*__ner_labels' and one '*__relation_labels' "
+            f"namespace; found NER={ner_prefixes}, relation={rel_prefixes}."
+        )
+    if ner_prefixes[0] != rel_prefixes[0]:
+        raise CheckpointError(
+            f"vocab.json NER prefix {ner_prefixes[0]!r} does not match relation prefix "
+            f"{rel_prefixes[0]!r}; the bundle is inconsistent."
+        )
+    return ner_prefixes[0]

radgraphit/inference/decoder.py

@@ -0,0 +1,53 @@
+"""Decode raw DyGIE backend output into immutable domain value objects.
+
+The backend speaks in plain tuples (the shape upstream DyGIE ``decode`` produced); this module is
+the single place that lifts those into typed :class:`ScoredEntity` / :class:`ScoredRelation`
+objects, producing one :class:`ReportPrediction` per report. Keeping this separate from both the
+torch backend and the serializer means the whole "raw predictions -> domain graph" step is pure
+and trivially testable with synthetic tuples.
+"""
+
+from __future__ import annotations
+
+from ..core.models import (
+    RawBackendOutput,
+    ReportPrediction,
+    ScoredEntity,
+    ScoredRelation,
+)
+
+
+def decode_report(
+    doc_key: str,
+    tokens: tuple[str, ...],
+    raw: RawBackendOutput,
+) -> ReportPrediction:
+    """Lift one report's raw backend tuples into a typed :class:`ReportPrediction`."""
+    entities = tuple(
+        ScoredEntity(
+            start_ix=int(start),
+            end_ix=int(end),
+            label=str(label),
+            raw_score=float(raw_score),
+            softmax_score=float(soft_score),
+        )
+        for (start, end, label, raw_score, soft_score) in raw.ner
+    )
+    relations = tuple(
+        ScoredRelation(
+            label=str(label),
+            source_start=int(s1),
+            source_end=int(e1),
+            target_start=int(s2),
+            target_end=int(e2),
+            raw_score=float(raw_score),
+            softmax_score=float(soft_score),
+        )
+        for (s1, e1, s2, e2, label, raw_score, soft_score) in raw.relations
+    )
+    return ReportPrediction(
+        doc_key=doc_key,
+        tokens=tokens,
+        entities=entities,
+        relations=relations,
+    )

radgraphit/inference/service.py

@@ -0,0 +1,80 @@
+"""Inference orchestration, with no global I/O.
+
+Ties the pure pieces together — validate input, tokenize, run the (injected) backend, decode to
+domain objects, serialize to RadGraph-XL — around any :class:`InferenceBackend`. The backend is a
+constructor argument, which is exactly where a fake backend is injected in tests; the service
+itself never imports torch.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+
+from ..core.errors import InvalidReportError
+from ..core.models import InferenceBackend, ReportPrediction
+from ..serialization.radgraph_xl import (
+    SerializationDiagnostics,
+    to_radgraph_xl,
+    to_radgraph_xl_with_diagnostics,
+)
+from .decoder import decode_report
+from .tokenizer import tokenize_report
+
+
+def normalize_reports(reports: str | Sequence[str]) -> list[str]:
+    """Validate and normalize the public input into a non-empty list of non-empty strings.
+
+    Accepts a single ``str`` or a ``Sequence[str]``. Rejects — deterministically, with
+    :class:`InvalidReportError` — anything else: the wrong type, an empty container, a non-string
+    element, or a string that is empty or whitespace-only. Order is preserved.
+
+    Unlike the upstream package, an empty report is never silently coerced to the string
+    ``"None"``; it is an error the caller must see.
+    """
+    if isinstance(reports, str):
+        items = [reports]
+    elif isinstance(reports, Sequence) and not isinstance(reports, bytes | bytearray):
+        items = list(reports)
+    else:
+        raise InvalidReportError(
+            f"reports must be a str or a sequence of str, got {type(reports).__name__}."
+        )
+
+    if not items:
+        raise InvalidReportError("reports is empty: provide at least one report.")
+
+    for index, item in enumerate(items):
+        if not isinstance(item, str):
+            raise InvalidReportError(
+                f"report at index {index} must be a str, got {type(item).__name__}."
+            )
+        if not item.strip():
+            raise InvalidReportError(f"report at index {index} is empty or whitespace-only.")
+    return items
+
+
+class InferenceService:
+    """Stateless orchestration around a single backend. One report is processed at a time."""
+
+    def __init__(self, backend: InferenceBackend):
+        self._backend = backend
+
+    def predict_graphs(self, reports: str | Sequence[str]) -> list[ReportPrediction]:
+        """Validate, tokenize, infer, and decode — returning the typed domain graphs (in order)."""
+        items = normalize_reports(reports)
+        predictions: list[ReportPrediction] = []
+        for index, text in enumerate(items):
+            tokenized = tokenize_report(text)
+            raw = self._backend.infer(tokenized.tokens)
+            predictions.append(decode_report(str(index), tokenized.tokens, raw))
+        return predictions
+
+    def predict(self, reports: str | Sequence[str]) -> dict[str, object]:
+        """Return the canonical RadGraph-XL dict (the public contract)."""
+        return to_radgraph_xl(self.predict_graphs(reports))
+
+    def predict_with_diagnostics(
+        self, reports: str | Sequence[str]
+    ) -> tuple[dict[str, object], SerializationDiagnostics]:
+        """Return the RadGraph-XL dict plus internal serialization diagnostics."""
+        return to_radgraph_xl_with_diagnostics(self.predict_graphs(reports))