Add transformers-compatible wrapper (AutoModel/AutoTokenizer via trust_remote_code)

README.md

@@ -0,0 +1,65 @@
+---
+license: apache-2.0
+language:
+- it
+tags:
+- radiology
+- information-extraction
+- named-entity-recognition
+- relation-extraction
+- medical
+- radgraph
+---
+
+# Radgraph-IT-v1
+
+A DyGIE++-style joint model for named entity recognition and relation extraction over Italian
+radiology reports, following the RadGraph schema. Encoder: [`IVN-RIN/medBIT-r3-plus`](https://huggingface.co/IVN-RIN/medBIT-r3-plus).
+
+Entity labels: `Anatomy` / `Observation`, each with `definitely present` / `definitely absent` /
+`uncertain`. Relation labels: `modify`, `located_at`, `suggestive_of`.
+
+## Usage
+
+```python
+from transformers import AutoModel
+
+model = AutoModel.from_pretrained("radgraphIT/Radgraph-IT-v1", trust_remote_code=True)
+model.eval()
+
+out = model.predict("Non versamento pleurico. Addensamento parenchimale al lobo inferiore del polmone destro.")
+print(out["words"])
+print(out["entities"])    # [start, end, label, raw_score, softmax_score] (word-index spans, inclusive)
+print(out["relations"])   # [start1, end1, start2, end2, label, raw_score, softmax_score]
+```
+
+`model.predict(text)` applies the same word-level preprocessing (regex cleanup +
+`nltk.wordpunct_tokenize`) used at training time. If you already have a word-tokenized list
+matching that exact scheme, use `model.predict_words(words)` to skip re-tokenizing.
+
+The bundled tokenizer (`AutoTokenizer.from_pretrained("radgraphIT/Radgraph-IT-v1")`) is the
+underlying subword tokenizer for the transformer encoder; it is not the word-level splitter
+`predict()` uses internally.
+
+Requires `transformers`, `torch`, `numpy`, and `nltk` (for `predict()`).
+
+## Architecture
+
+Word-level embeddings come from mean-pooling `medBIT-r3-plus` subwords back to one vector per
+word (documents longer than 512 tokens are split into word-aligned chunks). Span representations
+concatenate start/end word embeddings with a span-width embedding. NER scores every candidate
+span (width ≤ 12); the relation head prunes to the top spans-per-word and scores every pair.
+
+Trained on the full (non-cross-validation) split -- see repo commit history for provenance
+(`fold0/` → `model/`, port of a private AllenNLP-based training pipeline to a from-scratch,
+framework-free PyTorch + 🤗 `transformers` implementation).
+
+## Metrics (validation, best epoch 8)
+
+| | NER | Relation |
+|---|---|---|
+| Precision | 0.848 | 0.700 |
+| Recall | 0.868 | 0.680 |
+| F1 | 0.858 | 0.690 |
+
+Full per-class breakdown in `model/metrics.json`.

config.json

@@ -0,0 +1,49 @@
+{
+  "model_type": "radgraph_it",
+  "architectures": [
+    "RadgraphModel"
+  ],
+  "auto_map": {
+    "AutoConfig": "configuration_radgraph.RadgraphConfig",
+    "AutoModel": "modeling_radgraph.RadgraphModel"
+  },
+  "encoder_name": "IVN-RIN/medBIT-r3-plus",
+  "max_length": 512,
+  "max_span_width": 12,
+  "feature_size": 20,
+  "feedforward_params": {
+    "hidden_dims": [
+      150,
+      150
+    ],
+    "dropout": 0.4
+  },
+  "loss_weights": {
+    "ner": 0.2,
+    "relation": 1.0
+  },
+  "relation_spans_per_word": 0.5,
+  "train_encoder": true,
+  "span_pooling": false,
+  "transformer_params": null,
+  "relation_context": false,
+  "relation_feedforward_params": null,
+  "dataset": "radgraph-it",
+  "ner_labels": {
+    "": 0,
+    "Anatomy::definitely present": 1,
+    "Observation::definitely present": 2,
+    "Observation::definitely absent": 3,
+    "Observation::uncertain": 4,
+    "Anatomy::definitely absent": 5,
+    "Anatomy::uncertain": 6
+  },
+  "relation_labels": {
+    "": 0,
+    "modify": 1,
+    "located_at": 2,
+    "suggestive_of": 3
+  },
+  "torch_dtype": "float32",
+  "transformers_version": "5.14.1"
+}

configuration_radgraph.py

@@ -0,0 +1,65 @@
+"""transformers-compatible config for the RadGraph-IT DyGIE++ joint NER + relation model.
+
+Mirrors the hyperparameters of `training_v2/src/dygie/model.py`'s `DyGIEModel` plus the
+label vocabulary (`training_v2/src/dygie/vocab.py`'s `Vocabulary`), so a `RadgraphModel` can
+be reconstructed from `config.json` alone, matching the checkpoint trained by
+`run_medbit_full_cv.sh` (single train run on the full split, not actual cross-validation --
+see that script's own header comment).
+"""
+from transformers import PretrainedConfig
+
+
+class RadgraphConfig(PretrainedConfig):
+    model_type = "radgraph_it"
+
+    def __init__(
+        self,
+        encoder_name: str = "IVN-RIN/medBIT-r3-plus",
+        max_length: int = 512,
+        max_span_width: int = 12,
+        feature_size: int = 20,
+        feedforward_params: dict = None,
+        loss_weights: dict = None,
+        relation_spans_per_word: float = 0.5,
+        train_encoder: bool = True,
+        span_pooling: bool = False,
+        transformer_params: dict = None,
+        relation_context: bool = False,
+        relation_feedforward_params: dict = None,
+        dataset: str = "radgraph-it",
+        ner_labels: dict = None,
+        relation_labels: dict = None,
+        **kwargs,
+    ):
+        self.encoder_name = encoder_name
+        self.max_length = max_length
+        self.max_span_width = max_span_width
+        self.feature_size = feature_size
+        self.feedforward_params = feedforward_params or {"hidden_dims": [150, 150], "dropout": 0.4}
+        self.loss_weights = loss_weights or {"ner": 0.2, "relation": 1.0}
+        self.relation_spans_per_word = relation_spans_per_word
+        self.train_encoder = train_encoder
+        self.span_pooling = span_pooling
+        self.transformer_params = transformer_params
+        self.relation_context = relation_context
+        self.relation_feedforward_params = relation_feedforward_params
+        self.dataset = dataset
+        # Namespace-unqualified label -> index maps, null label "" pinned to 0. Namespaced as
+        # f"{dataset}__ner_labels" / f"{dataset}__relation_labels" when rebuilt into a Vocabulary
+        # (see modeling_radgraph.py), matching training_v2/src/dygie/vocab.py exactly.
+        self.ner_labels = ner_labels or {
+            "": 0,
+            "Anatomy::definitely present": 1,
+            "Observation::definitely present": 2,
+            "Observation::definitely absent": 3,
+            "Observation::uncertain": 4,
+            "Anatomy::definitely absent": 5,
+            "Anatomy::uncertain": 6,
+        }
+        self.relation_labels = relation_labels or {
+            "": 0,
+            "modify": 1,
+            "located_at": 2,
+            "suggestive_of": 3,
+        }
+        super().__init__(**kwargs)

dataset.py

@@ -0,0 +1,123 @@
+"""JSONL -> training examples.
+
+Port of the relevant half of `radgraph.dygie.data.dataset_readers.dygie.DyGIEReader`
+(enumerate candidate spans, attach NER/relation gold labels, drop gold spans wider than
+`max_span_width`) plus the coref/event fields it also reads but which this project never
+weights (see training/README.md: loss_weights.coref = loss_weights.events = 0 in every config).
+
+Simplification vs. v1: every document in this dataset is a single "sentence" spanning the
+whole report (see training/README.md), so there is no batching-over-sentences dimension here
+-- one `Example` = one document. `build_example` asserts this so a future multi-sentence
+document fails loudly instead of silently mis-training.
+
+Simplification vs. v1 (memory): AllenNLP's `AdjacencyField` materializes a dense
+(num_spans, num_spans) gold-relation matrix per document at collate time -- for a long report
+with ~2000 words and max_span_width=12 that's ~24k spans, i.e. a ~575M-cell tensor. We instead
+keep gold relations as a sparse {(span_i, span_j): label_id} dict and only ever materialize the
+small (num_kept, num_kept) submatrix after pruning (see model.py), which is mathematically
+identical (same loss, same gradients) at a fraction of the memory.
+"""
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Dict, List, Tuple
+
+import torch
+from torch.utils.data import Dataset as TorchDataset
+
+from .document import Document, Sentence
+from .vocab import Vocabulary
+
+
+def enumerate_spans(words: 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 = []
+    for start in range(n):
+        for end in range(start, min(start + max_span_width, n)):
+            spans.append((start, end))
+    return spans
+
+
+def _too_long(span: Tuple[int, int], max_span_width: int) -> bool:
+    return span[1] - span[0] + 1 > max_span_width
+
+
+@dataclass
+class Example:
+    doc_key: object
+    dataset: str
+    words: List[str]
+    spans: List[Tuple[int, int]]
+    span_index: Dict[Tuple[int, int], int]
+    ner_label_ids: torch.Tensor                       # (num_spans,) long, 0 = null
+    relation_gold: Dict[Tuple[int, int], int]          # (span_i, span_j) -> label id (>=1), sparse
+    weight: float
+    sentence: Sentence                                  # unfiltered gold, for relation eval (see module docstring)
+
+
+def build_example(doc: Document, vocab: Vocabulary, max_span_width: int) -> Example:
+    if len(doc.sentences) != 1:
+        raise NotImplementedError(
+            f"training_v2 assumes one sentence per document; {doc.doc_key} has {len(doc.sentences)}. "
+            "See training/README.md ('the whole report is a single sentence')."
+        )
+    sent = doc.sentences[0]
+    ner_ns = f"{doc.dataset}__ner_labels"
+    rel_ns = f"{doc.dataset}__relation_labels"
+
+    spans = enumerate_spans(sent.text, max_span_width)
+    span_index = {s: i for i, s in enumerate(spans)}
+
+    ner_label_ids = torch.zeros(len(spans), dtype=torch.long)
+    for span, label in (sent.ner_dict or {}).items():
+        if _too_long(span, max_span_width):
+            continue
+        ner_label_ids[span_index[span]] = vocab.get_token_index(label, ner_ns)
+
+    relation_gold: Dict[Tuple[int, int], int] = {}
+    for (span1, span2), label in (sent.relation_dict or {}).items():
+        if _too_long(span1, max_span_width) or _too_long(span2, max_span_width):
+            continue
+        relation_gold[(span_index[span1], span_index[span2])] = vocab.get_token_index(label, rel_ns)
+
+    return Example(doc_key=doc.doc_key, dataset=doc.dataset, words=sent.text, spans=spans,
+                    span_index=span_index, ner_label_ids=ner_label_ids, relation_gold=relation_gold,
+                    weight=doc.weight if doc.weight is not None else 1.0, sentence=sent)
+
+
+def read_documents(path: str) -> List[Document]:
+    import json
+    documents = []
+    with open(path) as f:
+        for line in f:
+            line = line.strip()
+            if line:
+                documents.append(Document.from_json(json.loads(line)))
+    return documents
+
+
+class ExampleDataset(TorchDataset):
+    """One dataset split (train/dev/test), pre-built into `Example`s.
+
+    `collate_fn` is the identity on the single element: every config in this project uses
+    `data_loader.batch_size = 1` (the joint model does not support batching multiple documents,
+    see model.py), so there is no padding/collation logic to write.
+    """
+
+    def __init__(self, documents: List[Document], vocab: Vocabulary, max_span_width: int):
+        self.examples = [build_example(doc, vocab, max_span_width) for doc in documents]
+
+    def __len__(self):
+        return len(self.examples)
+
+    def __getitem__(self, ix) -> Example:
+        return self.examples[ix]
+
+    @classmethod
+    def from_jsonl(cls, path: str, vocab: Vocabulary, max_span_width: int) -> "ExampleDataset":
+        return cls(read_documents(path), vocab, max_span_width)
+
+
+def collate_fn(batch: List[Example]) -> Example:
+    assert len(batch) == 1, "training_v2 only supports batch_size=1 (see module docstring)."
+    return batch[0]

document.py

@@ -0,0 +1,254 @@
+"""In-memory representation of a DyGIE-format document (JSONL).
+
+Trimmed port of the vendored radgraph.dygie.data.dataset_readers.document module: only the
+NER + relation pieces are kept (no coreference clusters, no events -- this project's configs
+always set loss_weights.coref = loss_weights.events = 0, so those heads never ran in v1
+either). See training/README.md for the schema.
+"""
+import re
+import json
+from typing import Any, Dict, List, Optional
+
+import numpy as np
+
+
+def _fields_to_batches(d: dict, keys_to_ignore=()):
+    """Inverse of `_batches_to_fields`: {"a": [1, 2], "b": [3, 4]} -> [{"a": 1, "b": 3}, {"a": 2, "b": 4}]."""
+    keys = [k for k in d.keys() if k not in keys_to_ignore]
+    lengths = {k: len(d[k]) for k in keys}
+    if len(set(lengths.values())) != 1:
+        raise ValueError(f"For document {d.get('doc_key')}, fields have different lengths: {lengths}.")
+    length = next(iter(lengths.values()))
+    return [{k: d[k][i] for k in keys} for i in range(length)]
+
+
+def _batches_to_fields(batches: List[dict]):
+    first_keys = batches[0].keys()
+    for entry in batches[1:]:
+        if set(entry.keys()) != set(first_keys):
+            raise ValueError("Keys do not match on all entries.")
+    res = {k: [] for k in first_keys}
+    for batch in batches:
+        for k, v in batch.items():
+            res[k].append(v)
+    return res
+
+
+class Span:
+    """A span, tracked both sentence-relative and document-relative."""
+
+    def __init__(self, start: int, end: int, sentence: "Sentence", sentence_offsets: bool = False):
+        self.sentence = sentence
+        # `sentence.text_joined` is memoized on Sentence (computed once) rather than rejoined
+        # here per span: the relation head can construct O(K^2) PredictedRelation/Span objects
+        # per document during decode (K = pruned span count), so re-joining a ~1000-word
+        # sentence per span turns into the dominant runtime cost otherwise.
+        self.sentence_text = sentence.text_joined
+        self.start_sent = start if sentence_offsets else start - sentence.sentence_start
+        self.end_sent = end if sentence_offsets else end - sentence.sentence_start
+
+    @property
+    def start_doc(self):
+        return self.start_sent + self.sentence.sentence_start
+
+    @property
+    def end_doc(self):
+        return self.end_sent + self.sentence.sentence_start
+
+    @property
+    def span_doc(self):
+        return (self.start_doc, self.end_doc)
+
+    @property
+    def span_sent(self):
+        return (self.start_sent, self.end_sent)
+
+    def __repr__(self):
+        return str(self.span_sent)
+
+    def __eq__(self, other):
+        return (self.span_doc == other.span_doc and self.span_sent == other.span_sent
+                and self.sentence == other.sentence)
+
+    def __hash__(self):
+        return hash(self.span_sent + (self.sentence_text,))
+
+
+class NER:
+    def __init__(self, ner, sentence: "Sentence", sentence_offsets: bool = False):
+        self.span = Span(ner[0], ner[1], sentence, sentence_offsets)
+        self.label = ner[2]
+
+    def __repr__(self):
+        return f"{self.span!r}: {self.label}"
+
+    def __eq__(self, other):
+        return self.span == other.span and self.label == other.label
+
+    def to_json(self):
+        return list(self.span.span_doc) + [self.label]
+
+
+def _format_float(x):
+    return round(x, 4)
+
+
+class PredictedNER(NER):
+    def __init__(self, ner, sentence, sentence_offsets: bool = False):
+        """`ner` = [span_start, span_end, label, raw_score, softmax_score]."""
+        super().__init__(ner, sentence, sentence_offsets)
+        self.raw_score = ner[3]
+        self.softmax_score = ner[4]
+
+    def to_json(self):
+        return super().to_json() + [_format_float(self.raw_score), _format_float(self.softmax_score)]
+
+
+class Relation:
+    def __init__(self, relation, sentence: "Sentence", sentence_offsets: bool = False):
+        start1, end1, start2, end2, label = relation
+        self.pair = (Span(start1, end1, sentence, sentence_offsets),
+                     Span(start2, end2, sentence, sentence_offsets))
+        self.label = label
+
+    def __repr__(self):
+        return f"{self.pair[0]!r}, {self.pair[1]!r}: {self.label}"
+
+    def __eq__(self, other):
+        return self.pair == other.pair and self.label == other.label
+
+    def to_json(self):
+        return list(self.pair[0].span_doc) + list(self.pair[1].span_doc) + [self.label]
+
+
+class PredictedRelation(Relation):
+    def __init__(self, relation, sentence, sentence_offsets: bool = False):
+        """`relation` = [start1, end1, start2, end2, label, raw_score, softmax_score]."""
+        super().__init__(relation[:5], sentence, sentence_offsets)
+        self.raw_score = relation[5]
+        self.softmax_score = relation[6]
+
+    def to_json(self):
+        return super().to_json() + [_format_float(self.raw_score), _format_float(self.softmax_score)]
+
+
+class Sentence:
+    """Despite the name, this project's documents always have exactly one "sentence" spanning
+    the whole report (see training/README.md); the multi-sentence machinery is kept because
+    the JSONL format is naturally list-of-sentences and nothing is gained by special-casing it.
+    """
+
+    def __init__(self, entry: dict, sentence_start: int, sentence_ix: int):
+        self.sentence_start = sentence_start
+        self.sentence_ix = sentence_ix
+        self.text = entry["sentences"]
+        self.text_joined = " ".join(self.text)  # memoized once; see Span.__init__
+        self.metadata = {k: v for k, v in entry.items() if k.startswith("_")}
+
+        if "ner" in entry:
+            self.ner = [NER(x, self) for x in entry["ner"]]
+            self.ner_dict = {e.span.span_sent: e.label for e in self.ner}
+        else:
+            self.ner, self.ner_dict = None, None
+
+        self.predicted_ner = ([PredictedNER(x, self) for x in entry["predicted_ner"]]
+                               if "predicted_ner" in entry else None)
+
+        if "relations" in entry:
+            self.relations = [Relation(x, self) for x in entry["relations"]]
+            self.relation_dict = {(r.pair[0].span_sent, r.pair[1].span_sent): r.label
+                                   for r in self.relations}
+        else:
+            self.relations, self.relation_dict = None, None
+
+        self.predicted_relations = ([PredictedRelation(x, self) for x in entry["predicted_relations"]]
+                                     if "predicted_relations" in entry else None)
+
+    def to_json(self):
+        res = {"sentences": self.text}
+        if self.ner is not None:
+            res["ner"] = [e.to_json() for e in self.ner]
+        if self.predicted_ner is not None:
+            res["predicted_ner"] = [e.to_json() for e in self.predicted_ner]
+        if self.relations is not None:
+            res["relations"] = [r.to_json() for r in self.relations]
+        if self.predicted_relations is not None:
+            res["predicted_relations"] = [r.to_json() for r in self.predicted_relations]
+        res.update(self.metadata)
+        return res
+
+    def __len__(self):
+        return len(self.text)
+
+    def __repr__(self):
+        return " ".join(self.text)
+
+
+class Document:
+    _ALLOWED_FIELDS = re.compile(r"doc_key|dataset|sentences|weight|.*ner$|.*relations$|^_.*")
+
+    def __init__(self, doc_key, dataset, sentences: List[Sentence], weight: Optional[float] = None):
+        self.doc_key = doc_key
+        self.dataset = dataset
+        self.sentences = sentences
+        self.weight = weight
+
+    @classmethod
+    def from_json(cls, js: Dict[str, Any]) -> "Document":
+        unexpected = [f for f in js if not cls._ALLOWED_FIELDS.match(f)]
+        if unexpected:
+            raise ValueError(f"Unexpected fields (prefix with '_' if intentional): {unexpected}")
+
+        doc_key = js["doc_key"]
+        dataset = js.get("dataset")
+        entries = _fields_to_batches(js, ("doc_key", "dataset", "weight"))
+        sentence_lengths = [len(e["sentences"]) for e in entries]
+        sentence_starts = np.roll(np.cumsum(sentence_lengths), 1)
+        sentence_starts[0] = 0
+        sentences = [Sentence(entry, int(start), ix)
+                     for ix, (entry, start) in enumerate(zip(entries, sentence_starts.tolist()))]
+        return cls(doc_key, dataset, sentences, js.get("weight"))
+
+    def to_json(self):
+        res = {"doc_key": self.doc_key, "dataset": self.dataset}
+        res.update(_batches_to_fields([s.to_json() for s in self.sentences]))
+        if self.weight is not None:
+            res["weight"] = self.weight
+        return res
+
+    @property
+    def n_tokens(self):
+        return sum(len(s) for s in self.sentences)
+
+    def __getitem__(self, ix):
+        return self.sentences[ix]
+
+    def __len__(self):
+        return len(self.sentences)
+
+    def __repr__(self):
+        return "\n".join(f"{i}: {' '.join(s.text)}" for i, s in enumerate(self.sentences))
+
+
+class Dataset:
+    def __init__(self, documents: List[Document]):
+        self.documents = documents
+
+    def __getitem__(self, i):
+        return self.documents[i]
+
+    def __len__(self):
+        return len(self.documents)
+
+    @classmethod
+    def from_jsonl(cls, fname):
+        documents = []
+        with open(fname) as f:
+            for line in f:
+                documents.append(Document.from_json(json.loads(line)))
+        return cls(documents)
+
+    def to_jsonl(self, fname):
+        with open(fname, "w") as f:
+            for doc in self.documents:
+                print(json.dumps(doc.to_json()), file=f)

feedforward.py

@@ -0,0 +1,23 @@
+"""Small MLP used by every scorer (NER, mention pruner, relation): Linear -> ReLU -> Dropout,
+one block per entry in `hidden_dims`, each block's width taken from that entry (see
+configs/medbit_span_pooling/fold0.json: feedforward_params = {hidden_dims: [300, 150],
+dropout: 0.4}).
+"""
+from typing import List
+
+from torch import nn
+
+
+class FeedForward(nn.Module):
+    def __init__(self, input_dim: int, hidden_dims: List[int], dropout: float):
+        super().__init__()
+        layers = []
+        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):
+        return self.net(x)

metrics.py

@@ -0,0 +1,91 @@
+"""Micro + per-class precision/recall/F1 accumulators. Port of ner_metrics.py / relation_metrics.py / f1.py.
+
+Per-class tp/fp/fn are accumulated alongside the micro totals so macro-F1 and the rare-class
+breakdown can be logged live during training (not only post-hoc in eval.py) -- same information
+the v1 exp2 class-weighting runs surfaced to W&B.
+"""
+from collections import defaultdict
+from typing import Dict, Tuple
+
+import torch
+
+
+def compute_f1(predicted: float, gold: float, matched: float) -> Tuple[float, float, float]:
+    precision = matched / predicted if predicted > 0 else 0.0
+    recall = matched / gold if gold > 0 else 0.0
+    f1 = 2 * precision * recall / (precision + recall) if (precision + recall) > 0 else 0.0
+    return precision, recall, f1
+
+
+class NERMetrics:
+    """Exact-match span+label F1, from tensors (no mask: batch_size=1 means no padding)."""
+
+    def __init__(self, num_classes: int, null_label: int = 0):
+        self.num_classes = num_classes
+        self.null_label = null_label
+        self.reset()
+
+    def reset(self):
+        self.tp = self.fp = self.fn = 0
+        self.per_class = {c: [0, 0, 0] for c in range(self.num_classes) if c != self.null_label}
+
+    def __call__(self, predictions: torch.Tensor, gold_labels: torch.Tensor):
+        for c in range(self.num_classes):
+            if c == self.null_label:
+                continue
+            tp = int(((predictions == c) & (gold_labels == c)).sum())
+            fp = int(((predictions == c) & (gold_labels != c)).sum())
+            fn = int(((predictions != c) & (gold_labels == c)).sum())
+            self.tp += tp
+            self.fp += fp
+            self.fn += fn
+            pc = self.per_class[c]
+            pc[0] += tp
+            pc[1] += fp
+            pc[2] += fn
+
+    def get_metric(self, reset: bool = False) -> Tuple[float, float, float]:
+        result = compute_f1(self.tp + self.fp, self.tp + self.fn, self.tp)
+        if reset:
+            self.reset()
+        return result
+
+    def get_per_class(self) -> Dict[int, float]:
+        """{class index -> F1}. Call BEFORE get_metric(reset=True), which clears the counts."""
+        return {c: compute_f1(v[0] + v[1], v[0] + v[2], v[0])[2] for c, v in self.per_class.items()}
+
+
+class RelationMetrics:
+    """Exact-match (span1, span2, label) F1, decoded predictions vs. the full (unfiltered) gold
+    relation dict -- see dataset.py's module docstring for why gold relations whose spans exceed
+    max_span_width still count as false negatives here."""
+
+    def __init__(self):
+        self.reset()
+
+    def reset(self):
+        self.total_gold = self.total_predicted = self.total_matched = 0
+        self.per_class = defaultdict(lambda: [0, 0, 0])  # label -> [tp, fp, fn]
+
+    def __call__(self, predicted: Dict[Tuple, str], gold: Dict[Tuple, str]):
+        self.total_gold += len(gold)
+        self.total_predicted += len(predicted)
+        for key, label in predicted.items():
+            if gold.get(key) == label:
+                self.total_matched += 1
+                self.per_class[label][0] += 1  # tp
+            else:
+                self.per_class[label][1] += 1  # fp: label wrong, or key absent from gold
+        for key, label in gold.items():
+            if predicted.get(key) != label:
+                self.per_class[label][2] += 1  # fn
+
+    def get_metric(self, reset: bool = False) -> Tuple[float, float, float]:
+        result = compute_f1(self.total_predicted, self.total_gold, self.total_matched)
+        if reset:
+            self.reset()
+        return result
+
+    def get_per_class(self) -> Dict[str, float]:
+        """{relation label -> F1}. Call BEFORE get_metric(reset=True), which clears the counts."""
+        return {lab: compute_f1(v[0] + v[1], v[0] + v[2], v[0])[2] for lab, v in self.per_class.items()}

model.py

@@ -0,0 +1,82 @@
+"""The joint model: encoder -> span representations -> NER head + relation head.
+
+Port of `radgraph.dygie.models.dygie.DyGIE`, coref/events removed (see package docstring in
+ner_head.py / relation_head.py for why: every config in this project sets their loss weight to
+0, so they never trained or predicted anything in v1 either).
+"""
+from typing import Dict
+
+import torch
+from torch import nn
+
+from .dataset import Example
+from .ner_head import NERHead
+from .relation_head import RelationHead
+from .span_extractor import EndpointSpanExtractor
+from .tokenizer_embedder import MismatchedEmbedder
+from .vocab import Vocabulary
+
+
+def _xavier_init_weights(module: nn.Module) -> None:
+    """xavier_normal_ on every 2-D '*.weight' / '*.weight_matrix' param -- the
+    `module_initializer` regexes every config in configs/medbit/ applies to the NER and
+    relation submodules (never to the pretrained encoder, which keeps its pretrained init)."""
+    for name, param in module.named_parameters():
+        if param.dim() >= 2 and (name.endswith("weight") or name.endswith("weight_matrix")):
+            nn.init.xavier_normal_(param)
+
+
+class DyGIEModel(nn.Module):
+    def __init__(self, vocab: Vocabulary, encoder_name: str, max_length: int, max_span_width: int,
+                 feature_size: int, feedforward_params: dict, loss_weights: Dict[str, float],
+                 relation_spans_per_word: float, train_encoder: bool = True,
+                 span_pooling: bool = False, transformer_params: dict = None,
+                 relation_context: bool = False, relation_feedforward_params: dict = None):
+        super().__init__()
+        self.vocab = vocab
+        self.loss_weights = loss_weights
+
+        self.embedder = MismatchedEmbedder(encoder_name, max_length, train_encoder)
+        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)
+
+        _xavier_init_weights(self.ner)
+        _xavier_init_weights(self.relation)
+        nn.init.xavier_normal_(self.span_extractor.width_embedding.weight)
+
+    def forward(self, example: Example) -> dict:
+        device = next(self.embedder.parameters()).device
+        word_embeddings = self.embedder(example.words)
+        spans_tensor = torch.tensor(example.spans, dtype=torch.long, device=device)
+        span_embeddings = self.span_extractor(word_embeddings, spans_tensor)
+
+        zero = word_embeddings.new_zeros(())
+        output_ner, output_relation = {"loss": zero}, {"loss": zero}
+
+        if self.loss_weights["ner"] > 0:
+            output_ner = self.ner(example.dataset, example.spans, span_embeddings, example.sentence,
+                                   example.ner_label_ids.to(device))
+        if self.loss_weights["relation"] > 0:
+            output_relation = self.relation(example.dataset, example.spans, span_embeddings,
+                                             len(example.words), word_embeddings,
+                                             example.relation_gold, example.sentence)
+
+        loss = (self.loss_weights["ner"] * output_ner.get("loss", zero) +
+                self.loss_weights["relation"] * output_relation.get("loss", zero))
+        loss = loss * example.weight
+
+        return {"loss": loss, "ner": output_ner, "relation": output_relation}
+
+    def get_metrics(self, reset: bool = False) -> Dict[str, float]:
+        res = {}
+        res.update(self.ner.get_metrics(reset))
+        res.update(self.relation.get_metrics(reset))
+        return res

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:31eb1746b8c7642d6421213911a05a522fffb026cd1818d49643f4fbe23dc4b2
+size 444688808

