Add transformers-compatible wrapper (AutoModel/AutoTokenizer via trust_remote_code)
0c48771 verified | """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") | |
| 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) | |
| 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 | |