File size: 4,997 Bytes
0c48771 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 | """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
|