modeling_radgraph.py

@@ -0,0 +1,107 @@
+"""transformers-compatible wrapper around `DyGIEModel` (training_v2/src/dygie/model.py).
+
+`RadgraphModel.model` is a plain `DyGIEModel`, so its state_dict keys are `model.<...>`,
+matching the `model.` prefix applied when this repo's checkpoint was converted from the
+original `best.pt` (`{"epoch": ..., "model_state_dict": ...}`, see observability.py's
+`HFCheckpointUploader` / trainer.py in the training repo) to `model.safetensors`.
+
+Two ways to run inference:
+  - `model.predict(text)` -- raw report text in, decoded entities/relations out. Applies the
+    exact word-level preprocessing training used (regex cleanup + wordpunct tokenization,
+    ported from radgraph/radgraph/utils.py's `radgraph_xl_preprocess_report`, verified against
+    the gold tokenization by check_tokenization_raw_it.py in the training repo).
+  - `model.predict_words(words)` -- skip preprocessing if you already have the exact
+    word-tokenized list (e.g. re-running on the original training data).
+"""
+import re
+from typing import List
+
+import torch
+from transformers import PreTrainedModel
+from nltk.tokenize import wordpunct_tokenize
+
+from .configuration_radgraph import RadgraphConfig
+from .document import Document
+from .model import DyGIEModel
+from .predictor import predict_document
+from .vocab import Vocabulary
+
+
+def _preprocess_report(text: str) -> str:
+    """Exact port of radgraph/radgraph/utils.py:34-48 (inference-time preprocessing)."""
+    text = text.replace("\\n", "  ")
+    text = text.replace("\\f", "  ")
+    text = text.replace("\\u2122", "      ")
+    text = text.replace("\n", " ")
+    text = text.replace("\\\"", "``")
+    text_sub = re.sub(r"\s+", " ", text)
+    t = " ".join(wordpunct_tokenize(text_sub))
+    t = t.replace(").", ") .")
+    t = t.replace("%.", "% .")
+    t = t.replace(".'", ". '")
+    t = t.replace("%,", "% ,")
+    t = t.replace("%)", "% )")
+    return t
+
+
+def _build_vocab(config: RadgraphConfig) -> Vocabulary:
+    vocab = Vocabulary()
+    ner_ns = f"{config.dataset}__ner_labels"
+    rel_ns = f"{config.dataset}__relation_labels"
+    vocab.add_namespace(ner_ns)
+    vocab.add_namespace(rel_ns)
+    for token, _ in sorted(config.ner_labels.items(), key=lambda kv: kv[1]):
+        if token != "":
+            vocab.add_token(token, ner_ns)
+    for token, _ in sorted(config.relation_labels.items(), key=lambda kv: kv[1]):
+        if token != "":
+            vocab.add_token(token, rel_ns)
+    return vocab
+
+
+class RadgraphModel(PreTrainedModel):
+    config_class = RadgraphConfig
+
+    def __init__(self, config: RadgraphConfig):
+        super().__init__(config)
+        self.vocab = _build_vocab(config)
+        self.model = DyGIEModel(
+            vocab=self.vocab,
+            encoder_name=config.encoder_name,
+            max_length=config.max_length,
+            max_span_width=config.max_span_width,
+            feature_size=config.feature_size,
+            feedforward_params=config.feedforward_params,
+            loss_weights=config.loss_weights,
+            relation_spans_per_word=config.relation_spans_per_word,
+            train_encoder=config.train_encoder,
+            span_pooling=config.span_pooling,
+            transformer_params=config.transformer_params,
+            relation_context=config.relation_context,
+            relation_feedforward_params=config.relation_feedforward_params,
+        )
+        # Required by transformers' from_pretrained machinery (sets up all_tied_weights_keys
+        # etc.) -- harmless here since DyGIEModel has no tied weights and the checkpoint's
+        # state dict gets loaded over whatever this touches anyway.
+        self.post_init()
+
+    def forward(self, words: List[str], doc_key: str = "doc"):
+        """Low-level forward pass: pre-tokenized words in, raw model output dict out
+        (see DyGIEModel.forward / predict_document)."""
+        doc = Document.from_json({"doc_key": doc_key, "dataset": self.config.dataset, "sentences": [words]})
+        return predict_document(self.model, self.vocab, doc, self.config.max_span_width)
+
+    @torch.no_grad()
+    def predict_words(self, words: List[str], doc_key: str = "doc") -> dict:
+        result = self.forward(words, doc_key=doc_key)
+        return {
+            "doc_key": result["doc_key"],
+            "words": words,
+            "entities": [e.to_json() for e in result["predicted_ner"]],
+            "relations": [r.to_json() for r in result["predicted_relations"]],
+        }
+
+    @torch.no_grad()
+    def predict(self, text: str, doc_key: str = "doc") -> dict:
+        words = _preprocess_report(text).split()
+        return self.predict_words(words, doc_key=doc_key)

