File size: 3,081 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 | """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
|