File size: 5,195 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 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 | """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]
|