ner_head.py

@@ -0,0 +1,99 @@
+"""NER head: score every candidate span, cross-entropy against the gold label (or null).
+
+Port of `radgraph.dygie.models.ner.NERTagger`. Simplified for batch_size=1 (dataset.py): no
+span mask, no per-sentence looping -- one document's worth of spans at a time.
+
+`transformer_params` (medbit_transformer arm, see notes/v2_plan.md item 1) swaps the plain
+FeedForward scorer for a TransformerBlock: every candidate span attends to every other
+candidate span in the document before classification, instead of each span being scored in
+isolation. NER scores every candidate span unpruned, so this is full O(n^2) self-attention
+over the whole doc's span pool (accepted cost, not capped -- see v2_plan.md item 1).
+"""
+from typing import Dict, List
+
+import torch
+from torch import nn
+
+from .document import PredictedNER, Sentence
+from .feedforward import FeedForward
+from .metrics import NERMetrics
+from .transformer_block import TransformerBlock
+from .vocab import Vocabulary
+
+
+class NERHead(nn.Module):
+    def __init__(self, vocab: Vocabulary, span_emb_dim: int, feedforward_params: dict,
+                 transformer_params: dict = 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()
+        self.metrics: Dict[str, NERMetrics] = {}
+        for ns in self.namespaces:
+            if self.use_transformer:
+                body = TransformerBlock(span_emb_dim, **transformer_params)
+            else:
+                body = FeedForward(span_emb_dim, feedforward_params["hidden_dims"],
+                                    feedforward_params["dropout"])
+            self.scorers[ns] = body
+            self.classifiers[ns] = nn.Linear(body.output_dim, self.n_labels[ns] - 1)
+            self.metrics[ns] = NERMetrics(self.n_labels[ns])
+
+        self._loss = nn.CrossEntropyLoss(reduction="sum")
+
+    def forward(self, dataset: str, spans: List[tuple], span_embeddings: torch.Tensor,
+                sentence: Sentence, ner_label_ids: torch.Tensor = None) -> dict:
+        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)
+        predicted = scores.argmax(dim=-1)
+
+        predictions = self.decode(ns, scores.detach(), spans, sentence)
+        output = {"namespace": ns, "scores": scores, "predicted": predicted, "predictions": predictions}
+        if ner_label_ids is not None:
+            self.metrics[ns](predicted, ner_label_ids)
+            output["loss"] = self._loss(scores, ner_label_ids)
+        return output
+
+    def decode(self, namespace: str, scores: torch.Tensor, spans: List[tuple],
+               sentence: Sentence) -> List[PredictedNER]:
+        softmax_scores = torch.softmax(scores, dim=-1)
+        raw, predicted = scores.max(dim=-1)
+        soft, _ = softmax_scores.max(dim=-1)
+
+        predictions = []
+        for i in (predicted != 0).nonzero(as_tuple=True)[0].tolist():
+            label = self.vocab.get_token_from_index(int(predicted[i]), namespace)
+            start, end = spans[i]
+            entry = [start, end, label, float(raw[i]), float(soft[i])]
+            predictions.append(PredictedNER(entry, sentence, sentence_offsets=True))
+        return predictions
+
+    def get_metrics(self, reset: bool = False) -> Dict[str, float]:
+        res = {}
+        for ns, metric in self.metrics.items():
+            prefix = ns.replace("_labels", "")
+            per_class = metric.get_per_class()  # {class idx -> F1}, read BEFORE the reset below
+            for idx, class_f1 in per_class.items():
+                res[f"ner_perclass/{self.vocab.get_token_from_index(idx, ns)}"] = class_f1
+            res["ner_macro_f1"] = sum(per_class.values()) / len(per_class) if per_class else 0.0
+            precision, recall, f1 = metric.get_metric(reset)
+            res[f"{prefix}_precision"] = precision
+            res[f"{prefix}_recall"] = recall
+            res[f"{prefix}_f1"] = f1
+        for name in ("precision", "recall", "f1"):
+            values = [v for k, v in res.items()
+                      if k.endswith(f"_{name}") and "perclass/" not in k and "_macro_" not in k]
+            res[f"MEAN__ner_{name}"] = sum(values) / len(values) if values else 0.0
+        return res

