File size: 2,231 Bytes
7c5e40e | 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 | """Canonical Universal Dependencies label vocabularies (basic UD, v1).
These fixed id maps make the model's graph heads supervised targets stable and
reproducible. Two label spaces:
* Node types <- UPOS (17 universal POS tags) -> ``node_type_logits``.
* Relations <- universal deprel (subtype after ':' dropped) -> arc/rel head.
A model trained with UD supervision must have ``node_type_vocab_size`` >=
``NUM_UPOS`` and ``graph_relation_types`` >= ``NUM_DEPREL``.
"""
from __future__ import annotations
# 17 universal POS tags (UD v2).
UPOS_TAGS: tuple[str, ...] = (
"ADJ", "ADP", "ADV", "AUX", "CCONJ", "DET", "INTJ", "NOUN", "NUM",
"PART", "PRON", "PROPN", "PUNCT", "SCONJ", "SYM", "VERB", "X",
)
UPOS_TO_ID: dict[str, int] = {tag: i for i, tag in enumerate(UPOS_TAGS)}
NUM_UPOS = len(UPOS_TAGS)
PUNCT_UPOS_ID = UPOS_TO_ID["PUNCT"]
# 37 universal syntactic relations (UD v2), including ``root``.
UD_DEPRELS: tuple[str, ...] = (
"acl", "advcl", "advmod", "amod", "appos", "aux", "case", "cc", "ccomp",
"clf", "compound", "conj", "cop", "csubj", "dep", "det", "discourse",
"dislocated", "expl", "fixed", "flat", "goeswith", "iobj", "list", "mark",
"nmod", "nsubj", "nummod", "obj", "obl", "orphan", "parataxis", "punct",
"reparandum", "root", "vocative", "xcomp",
)
DEPREL_TO_ID: dict[str, int] = {rel: i for i, rel in enumerate(UD_DEPRELS)}
NUM_DEPREL = len(UD_DEPRELS)
PUNCT_DEPREL_ID = DEPREL_TO_ID["punct"]
ROOT_DEPREL_ID = DEPREL_TO_ID["root"]
# Ignore index shared with torch cross-entropy for masked/unsupervised positions.
IGNORE_INDEX = -100
def universal_deprel(deprel: str) -> str:
"""Drop the language-specific subtype (``nsubj:pass`` -> ``nsubj``)."""
return deprel.split(":", 1)[0].lower()
def deprel_to_id(deprel: str) -> int:
"""Map a (possibly subtyped) deprel to a universal relation id.
Unknown relations fall back to the generic ``dep`` bucket rather than
failing, so a new treebank cannot crash training.
"""
return DEPREL_TO_ID.get(universal_deprel(deprel), DEPREL_TO_ID["dep"])
def upos_to_id(upos: str) -> int | None:
"""Map a UPOS tag to a node-type id, or ``None`` if unrecognised."""
return UPOS_TO_ID.get(upos.upper())
|