predictor.py

@@ -0,0 +1,25 @@
+"""Run a trained model on a document and get back decoded predictions.
+
+Port of the one thing `radgraph.dygie.predictors.dygie.DyGIEPredictor` did that this project
+actually used: forward pass -> `PredictedNER` / `PredictedRelation` objects, in document
+(here: sentence, since there's always exactly one) coordinates.
+"""
+import torch
+
+from .dataset import Example, build_example
+from .document import Document
+from .model import DyGIEModel
+from .vocab import Vocabulary
+
+
+@torch.no_grad()
+def predict_document(model: DyGIEModel, vocab: Vocabulary, doc: Document, max_span_width: int) -> dict:
+    model.eval()
+    example: Example = build_example(doc, vocab, max_span_width)
+    output = model(example)
+    return {
+        "doc_key": doc.doc_key,
+        "example": example,
+        "predicted_ner": output["ner"].get("predictions", []),
+        "predicted_relations": output["relation"].get("predictions", []),
+    }

pruner.py

@@ -0,0 +1,26 @@
+"""Mention pruner: score every span, keep the top-k, in original span order.
+
+Port of `radgraph.dygie.models.entity_beam_pruner.Pruner`'s default path (no entity_beam, no
+gold_beam -- this project's relation module never sets either, see relation_head.py).
+Simplified for batch_size=1: no masking, no batch dimension.
+"""
+from typing import Tuple
+
+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.LongTensor, 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

relation_head.py

@@ -0,0 +1,203 @@
+"""Relation head: prune to top-K spans, score every pair, cross-entropy against gold.
+
+Port of `radgraph.dygie.models.relation.RelationExtractor`. Two deliberate simplifications
+made possible by batch_size=1 (see dataset.py):
+
+  1. No span mask / no ignore_index: every one of the K pruned spans is a real span (there's
+     no cross-document padding to mask out), so every (i, j) pair among them is a legitimate
+     loss target -- either a real relation or "null".
+  2. The scorer outputs `n_labels - 1` channels (real labels only) with a null column
+     prepended, exactly like ner_head.py already does. v1's RelationExtractor scores the
+     *full* `n_labels` (which already includes the null vocab slot at index 0) and *also*
+     prepends a dummy null column, leaving one dead, never-targeted output channel that gets
+     trained toward nothing and just gets suppressed as a side effect of softmax competition.
+     Sizing it like NER removes that dead channel and, as a bonus, means predicted/gold label
+     ids line up with vocabulary indices directly -- no "+1 shift so null is class 0" dance.
+
+`transformer_params` (medbit_transformer arm, see notes/v2_plan.md item 1) swaps the pairwise
+`relation_feedforwards` scorer for a TransformerBlock applied to the pruner's already-pruned
+top-K spans, *before* pair formation: the K surviving spans attend to each other, then pairs
+are built from the now-contextualized embeddings and scored with a bare Linear (the
+transformer's own internal FFN sublayer already did the nonlinear work `relation_feedforwards`
+used to do). The pruner's own mention scorer stays a plain FeedForward -- it's a preliminary
+filter picking which K spans survive, not the classifier itself.
+
+`use_between_context` (medbit_relations_pooling arm, see notes/v2_plan.md item 2) adds a 4th
+block to the pair representation: the mean-pooled *raw* word embeddings strictly between the
+two spans (not the pooled span embeddings -- those don't preserve individual word positions).
+Computed via a prefix sum over `word_embeddings` (`_between_context`) so all K^2 pairs are done
+in one vectorized pass instead of a per-pair loop. Overlapping/adjacent/nested spans have no
+words between them -> zero vector for that pair. Orthogonal to `transformer_params`: the arm
+forks from the plain medbit baseline, not medbit_transformer, so both flags are never on
+together in practice, but the code composes correctly either way (the between-block just adds
+`word_emb_dim` to whichever pair representation -- FeedForward or transformer -- is in use).
+
+`relation_feedforward_params` (medbit_relations_pooling arm only, defaults to `feedforward_params`
+for every other arm) sizes just the relation pair scorer's hidden layers, kept separate from the
+pruner's `mention_ff` -- the pruner still consumes plain `span_emb_dim` (unaffected by
+`use_between_context`), so only the layers that actually see the enlarged pair embedding get
+widened, mirroring the "pruner untouched" principle already applied to `transformer_params`.
+"""
+import math
+from typing import Dict, List, Tuple
+
+import torch
+from torch import nn
+
+from .document import PredictedRelation, Sentence
+from .feedforward import FeedForward
+from .metrics import RelationMetrics
+from .pruner import Pruner
+from .transformer_block import TransformerBlock
+from .vocab import Vocabulary
+
+
+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,
+                 use_between_context: bool = False, relation_feedforward_params: dict = 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()
+        self.metrics: Dict[str, RelationMetrics] = {}
+
+        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:
+                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)
+            self.metrics[ns] = RelationMetrics()
+
+        self._loss = nn.CrossEntropyLoss(reduction="sum")
+
+    @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:
+        """Mean of the raw word embeddings strictly between two spans, one block per (i, j) pair
+        -- vectorized over the whole K x K grid via a prefix sum (no per-pair loop; K^2 pairs
+        can run into the tens of thousands per doc). `low`/`high` don't depend on pair order (i,
+        j) vs (j, i), so the between-region -- and this block -- is symmetric by construction.
+        Overlapping/adjacent/nested spans (count <= 0) get a zero vector."""
+        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))       # end of whichever span is first
+        high = torch.maximum(starts.unsqueeze(1), starts.unsqueeze(0))  # start of whichever span is second
+        count = (high - low - 1).clamp(min=0)                          # words strictly between
+
+        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 forward(self, dataset: str, spans: List[Tuple[int, int]], span_embeddings: torch.Tensor,
+                num_words: int, word_embeddings: torch.Tensor = None,
+                relation_gold: Dict[Tuple[int, int], int] = None,
+                sentence: Sentence = None) -> dict:
+        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_indices_list = top_indices.tolist()
+        top_spans = [spans[i] for i in top_indices_list]
+
+        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
+
+        predicted_dict, predictions = self.decode(ns, scores.detach(), top_spans, sentence)
+        output = {"namespace": ns, "scores": scores, "top_spans": top_spans, "predictions": predictions}
+
+        if relation_gold is not None:
+            pos_of = {orig: pos for pos, orig in enumerate(top_indices_list)}
+            gold_matrix = scores.new_zeros(k, k, dtype=torch.long)
+            for (oi, oj), label_id in relation_gold.items():
+                pi, pj = pos_of.get(oi), pos_of.get(oj)
+                if pi is not None and pj is not None:
+                    gold_matrix[pi, pj] = label_id
+
+            output["loss"] = self._loss(scores.view(k * k, -1), gold_matrix.view(-1))
+            self.metrics[ns](predicted_dict, sentence.relation_dict or {})
+
+        return output
+
+    def decode(self, namespace: str, scores: torch.Tensor, top_spans: List[Tuple[int, int]],
+               sentence: Sentence):
+        raw, predicted = scores.max(dim=-1)
+        soft, _ = torch.softmax(scores, dim=-1).max(dim=-1)
+
+        res_dict, predictions = {}, []
+        for i, j in (predicted != 0).nonzero(as_tuple=False).tolist():
+            label = self.vocab.get_token_from_index(int(predicted[i, j]), namespace)
+            span1, span2 = top_spans[i], top_spans[j]
+            res_dict[(span1, span2)] = label
+            entry = [span1[0], span1[1], span2[0], span2[1], label, float(raw[i, j]), float(soft[i, j])]
+            predictions.append(PredictedRelation(entry, sentence, sentence_offsets=True))
+        return res_dict, predictions
+
+    def get_metrics(self, reset: bool = False) -> Dict[str, float]:
+        res = {}
+        for ns, metric in self.metrics.items():
+            prefix = ns.replace("_labels", "")
+            per_class = metric.get_per_class()  # {relation label -> F1}, read BEFORE the reset below
+            for label, class_f1 in per_class.items():
+                res[f"relation_perclass/{label}"] = class_f1
+            res["relation_macro_f1"] = sum(per_class.values()) / len(per_class) if per_class else 0.0
+            precision, recall, f1 = metric.get_metric(reset)
+            res[f"{prefix}_precision"] = precision
+            res[f"{prefix}_recall"] = recall
+            res[f"{prefix}_f1"] = f1
+        for name in ("precision", "recall", "f1"):
+            values = [v for k, v in res.items()
+                      if k.endswith(f"_{name}") and "perclass/" not in k and "_macro_" not in k]
+            res[f"MEAN__relation_{name}"] = sum(values) / len(values) if values else 0.0
+        return res

span_extractor.py

@@ -0,0 +1,48 @@
+"""Span representations: concat(start word embedding, end word embedding, [mean-pooled interior
+embedding], width embedding).
+
+Port of AllenNLP's `EndpointSpanExtractor(combination="x,y", bucket_widths=False)`, the only
+configuration v1 used. v2 adds an optional mean-pooled block -- the average of every word
+embedding inside [start, end] inclusive -- gated by `mean_pool` (see notes/v2_plan.md item 3,
+`configs/medbit_span_pooling/`). Simplified for the batch_size=1 case (see dataset.py): no batch
+dimension, no span masking -- every document has an exact, un-padded list of spans.
+"""
+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.LongTensor,
+                        ends: torch.LongTensor) -> torch.Tensor:
+        """Mean of word_embeddings[start:end+1] per span, vectorized over the whole (bounded-
+        width) span list: gather a (num_spans, max_width) index grid and masked-average it. No
+        cumsum needed -- widths are capped by max_span_width (e.g. 12)."""
+        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)

tokenizer_config.json

@@ -0,0 +1,19 @@
+{
+  "backend": "tokenizers",
+  "cls_token": "[CLS]",
+  "do_basic_tokenize": true,
+  "do_lower_case": false,
+  "is_local": false,
+  "local_files_only": false,
+  "mask_token": "[MASK]",
+  "max_len": 512,
+  "model_max_length": 512,
+  "never_split": null,
+  "pad_token": "[PAD]",
+  "sep_token": "[SEP]",
+  "strip_accents": null,
+  "tokenize_chinese_chars": true,
+  "tokenizer_class": "BertTokenizer",
+  "truncation": true,
+  "unk_token": "[UNK]"
+}

tokenizer_embedder.py

@@ -0,0 +1,114 @@
+"""Word-level embeddings from a subword transformer encoder.
+
+Port of AllenNLP's `pretrained_transformer_mismatched` (indexer + embedder pair): tokenize
+each word into subwords, run the transformer, mean-pool subwords back to one vector per word
+(see training/README.md: word tokenization is fixed upstream, so the model always sees
+word-level spans). Long documents are split into non-overlapping, word-aligned chunks of at
+most `max_length` transformer positions (the medbit arm's "context folding": medBIT-r3-plus's
+native window is 512 tokens, most reports are longer).
+
+Deliberate deviation from v1: AllenNLP folds the flat *wordpiece* stream into fixed windows,
+so a word's subwords can occasionally land split across two windows, encoded with no shared
+context. Here chunks are packed at word granularity, so a word's subwords are always encoded
+together -- strictly better representation quality at chunk boundaries, same architecture.
+
+Each chunk is re-tokenized and special-token-wrapped through the tokenizer's own public
+`__call__(..., is_split_into_words=True)` API rather than hand-rolling special-token insertion:
+transformers 5.x dropped the old `build_inputs_with_special_tokens` / `prepare_for_model`
+helpers from its fast-tokenizer wrapper (see training_v2/README.md), and going through the
+stable public API is more robust across tokenizer families anyway.
+"""
+import logging
+from typing import List
+
+import torch
+from torch import nn
+from transformers import AutoConfig, AutoModel, AutoTokenizer
+
+logger = logging.getLogger(__name__)
+
+
+class MismatchedEmbedder(nn.Module):
+    def __init__(self, model_name: str, max_length: int, train_parameters: bool = True):
+        super().__init__()
+        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
+        # Built from config, not `AutoModel.from_pretrained`: this HF-packaged copy is only ever
+        # constructed either directly (fresh pretrained-encoder init, matching training_v2's
+        # behavior) or via `RadgraphModel.from_pretrained(...)`, which builds the whole module
+        # tree under a `torch.device("meta")` context and then loads our fine-tuned
+        # model.safetensors over it -- a nested `from_pretrained` call in that path hits
+        # transformers' "meta device anti-pattern" guard. Building from config sidesteps that
+        # and is equivalent here since the encoder weights get overwritten by the checkpoint
+        # load either way.
+        self.encoder = AutoModel.from_config(AutoConfig.from_pretrained(model_name))
+        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] * len(words)  # type: ignore[list-item]
+        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 = {}
+            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)
+                word_embeddings[i] = hidden.new_zeros(self.output_dim)
+
+        return torch.stack(word_embeddings, dim=0)

transformer_block.py

@@ -0,0 +1,50 @@
+"""Transformer-block classifier body: self-attention across the span dimension, used by
+ner_head.py / relation_head.py in the `medbit_transformer` arm in place of the plain
+`FeedForward` scorer (see notes/v2_plan.md item 1).
+
+A standard `nn.TransformerEncoderLayer` already bundles self-attention with its own
+position-wise FFN sublayer (residual + norm around both), so this fully replaces
+`FeedForward` rather than wrapping it -- no separate downstream MLP survives; callers add
+only a final `nn.Linear` to project to label logits, same shape as today.
+
+Spans are the "sequence" here (length = number of candidate/pruned spans, not number of
+words); there is no batch dimension and no padding mask, matching every other module's
+batch_size=1 simplification (see pruner.py, span_extractor.py). A sinusoidal positional
+encoding keyed on each span's start-word index gives attention a notion of span order/distance
+-- the same choice already made in v2_plan.md item 1 for the (superseded) word-level
+context_layer idea, just applied to span positions instead of word positions.
+"""
+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:
+        """span_embeddings: (num_spans, input_dim). start_positions: (num_spans,) int64.
+        No batch dim on the way in/out -- one fake batch of size 1 for nn.TransformerEncoder."""
+        x = self.input_proj(span_embeddings) + sinusoidal_position_encoding(start_positions, self.d_model)
+        return self.encoder(x.unsqueeze(0)).squeeze(0)

vocab.py

@@ -0,0 +1,78 @@
+"""Minimal label vocabulary: per-namespace string<->int maps, null label pinned to index 0.
+
+Port of the one piece of AllenNLP's `Vocabulary` this project actually uses: separate NER /
+relation label namespaces per `dataset` value (e.g. "radgraph-it__ner_labels"), with a
+non-embedded null label "" always at index 0 so "argmax == 0" means "no label" throughout the
+model. Built once from the training split and saved alongside the checkpoint so dev/test/
+inference use the exact same label<->index mapping.
+"""
+import json
+from pathlib import Path
+from typing import Dict, Iterable, List
+
+NULL_LABEL = ""
+
+
+class Vocabulary:
+    def __init__(self):
+        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 from_documents(cls, documents: Iterable, ner_suffix="ner_labels",
+                        relation_suffix="relation_labels") -> "Vocabulary":
+        vocab = cls()
+        for doc in documents:
+            ner_ns = f"{doc.dataset}__{ner_suffix}"
+            rel_ns = f"{doc.dataset}__{relation_suffix}"
+            vocab.add_namespace(ner_ns)
+            vocab.add_namespace(rel_ns)
+            for sent in doc.sentences:
+                for e in sent.ner or []:
+                    vocab.add_token(e.label, ner_ns)
+                for r in sent.relations or []:
+                    vocab.add_token(r.label, rel_ns)
+        return vocab
+
+    def to_dict(self) -> dict:
+        return {ns: dict(sorted(t2i.items(), key=lambda kv: kv[1]))
+                for ns, t2i in self._token_to_index.items()}
+
+    def save(self, path: str) -> None:
+        Path(path).write_text(json.dumps(self.to_dict(), ensure_ascii=False, indent=2))
+
+    @classmethod
+    def load(cls, path: str) -> "Vocabulary":
+        vocab = cls()
+        for ns, t2i in json.loads(Path(path).read_text()).items():
+            vocab.add_namespace(ns)
+            for token, idx in sorted(t2i.items(), key=lambda kv: kv[1]):
+                if token != NULL_LABEL:
+                    vocab.add_token(token, ns)
+        return vocab