Model save

README.md

@@ -0,0 +1,57 @@
+---
+base_model: xlm-roberta-base
+datasets: CoBaLD/enhanced-ud-syntax
+language: en
+library_name: transformers
+license: gpl-3.0
+metrics:
+- accuracy
+- f1
+pipeline_tag: cobald-parsing
+tags:
+- pytorch
+model-index:
+- name: CoBaLD/cobald-parser-pretrain-en
+  results:
+  - task:
+      type: token-classification
+    dataset:
+      name: enhanced-ud-syntax
+      type: CoBaLD/enhanced-ud-syntax
+      split: validation
+    metrics:
+    - type: f1
+      value: 0.2499754084433992
+      name: Null F1
+    - type: accuracy
+      value: 0.6778357854769815
+      name: Ud Jaccard
+    - type: accuracy
+      value: 0.74552842927079
+      name: Eud Jaccard
+---
+
+# Model Card for cobald-parser-pretrain-en
+
+A transformer-based multihead parser for CoBaLD annotation.
+
+This model parses a pre-tokenized CoNLL-U text and jointly labels each token with three tiers of tags:
+* Grammatical tags (lemma, UPOS, XPOS, morphological features),
+* Syntactic tags (basic and enhanced Universal Dependencies),
+* Semantic tags (deep slot and semantic class).
+
+## Model Sources
+
+- **Repository:** https://github.com/CobaldAnnotation/CobaldParser
+- **Paper:** https://dialogue-conf.org/wp-content/uploads/2025/04/BaiukIBaiukAPetrovaM.009.pdf
+- **Demo:** [coming soon]
+
+## Citation
+
+@inproceedings{baiuk2025cobald,
+  title={CoBaLD Parser: Joint Morphosyntactic and Semantic Annotation},
+  author={Baiuk, Ilia and Baiuk, Alexandra and Petrova, Maria},
+  booktitle={Proceedings of the International Conference "Dialogue"},
+  volume={I},
+  year={2025}
+}

config.json

@@ -0,0 +1,512 @@
+{
+  "activation": "relu",
+  "architectures": [
+    "CobaldParser"
+  ],
+  "auto_map": {
+    "AutoConfig": "configuration.CobaldParserConfig",
+    "AutoModel": "modeling_parser.CobaldParser"
+  },
+  "consecutive_null_limit": 3,
+  "custom_pipelines": {
+    "cobald-parsing": {
+      "impl": "pipeline.ConlluTokenClassificationPipeline",
+      "pt": "CobaldParser"
+    }
+  },
+  "deepslot_classifier_hidden_size": 256,
+  "dependency_classifier_hidden_size": 128,
+  "dropout": 0.1,
+  "encoder_model_name": "xlm-roberta-base",
+  "lemma_classifier_hidden_size": 512,
+  "misc_classifier_hidden_size": 512,
+  "model_type": "cobald_parser",
+  "morphology_classifier_hidden_size": 512,
+  "null_classifier_hidden_size": 512,
+  "semclass_classifier_hidden_size": 512,
+  "torch_dtype": "float32",
+  "transformers_version": "4.51.3",
+  "vocabulary": {
+    "eud_deprel": {
+      "0": "acl",
+      "1": "acl:about",
+      "2": "acl:about_whether",
+      "3": "acl:after",
+      "4": "acl:against",
+      "5": "acl:as",
+      "6": "acl:as_if",
+      "7": "acl:as_to",
+      "8": "acl:at",
+      "9": "acl:before",
+      "10": "acl:behind",
+      "11": "acl:between",
+      "12": "acl:beyond",
+      "13": "acl:but",
+      "14": "acl:but_to",
+      "15": "acl:concerning",
+      "16": "acl:except_that",
+      "17": "acl:for",
+      "18": "acl:for_to",
+      "19": "acl:from",
+      "20": "acl:if",
+      "21": "acl:in",
+      "22": "acl:including",
+      "23": "acl:including_whether",
+      "24": "acl:inside",
+      "25": "acl:instead_of",
+      "26": "acl:into",
+      "27": "acl:like",
+      "28": "acl:of",
+      "29": "acl:of_if",
+      "30": "acl:of_why",
+      "31": "acl:on",
+      "32": "acl:once",
+      "33": "acl:over",
+      "34": "acl:prior_to",
+      "35": "acl:regarding",
+      "36": "acl:relcl",
+      "37": "acl:relcl:to",
+      "38": "acl:since",
+      "39": "acl:such_as",
+      "40": "acl:than",
+      "41": "acl:that",
+      "42": "acl:though",
+      "43": "acl:to",
+      "44": "acl:toward",
+      "45": "acl:towards",
+      "46": "acl:under",
+      "47": "acl:until",
+      "48": "acl:upon",
+      "49": "acl:when",
+      "50": "acl:where",
+      "51": "acl:whether",
+      "52": "acl:why",
+      "53": "acl:with",
+      "54": "advcl",
+      "55": "advcl:about",
+      "56": "advcl:about_whether",
+      "57": "advcl:after",
+      "58": "advcl:against",
+      "59": "advcl:albeit",
+      "60": "advcl:along_with",
+      "61": "advcl:although",
+      "62": "advcl:as",
+      "63": "advcl:as_if",
+      "64": "advcl:as_in",
+      "65": "advcl:as_though",
+      "66": "advcl:as_to",
+      "67": "advcl:as_well_as",
+      "68": "advcl:as_with",
+      "69": "advcl:at",
+      "70": "advcl:because",
+      "71": "advcl:before",
+      "72": "advcl:behind",
+      "73": "advcl:besides",
+      "74": "advcl:between",
+      "75": "advcl:beyond",
+      "76": "advcl:but",
+      "77": "advcl:by",
+      "78": "advcl:cause",
+      "79": "advcl:despite",
+      "80": "advcl:due_to",
+      "81": "advcl:except",
+      "82": "advcl:except_for",
+      "83": "advcl:except_that",
+      "84": "advcl:for",
+      "85": "advcl:for_if",
+      "86": "advcl:for_to",
+      "87": "advcl:from",
+      "88": "advcl:given",
+      "89": "advcl:if",
+      "90": "advcl:if_to",
+      "91": "advcl:in",
+      "92": "advcl:in_between",
+      "93": "advcl:in_case",
+      "94": "advcl:in_order",
+      "95": "advcl:in_order_for",
+      "96": "advcl:in_order_to",
+      "97": "advcl:in_that",
+      "98": "advcl:including_by",
+      "99": "advcl:inside",
+      "100": "advcl:insofar_as",
+      "101": "advcl:instead_of",
+      "102": "advcl:into",
+      "103": "advcl:lest",
+      "104": "advcl:like",
+      "105": "advcl:of",
+      "106": "advcl:of_whether",
+      "107": "advcl:on",
+      "108": "advcl:on_whether",
+      "109": "advcl:once",
+      "110": "advcl:out",
+      "111": "advcl:over",
+      "112": "advcl:past",
+      "113": "advcl:prior_to",
+      "114": "advcl:provided",
+      "115": "advcl:rather_than",
+      "116": "advcl:relcl",
+      "117": "advcl:relcl:because",
+      "118": "advcl:since",
+      "119": "advcl:so",
+      "120": "advcl:so_as_to",
+      "121": "advcl:so_that",
+      "122": "advcl:such_as",
+      "123": "advcl:than",
+      "124": "advcl:than_if",
+      "125": "advcl:that",
+      "126": "advcl:the",
+      "127": "advcl:though",
+      "128": "advcl:through",
+      "129": "advcl:till",
+      "130": "advcl:to",
+      "131": "advcl:toward",
+      "132": "advcl:towards",
+      "133": "advcl:under",
+      "134": "advcl:unless",
+      "135": "advcl:until",
+      "136": "advcl:upon",
+      "137": "advcl:when",
+      "138": "advcl:where",
+      "139": "advcl:whereas",
+      "140": "advcl:whether",
+      "141": "advcl:while",
+      "142": "advcl:whilst",
+      "143": "advcl:whither",
+      "144": "advcl:with",
+      "145": "advcl:without",
+      "146": "advmod",
+      "147": "amod",
+      "148": "appos",
+      "149": "aux",
+      "150": "aux:pass",
+      "151": "case",
+      "152": "case:of",
+      "153": "cc",
+      "154": "cc:preconj",
+      "155": "ccomp",
+      "156": "compound",
+      "157": "compound:prt",
+      "158": "conj",
+      "159": "conj:and",
+      "160": "conj:and_or",
+      "161": "conj:and_yet",
+      "162": "conj:as_well_as",
+      "163": "conj:but",
+      "164": "conj:et",
+      "165": "conj:for",
+      "166": "conj:let_alone",
+      "167": "conj:minus",
+      "168": "conj:nor",
+      "169": "conj:not",
+      "170": "conj:not_to_mention",
+      "171": "conj:or",
+      "172": "conj:plus",
+      "173": "conj:plus_minus",
+      "174": "conj:rather_than",
+      "175": "conj:slash",
+      "176": "conj:though",
+      "177": "conj:yet",
+      "178": "cop",
+      "179": "csubj",
+      "180": "csubj:outer",
+      "181": "csubj:pass",
+      "182": "csubj:xsubj",
+      "183": "dep",
+      "184": "det",
+      "185": "det:predet",
+      "186": "discourse",
+      "187": "dislocated",
+      "188": "expl",
+      "189": "fixed",
+      "190": "flat",
+      "191": "goeswith",
+      "192": "iobj",
+      "193": "list",
+      "194": "mark",
+      "195": "nmod",
+      "196": "nmod:a_la",
+      "197": "nmod:aboard",
+      "198": "nmod:about",
+      "199": "nmod:above",
+      "200": "nmod:according_to",
+      "201": "nmod:across",
+      "202": "nmod:after",
+      "203": "nmod:against",
+      "204": "nmod:along",
+      "205": "nmod:alongside",
+      "206": "nmod:amidst",
+      "207": "nmod:among",
+      "208": "nmod:amongst",
+      "209": "nmod:around",
+      "210": "nmod:as",
+      "211": "nmod:as_for",
+      "212": "nmod:as_in",
+      "213": "nmod:as_opposed_to",
+      "214": "nmod:as_to",
+      "215": "nmod:astride",
+      "216": "nmod:at",
+      "217": "nmod:atop",
+      "218": "nmod:barring",
+      "219": "nmod:because_of",
+      "220": "nmod:before",
+      "221": "nmod:behind",
+      "222": "nmod:below",
+      "223": "nmod:besides",
+      "224": "nmod:between",
+      "225": "nmod:beyond",
+      "226": "nmod:but",
+      "227": "nmod:by",
+      "228": "nmod:circa",
+      "229": "nmod:colon",
+      "230": "nmod:concerning",
+      "231": "nmod:desc",
+      "232": "nmod:despite",
+      "233": "nmod:down",
+      "234": "nmod:due_to",
+      "235": "nmod:during",
+      "236": "nmod:except",
+      "237": "nmod:except_for",
+      "238": "nmod:excluding",
+      "239": "nmod:following",
+      "240": "nmod:for",
+      "241": "nmod:from",
+      "242": "nmod:from_across",
+      "243": "nmod:from_below",
+      "244": "nmod:from_outside",
+      "245": "nmod:from_over",
+      "246": "nmod:in",
+      "247": "nmod:including",
+      "248": "nmod:inside",
+      "249": "nmod:instead_of",
+      "250": "nmod:into",
+      "251": "nmod:like",
+      "252": "nmod:minus",
+      "253": "nmod:near",
+      "254": "nmod:next_to",
+      "255": "nmod:of",
+      "256": "nmod:off",
+      "257": "nmod:on",
+      "258": "nmod:onto",
+      "259": "nmod:opposite",
+      "260": "nmod:out",
+      "261": "nmod:out_of",
+      "262": "nmod:outside",
+      "263": "nmod:over",
+      "264": "nmod:past",
+      "265": "nmod:per",
+      "266": "nmod:poss",
+      "267": "nmod:post",
+      "268": "nmod:prior_to",
+      "269": "nmod:pro",
+      "270": "nmod:re",
+      "271": "nmod:regarding",
+      "272": "nmod:round",
+      "273": "nmod:save",
+      "274": "nmod:since",
+      "275": "nmod:slash",
+      "276": "nmod:such_as",
+      "277": "nmod:than",
+      "278": "nmod:through",
+      "279": "nmod:throughout",
+      "280": "nmod:thru",
+      "281": "nmod:times",
+      "282": "nmod:to",
+      "283": "nmod:toward",
+      "284": "nmod:towards",
+      "285": "nmod:under",
+      "286": "nmod:unlike",
+      "287": "nmod:unmarked",
+      "288": "nmod:until",
+      "289": "nmod:up",
+      "290": "nmod:upon",
+      "291": "nmod:versus",
+      "292": "nmod:via",
+      "293": "nmod:with",
+      "294": "nmod:within",
+      "295": "nmod:without",
+      "296": "nmod:x",
+      "297": "nsubj",
+      "298": "nsubj:outer",
+      "299": "nsubj:pass",
+      "300": "nsubj:pass:xsubj",
+      "301": "nsubj:xsubj",
+      "302": "nummod",
+      "303": "obj",
+      "304": "obl",
+      "305": "obl:aboard",
+      "306": "obl:about",
+      "307": "obl:above",
+      "308": "obl:according_to",
+      "309": "obl:across",
+      "310": "obl:after",
+      "311": "obl:against",
+      "312": "obl:agent",
+      "313": "obl:along",
+      "314": "obl:along_with",
+      "315": "obl:alongside",
+      "316": "obl:amid",
+      "317": "obl:amidst",
+      "318": "obl:among",
+      "319": "obl:amongst",
+      "320": "obl:around",
+      "321": "obl:as",
+      "322": "obl:as_for",
+      "323": "obl:as_in",
+      "324": "obl:as_of",
+      "325": "obl:as_opposed_to",
+      "326": "obl:as_to",
+      "327": "obl:aside",
+      "328": "obl:at",
+      "329": "obl:atop",
+      "330": "obl:because_of",
+      "331": "obl:before",
+      "332": "obl:behind",
+      "333": "obl:below",
+      "334": "obl:beneath",
+      "335": "obl:beside",
+      "336": "obl:besides",
+      "337": "obl:between",
+      "338": "obl:beyond",
+      "339": "obl:but",
+      "340": "obl:by",
+      "341": "obl:circa",
+      "342": "obl:concerning",
+      "343": "obl:depending",
+      "344": "obl:depending_on",
+      "345": "obl:depending_upon",
+      "346": "obl:despite",
+      "347": "obl:down",
+      "348": "obl:due_to",
+      "349": "obl:during",
+      "350": "obl:except",
+      "351": "obl:except_for",
+      "352": "obl:following",
+      "353": "obl:for",
+      "354": "obl:for_post",
+      "355": "obl:from",
+      "356": "obl:from_across",
+      "357": "obl:from_among",
+      "358": "obl:from_behind",
+      "359": "obl:from_over",
+      "360": "obl:given",
+      "361": "obl:in",
+      "362": "obl:in_between",
+      "363": "obl:in_case_of",
+      "364": "obl:including",
+      "365": "obl:including_before",
+      "366": "obl:including_for",
+      "367": "obl:including_in",
+      "368": "obl:inside",
+      "369": "obl:instead_of",
+      "370": "obl:into",
+      "371": "obl:like",
+      "372": "obl:minus",
+      "373": "obl:near",
+      "374": "obl:nearby",
+      "375": "obl:nigh",
+      "376": "obl:notwithstanding",
+      "377": "obl:of",
+      "378": "obl:off",
+      "379": "obl:off_of",
+      "380": "obl:on",
+      "381": "obl:on_board",
+      "382": "obl:onto",
+      "383": "obl:opposite",
+      "384": "obl:out",
+      "385": "obl:out_of",
+      "386": "obl:outside",
+      "387": "obl:over",
+      "388": "obl:past",
+      "389": "obl:per",
+      "390": "obl:post",
+      "391": "obl:prior_to",
+      "392": "obl:re",
+      "393": "obl:regarding",
+      "394": "obl:round",
+      "395": "obl:since",
+      "396": "obl:than",
+      "397": "obl:through",
+      "398": "obl:throughout",
+      "399": "obl:thru",
+      "400": "obl:till",
+      "401": "obl:to",
+      "402": "obl:to_before",
+      "403": "obl:toward",
+      "404": "obl:towards",
+      "405": "obl:under",
+      "406": "obl:underneath",
+      "407": "obl:unlike",
+      "408": "obl:unmarked",
+      "409": "obl:until",
+      "410": "obl:unto",
+      "411": "obl:up",
+      "412": "obl:up_to",
+      "413": "obl:upon",
+      "414": "obl:versus",
+      "415": "obl:via",
+      "416": "obl:with",
+      "417": "obl:within",
+      "418": "obl:without",
+      "419": "parataxis",
+      "420": "punct",
+      "421": "ref",
+      "422": "reparandum",
+      "423": "root",
+      "424": "vocative",
+      "425": "xcomp"
+    },
+    "ud_deprel": {
+      "0": "acl",
+      "1": "acl:relcl",
+      "2": "advcl",
+      "3": "advcl:relcl",
+      "4": "advmod",
+      "5": "amod",
+      "6": "appos",
+      "7": "aux",
+      "8": "aux:pass",
+      "9": "case",
+      "10": "cc",
+      "11": "cc:preconj",
+      "12": "ccomp",
+      "13": "compound",
+      "14": "compound:prt",
+      "15": "conj",
+      "16": "cop",
+      "17": "csubj",
+      "18": "csubj:outer",
+      "19": "csubj:pass",
+      "20": "dep",
+      "21": "det",
+      "22": "det:predet",
+      "23": "discourse",
+      "24": "dislocated",
+      "25": "expl",
+      "26": "fixed",
+      "27": "flat",
+      "28": "goeswith",
+      "29": "iobj",
+      "30": "list",
+      "31": "mark",
+      "32": "nmod",
+      "33": "nmod:desc",
+      "34": "nmod:poss",
+      "35": "nmod:unmarked",
+      "36": "nsubj",
+      "37": "nsubj:outer",
+      "38": "nsubj:pass",
+      "39": "nummod",
+      "40": "obj",
+      "41": "obl",
+      "42": "obl:agent",
+      "43": "obl:unmarked",
+      "44": "orphan",
+      "45": "parataxis",
+      "46": "punct",
+      "47": "reparandum",
+      "48": "root",
+      "49": "vocative",
+      "50": "xcomp"
+    }
+  }
+}

configuration.py

@@ -0,0 +1,48 @@
+from transformers import PretrainedConfig
+
+
+class CobaldParserConfig(PretrainedConfig):
+    model_type = "cobald_parser"
+
+    def __init__(
+        self,
+        encoder_model_name: str = None,
+        null_classifier_hidden_size: int = 0,
+        lemma_classifier_hidden_size: int = 0,
+        morphology_classifier_hidden_size: int = 0,
+        dependency_classifier_hidden_size: int = 0,
+        misc_classifier_hidden_size: int = 0,
+        deepslot_classifier_hidden_size: int = 0,
+        semclass_classifier_hidden_size: int = 0,
+        activation: str = 'relu',
+        dropout: float = 0.1,
+        consecutive_null_limit: int = 0,
+        vocabulary: dict[dict[int, str]] = {},
+        **kwargs
+    ):
+        self.encoder_model_name = encoder_model_name
+        self.null_classifier_hidden_size = null_classifier_hidden_size
+        self.consecutive_null_limit = consecutive_null_limit
+        self.lemma_classifier_hidden_size = lemma_classifier_hidden_size
+        self.morphology_classifier_hidden_size = morphology_classifier_hidden_size
+        self.dependency_classifier_hidden_size = dependency_classifier_hidden_size
+        self.misc_classifier_hidden_size = misc_classifier_hidden_size
+        self.deepslot_classifier_hidden_size = deepslot_classifier_hidden_size
+        self.semclass_classifier_hidden_size = semclass_classifier_hidden_size
+        self.activation = activation
+        self.dropout = dropout
+        # The serialized config stores mappings as strings,
+        # e.g. {"0": "acl", "1": "conj"}, so we have to convert them to int.
+        self.vocabulary = {
+            column: {int(k): v for k, v in labels.items()}
+            for column, labels in vocabulary.items()
+        }
+        # HACK: Tell HF hub about custom pipeline.
+        # It should not be hardcoded like this but other workaround are worse imo.
+        self.custom_pipelines = {
+            "cobald-parsing": {
+                "impl": "pipeline.ConlluTokenClassificationPipeline",
+                "pt": "CobaldParser",
+            }
+        }
+        super().__init__(**kwargs)

dependency_classifier.py

@@ -0,0 +1,299 @@
+from typing import override
+from copy import deepcopy
+
+import numpy as np
+
+import torch
+from torch import nn
+from torch import Tensor, FloatTensor, BoolTensor, LongTensor
+import torch.nn.functional as F
+
+from transformers.activations import ACT2FN
+
+from cobald_parser.bilinear_matrix_attention import BilinearMatrixAttention
+from cobald_parser.chu_liu_edmonds import decode_mst
+from cobald_parser.utils import pairwise_mask, replace_masked_values
+
+
+class DependencyHeadBase(nn.Module):
+    """
+    Base class for scoring arcs and relations between tokens in a dependency tree/graph.
+    """
+
+    def __init__(self, hidden_size: int, n_rels: int):
+        super().__init__()
+
+        self.arc_attention = BilinearMatrixAttention(
+            hidden_size,
+            hidden_size,
+            use_input_biases=True,
+            n_labels=1
+        )
+        self.rel_attention = BilinearMatrixAttention(
+            hidden_size,
+            hidden_size,
+            use_input_biases=True,
+            n_labels=n_rels
+        )
+
+    def forward(
+        self,
+        h_arc_head: Tensor,    # [batch_size, seq_len, hidden_size]
+        h_arc_dep: Tensor,     # ...
+        h_rel_head: Tensor,    # ...
+        h_rel_dep: Tensor,     # ...
+        gold_arcs: LongTensor, # [batch_size, seq_len, seq_len]
+        mask: BoolTensor       # [batch_size, seq_len]
+    ) -> dict[str, Tensor]:
+
+        # Score arcs.
+        # s_arc[:, i, j] = score of edge j -> i.
+        s_arc = self.arc_attention(h_arc_head, h_arc_dep)
+        # Mask undesirable values (padding, nulls, etc.) with -inf.
+        replace_masked_values(s_arc, pairwise_mask(mask), replace_with=-1e8)
+        # Score arcs' relations.
+        # [batch_size, seq_len, seq_len, num_labels]
+        s_rel = self.rel_attention(h_rel_head, h_rel_dep).permute(0, 2, 3, 1)
+
+        # Calculate loss.
+        loss = 0.0
+        if gold_arcs is not None:
+            loss += self.calc_arc_loss(s_arc, gold_arcs)
+            loss += self.calc_rel_loss(s_rel, gold_arcs)
+
+        # Predict arcs based on the scores.
+        # [batch_size, seq_len, seq_len]
+        pred_arcs_3d = self.predict_arcs(s_arc, mask)
+        # [batch_size, seq_len, seq_len]
+        pred_rels_3d = self.predict_rels(s_rel)
+        # [n_pred_arcs, 4]
+        preds_combined = self.combine_arcs_rels(pred_arcs_3d, pred_rels_3d)
+        return {
+            'preds': preds_combined,
+            'loss': loss
+        }
+
+    @staticmethod
+    def calc_arc_loss(
+        s_arc: Tensor,         # [batch_size, seq_len, seq_len]
+        gold_arcs: LongTensor  # [n_arcs, 4]
+    ) -> Tensor:
+        """Calculate arc loss."""
+        raise NotImplementedError
+
+    @staticmethod
+    def calc_rel_loss(
+        s_rel: Tensor,         # [batch_size, seq_len, seq_len, num_labels]
+        gold_arcs: LongTensor  # [n_arcs, 4]
+    ) -> Tensor:
+        batch_idxs, arcs_from, arcs_to, rels = gold_arcs.T
+        return F.cross_entropy(s_rel[batch_idxs, arcs_from, arcs_to], rels)
+    
+    def predict_arcs(
+        self,
+        s_arc: Tensor,   # [batch_size, seq_len, seq_len]
+        mask: BoolTensor # [batch_size, seq_len]
+    ) -> LongTensor:
+        """Predict arcs from scores."""
+        raise NotImplementedError
+    
+    def predict_rels(
+        self,
+        s_rel: FloatTensor
+    ) -> LongTensor:
+        return s_rel.argmax(dim=-1).long()
+    
+    @staticmethod
+    def combine_arcs_rels(
+        pred_arcs: LongTensor,
+        pred_rels: LongTensor
+    ) -> LongTensor:
+        """Select relations towards predicted arcs."""
+        assert pred_arcs.shape == pred_rels.shape
+        # Get indices where arcs exist
+        indices = pred_arcs.nonzero(as_tuple=True)
+        batch_idxs, from_idxs, to_idxs = indices
+        # Get corresponding relation types
+        rel_types = pred_rels[batch_idxs, from_idxs, to_idxs]
+        # Stack as [batch_idx, from_idx, to_idx, rel_type]
+        return torch.stack([batch_idxs, from_idxs, to_idxs, rel_types], dim=1)
+
+
+class DependencyHead(DependencyHeadBase):
+    """
+    Basic UD syntax specialization that predicts single edge for each token.
+    """
+
+    @override
+    def predict_arcs(
+        self,
+        s_arc: Tensor,   # [batch_size, seq_len, seq_len]
+        mask: BoolTensor # [batch_size, seq_len]
+    ) -> Tensor:
+
+        if self.training:
+            # During training, use fast greedy decoding.
+            # - [batch_size, seq_len]
+            pred_arcs_seq = s_arc.argmax(dim=-1)
+        else:
+            # During inference, diligently decode Maximum Spanning Tree.
+            pred_arcs_seq = self._mst_decode(s_arc, mask)
+            # FIXME
+            # pred_arcs_seq = s_arc.argmax(dim=-1)
+
+        # Upscale arcs sequence of shape [batch_size, seq_len]
+        # to matrix of shape [batch_size, seq_len, seq_len].
+        pred_arcs = F.one_hot(pred_arcs_seq, num_classes=pred_arcs_seq.size(1)).long()
+        return pred_arcs
+
+    def _mst_decode(
+        self,
+        s_arc: Tensor, # [batch_size, seq_len, seq_len]
+        mask: Tensor   # [batch_size, seq_len]
+    ) -> tuple[Tensor, Tensor]:
+
+        batch_size = s_arc.size(0)
+        device = s_arc.device
+        s_arc = s_arc.cpu()
+
+        # Convert scores to probabilities, as `decode_mst` expects non-negative values.
+        arc_probs = nn.functional.softmax(s_arc, dim=-1)
+        # Transpose arcs, because decode_mst defines 'energy' matrix as
+        #  energy[i,j] = "Score that `i` is the head of `j`",
+        # whereas
+        #  arc_probs[i,j] = "Probability that `j` is the head of `i`".
+        arc_probs = arc_probs.transpose(1, 2)
+
+        # `decode_mst` knows nothing about UD and ROOT, so we have to manually
+        # zero probabilities of arcs leading to ROOT to make sure ROOT is a source node
+        # of a graph.
+
+        # Decode ROOT positions from diagonals.
+        # shape: [batch_size]
+        root_idxs = arc_probs.diagonal(dim1=1, dim2=2).argmax(dim=-1)
+        # Zero out arcs leading to ROOTs.
+        arc_probs[torch.arange(batch_size), :, root_idxs] = 0.0
+
+        pred_arcs = []
+        for sample_idx in range(batch_size):
+            energy = arc_probs[sample_idx]
+            # has_labels=False because we will decode them manually later.
+            lengths = mask[sample_idx].sum()
+            heads, _ = decode_mst(energy, lengths, has_labels=False)
+            # Some nodes may be isolated. Pick heads greedily in this case.
+            heads[heads <= 0] = s_arc[sample_idx].argmax(dim=-1)[heads <= 0]
+            pred_arcs.append(heads)
+
+        # shape: [batch_size, seq_len]
+        pred_arcs = torch.from_numpy(np.stack(pred_arcs)).long().to(device)
+        return pred_arcs
+
+    @staticmethod
+    @override
+    def calc_arc_loss(
+        s_arc: Tensor,         # [batch_size, seq_len, seq_len]
+        gold_arcs: LongTensor  # [n_arcs, 4]
+    ) -> tuple[Tensor, Tensor]:
+        batch_idxs, from_idxs, to_idxs, _ = gold_arcs.T
+        return F.cross_entropy(s_arc[batch_idxs, from_idxs], to_idxs)
+
+
+class MultiDependencyHead(DependencyHeadBase):
+    """
+    Enhanced UD syntax specialization that predicts multiple edges for each token.
+    """
+
+    @override
+    def predict_arcs(
+        self,
+        s_arc: Tensor,   # [batch_size, seq_len, seq_len]
+        mask: BoolTensor # [batch_size, seq_len]
+    ) -> Tensor:
+        # Convert scores to probabilities.
+        arc_probs = torch.sigmoid(s_arc)
+        # Find confident arcs (with prob > 0.5).
+        return arc_probs.round().long()
+
+    @staticmethod
+    @override
+    def calc_arc_loss(
+        s_arc: Tensor,         # [batch_size, seq_len, seq_len]
+        gold_arcs: LongTensor  # [n_arcs, 4]
+    ) -> Tensor:
+        batch_idxs, from_idxs, to_idxs, _ = gold_arcs.T
+        # Gold arcs but as a matrix, where matrix[i, arcs_from, arc_to] = 1.0 if arcs is present.
+        gold_arcs_matrix = torch.zeros_like(s_arc)
+        gold_arcs_matrix[batch_idxs, from_idxs, to_idxs] = 1.0
+        # Padded arcs's logits are huge negative values that doesn't contribute to the loss.
+        return F.binary_cross_entropy_with_logits(s_arc, gold_arcs_matrix)
+
+
+class DependencyClassifier(nn.Module):
+    """
+    Dozat and Manning's biaffine dependency classifier.
+    """
+
+    def __init__(
+        self,
+        input_size: int,
+        hidden_size: int,
+        n_rels_ud: int,
+        n_rels_eud: int,
+        activation: str,
+        dropout: float,
+    ):
+        super().__init__()
+
+        self.arc_dep_mlp = nn.Sequential(
+            nn.Dropout(dropout),
+            nn.Linear(input_size, hidden_size),
+            ACT2FN[activation],
+            nn.Dropout(dropout)
+        )
+        # All mlps are equal.
+        self.arc_head_mlp = deepcopy(self.arc_dep_mlp)
+        self.rel_dep_mlp = deepcopy(self.arc_dep_mlp)
+        self.rel_head_mlp = deepcopy(self.arc_dep_mlp)
+
+        self.dependency_head_ud = DependencyHead(hidden_size, n_rels_ud)
+        self.dependency_head_eud = MultiDependencyHead(hidden_size, n_rels_eud)
+
+    def forward(
+        self,
+        embeddings: Tensor,    # [batch_size, seq_len, embedding_size]
+        gold_ud: Tensor,       # [n_ud_arcs, 4]
+        gold_eud: Tensor,      # [n_eud_arcs, 4]
+        mask_ud: Tensor,       # [batch_size, seq_len]
+        mask_eud: Tensor       # [batch_size, seq_len]
+    ) -> dict[str, Tensor]:
+
+        # - [batch_size, seq_len, hidden_size]
+        h_arc_head = self.arc_head_mlp(embeddings)
+        h_arc_dep = self.arc_dep_mlp(embeddings)
+        h_rel_head = self.rel_head_mlp(embeddings)
+        h_rel_dep = self.rel_dep_mlp(embeddings)
+
+        # Share the h vectors between dependency and multi-dependency heads.
+        output_ud = self.dependency_head_ud(
+            h_arc_head,
+            h_arc_dep,
+            h_rel_head,
+            h_rel_dep,
+            gold_arcs=gold_ud,
+            mask=mask_ud
+        )
+        output_eud = self.dependency_head_eud(
+            h_arc_head,
+            h_arc_dep,
+            h_rel_head,
+            h_rel_dep,
+            gold_arcs=gold_eud,
+            mask=mask_eud
+        )
+
+        return {
+            'preds_ud': output_ud["preds"],
+            'preds_eud': output_eud["preds"],
+            'loss_ud': output_ud["loss"],
+            'loss_eud': output_eud["loss"]
+        }

encoder.py

@@ -0,0 +1,109 @@
+import torch
+from torch import nn
+from torch import Tensor, LongTensor
+
+from transformers import AutoTokenizer, AutoModel
+
+
+class WordTransformerEncoder(nn.Module):
+    """
+    Encodes sentences into word-level embeddings using a pretrained MLM transformer.
+    """
+    def __init__(self, model_name: str):
+        super().__init__()
+        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
+        # Model like BERT, RoBERTa, etc.
+        self.model = AutoModel.from_pretrained(model_name)
+
+    def forward(self, words: list[list[str]]) -> Tensor:
+        """
+        Build words embeddings.
+
+        - Tokenizes input sentences into subtokens.
+        - Passes the subtokens through the pre-trained transformer model.
+        - Aggregates subtoken embeddings into word embeddings using mean pooling.
+        """
+        batch_size = len(words)
+
+        # BPE tokenization: split words into subtokens, e.g. ['kidding'] -> ['▁ki', 'dding'].
+        subtokens = self.tokenizer(
+            words,
+            padding=True,
+            truncation=True,
+            is_split_into_words=True,
+            return_tensors='pt'
+        )
+        subtokens = subtokens.to(self.model.device)
+        # Index words from 1 and reserve 0 for special subtokens (e.g. <s>, </s>, padding, etc.).
+        # Such numeration makes a following aggregation easier.
+        words_ids = torch.stack([
+            torch.tensor(
+                [word_id + 1 if word_id is not None else 0 for word_id in subtokens.word_ids(batch_idx)],
+                dtype=torch.long,
+                device=self.model.device
+            )
+            for batch_idx in range(batch_size)
+        ])
+
+        # Run model and extract subtokens embeddings from the last layer.
+        subtokens_embeddings = self.model(**subtokens).last_hidden_state
+
+        # Aggreate subtokens embeddings into words embeddings.
+        # [batch_size, n_words, embedding_size]
+        words_emeddings = self._aggregate_subtokens_embeddings(subtokens_embeddings, words_ids)
+        return words_emeddings
+
+    def _aggregate_subtokens_embeddings(
+        self,
+        subtokens_embeddings: Tensor, # [batch_size, n_subtokens, embedding_size]
+        words_ids: LongTensor          # [batch_size, n_subtokens]
+    ) -> Tensor:
+        """
+        Aggregate subtoken embeddings into word embeddings by averaging.
+
+        This method ensures that multiple subtokens corresponding to a single word are combined
+        into a single embedding.
+        """
+        batch_size, n_subtokens, embedding_size = subtokens_embeddings.shape
+        # The number of words in a sentence plus an "auxiliary" word in the beginnig.
+        n_words = torch.max(words_ids) + 1
+
+        words_embeddings = torch.zeros(
+            size=(batch_size, n_words, embedding_size),
+            dtype=subtokens_embeddings.dtype,
+            device=self.model.device
+        )
+        words_ids_expanded = words_ids.unsqueeze(-1).expand(batch_size, n_subtokens, embedding_size)
+
+        # Use scatter_reduce_ to average embeddings of subtokens corresponding to the same word.
+        # All the padding and special subtokens will be aggregated into an "auxiliary" first embedding,
+        # namely into words_embeddings[:, 0, :].
+        words_embeddings.scatter_reduce_(
+            dim=1,
+            index=words_ids_expanded,
+            src=subtokens_embeddings,
+            reduce="mean",
+            include_self=False
+        )
+        # Now remove the auxiliary word in the beginning.
+        words_embeddings = words_embeddings[:, 1:, :]
+        return words_embeddings
+
+    def get_embedding_size(self) -> int:
+        """Returns the embedding size of the transformer model, e.g. 768 for BERT."""
+        return self.model.config.hidden_size
+
+    def get_embeddings_layer(self):
+        """Returns the embeddings model."""
+        return self.model.embeddings
+        
+    def get_transformer_layers(self) -> list[nn.Module]:
+        """
+        Return a flat list of all transformer-*block* layers, excluding embeddings/poolers, etc.
+        """
+        layers = []
+        for sub in self.model.modules():
+            # find all ModuleLists (these always hold the actual block layers)
+            if isinstance(sub, nn.ModuleList):
+                layers.extend(list(sub))
+        return layers

mlp_classifier.py

@@ -0,0 +1,46 @@
+import torch
+from torch import nn
+from torch import Tensor, LongTensor
+
+from transformers.activations import ACT2FN
+
+
+class MlpClassifier(nn.Module):
+    """ Simple feed-forward multilayer perceptron classifier. """
+
+    def __init__(
+        self,
+        input_size: int,
+        hidden_size: int,
+        n_classes: int,
+        activation: str,
+        dropout: float,
+        class_weights: list[float] = None,
+    ):
+        super().__init__()
+
+        self.n_classes = n_classes
+        self.classifier = nn.Sequential(
+            nn.Dropout(dropout),
+            nn.Linear(input_size, hidden_size),
+            ACT2FN[activation],
+            nn.Dropout(dropout),
+            nn.Linear(hidden_size, n_classes)
+        )
+        if class_weights is not None:
+            class_weights = torch.tensor(class_weights, dtype=torch.long)
+        self.cross_entropy = nn.CrossEntropyLoss(weight=class_weights)
+
+    def forward(self, embeddings: Tensor, labels: LongTensor = None) -> dict:
+        logits = self.classifier(embeddings)
+        # Calculate loss.
+        loss = 0.0
+        if labels is not None:
+            # Reshape tensors to match expected dimensions
+            loss = self.cross_entropy(
+                logits.view(-1, self.n_classes),
+                labels.view(-1)
+            )
+        # Predictions.
+        preds = logits.argmax(dim=-1)
+        return {'preds': preds, 'loss': loss}

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:dfb31d300bc938f0353cb18e733a5c611ef8a6bbbbea8766e75878af3c304b06
+size 1147244460

modeling_parser.py

@@ -0,0 +1,190 @@
+from torch import nn
+from torch import LongTensor
+from transformers import PreTrainedModel
+from transformers.modeling_outputs import ModelOutput
+from dataclasses import dataclass
+
+from .configuration import CobaldParserConfig
+from .encoder import WordTransformerEncoder
+from .mlp_classifier import MlpClassifier
+from .dependency_classifier import DependencyClassifier
+from .utils import (
+    build_padding_mask,
+    build_null_mask,
+    prepend_cls,
+    remove_nulls,
+    add_nulls
+)
+
+
+@dataclass
+class CobaldParserOutput(ModelOutput):
+    """
+    Output type for CobaldParser.
+    """
+    loss: float = None
+    words: list = None
+    counting_mask: LongTensor = None
+    lemma_rules: LongTensor = None
+    joint_feats: LongTensor = None
+    deps_ud: LongTensor = None
+    deps_eud: LongTensor = None
+    miscs: LongTensor = None
+    deepslots: LongTensor = None
+    semclasses: LongTensor = None
+
+
+class CobaldParser(PreTrainedModel):
+    """Morpho-Syntax-Semantic Parser."""
+
+    config_class = CobaldParserConfig
+
+    def __init__(self, config: CobaldParserConfig):
+        super().__init__(config)
+
+        self.encoder = WordTransformerEncoder(
+            model_name=config.encoder_model_name
+        )
+        embedding_size = self.encoder.get_embedding_size()
+
+        self.classifiers = nn.ModuleDict()
+        self.classifiers["null"] = MlpClassifier(
+            input_size=self.encoder.get_embedding_size(),
+            hidden_size=config.null_classifier_hidden_size,
+            n_classes=config.consecutive_null_limit + 1,
+            activation=config.activation,
+            dropout=config.dropout
+        )
+        if "lemma_rule" in config.vocabulary:
+            self.classifiers["lemma_rule"] = MlpClassifier(
+                input_size=embedding_size,
+                hidden_size=config.lemma_classifier_hidden_size,
+                n_classes=len(config.vocabulary["lemma_rule"]),
+                activation=config.activation,
+                dropout=config.dropout
+            )
+        if "joint_feats" in config.vocabulary:
+            self.classifiers["joint_feats"] = MlpClassifier(
+                input_size=embedding_size,
+                hidden_size=config.morphology_classifier_hidden_size,
+                n_classes=len(config.vocabulary["joint_feats"]),
+                activation=config.activation,
+                dropout=config.dropout
+            )
+        if "ud_deprel" in config.vocabulary or "eud_deprel" in config.vocabulary:
+            self.classifiers["syntax"] = DependencyClassifier(
+                input_size=embedding_size,
+                hidden_size=config.dependency_classifier_hidden_size,
+                n_rels_ud=len(config.vocabulary["ud_deprel"]),
+                n_rels_eud=len(config.vocabulary["eud_deprel"]),
+                activation=config.activation,
+                dropout=config.dropout
+            )
+        if "misc" in config.vocabulary:
+            self.classifiers["misc"] = MlpClassifier(
+                input_size=embedding_size,
+                hidden_size=config.misc_classifier_hidden_size,
+                n_classes=len(config.vocabulary["misc"]),
+                activation=config.activation,
+                dropout=config.dropout
+            )
+        if "deepslot" in config.vocabulary:
+            self.classifiers["deepslot"] = MlpClassifier(
+                input_size=embedding_size,
+                hidden_size=config.deepslot_classifier_hidden_size,
+                n_classes=len(config.vocabulary["deepslot"]),
+                activation=config.activation,
+                dropout=config.dropout
+            )
+        if "semclass" in config.vocabulary:
+            self.classifiers["semclass"] = MlpClassifier(
+                input_size=embedding_size,
+                hidden_size=config.semclass_classifier_hidden_size,
+                n_classes=len(config.vocabulary["semclass"]),
+                activation=config.activation,
+                dropout=config.dropout
+            )
+
+    def forward(
+        self,
+        words: list[list[str]],
+        counting_masks: LongTensor = None,
+        lemma_rules: LongTensor = None,
+        joint_feats: LongTensor = None,
+        deps_ud: LongTensor = None,
+        deps_eud: LongTensor = None,
+        miscs: LongTensor = None,
+        deepslots: LongTensor = None,
+        semclasses: LongTensor = None,
+        sent_ids: list[str] = None,
+        texts: list[str] = None,
+        inference_mode: bool = False
+    ) -> CobaldParserOutput:
+        result = {}
+
+        # Extra [CLS] token accounts for the case when #NULL is the first token in a sentence.
+        words_with_cls = prepend_cls(words)
+        words_without_nulls = remove_nulls(words_with_cls)
+        # Embeddings of words without nulls.
+        embeddings_without_nulls = self.encoder(words_without_nulls)
+        # Predict nulls.
+        null_output = self.classifiers["null"](embeddings_without_nulls, counting_masks)
+        result["counting_mask"] = null_output['preds']
+        result["loss"] = null_output["loss"]
+
+        # "Teacher forcing": during training, pass the original words (with gold nulls)
+        # to the classification heads, so that they are trained upon correct sentences.
+        if inference_mode:
+            # Restore predicted nulls in the original sentences.
+            result["words"] = add_nulls(words, null_output["preds"])
+        else:
+            result["words"] = words
+
+        # Encode words with nulls.
+        # [batch_size, seq_len, embedding_size]
+        embeddings = self.encoder(result["words"])
+
+        # Predict lemmas and morphological features.
+        if "lemma_rule" in self.classifiers:
+            lemma_output = self.classifiers["lemma_rule"](embeddings, lemma_rules)
+            result["lemma_rules"] = lemma_output['preds']
+            result["loss"] += lemma_output['loss']
+
+        if "joint_feats" in self.classifiers:
+            joint_feats_output = self.classifiers["joint_feats"](embeddings, joint_feats)
+            result["joint_feats"] = joint_feats_output['preds']
+            result["loss"] += joint_feats_output['loss']
+
+        # Predict syntax.
+        if "syntax" in self.classifiers:
+            padding_mask = build_padding_mask(result["words"], self.device)
+            null_mask = build_null_mask(result["words"], self.device)
+            deps_output = self.classifiers["syntax"](
+                embeddings,
+                deps_ud,
+                deps_eud,
+                mask_ud=(padding_mask & ~null_mask),
+                mask_eud=padding_mask
+            )
+            result["deps_ud"] = deps_output['preds_ud']
+            result["deps_eud"] = deps_output['preds_eud']
+            result["loss"] += deps_output['loss_ud'] + deps_output['loss_eud']
+
+        # Predict miscellaneous features.
+        if "misc" in self.classifiers:
+            misc_output = self.classifiers["misc"](embeddings, miscs)
+            result["miscs"] = misc_output['preds']
+            result["loss"] += misc_output['loss']
+
+        # Predict semantics.
+        if "deepslot" in self.classifiers:
+            deepslot_output = self.classifiers["deepslot"](embeddings, deepslots)
+            result["deepslots"] = deepslot_output['preds']
+            result["loss"] += deepslot_output['loss']
+
+        if "semclass" in self.classifiers:
+            semclass_output = self.classifiers["semclass"](embeddings, semclasses)
+            result["semclasses"] = semclass_output['preds']
+            result["loss"] += semclass_output['loss']
+
+        return CobaldParserOutput(**result)

training_args.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:fb86bfc1859cb8830ba649eb0e08e3be821f421c3c9c251c0e1fe160e95afe74
+size 5432

utils.py

@@ -0,0 +1,66 @@
+import torch
+from torch import Tensor
+
+
+def pad_sequences(sequences: list[Tensor], padding_value: int) -> Tensor:
+    """
+    Stack 1d tensors (sequences) into a single 2d tensor so that each sequence is padded on the
+    right.
+    """
+    return torch.nn.utils.rnn.pad_sequence(sequences, padding_value=padding_value, batch_first=True)
+
+
+def _build_condition_mask(sentences: list[list[str]], condition_fn: callable, device) -> Tensor:
+    masks = [
+        torch.tensor([condition_fn(word) for word in sentence], dtype=bool, device=device)
+        for sentence in sentences
+    ]
+    return pad_sequences(masks, padding_value=False)
+
+def build_padding_mask(sentences: list[list[str]], device) -> Tensor:
+    return _build_condition_mask(sentences, condition_fn=lambda word: True, device=device)
+
+def build_null_mask(sentences: list[list[str]], device) -> Tensor:
+    return _build_condition_mask(sentences, condition_fn=lambda word: word == "#NULL", device=device)
+
+
+def pairwise_mask(masks1d: Tensor) -> Tensor:
+    """
+    Calculate an outer product of a mask, i.e. masks2d[:, i, j] = masks1d[:, i] & masks1d[:, j].
+    """
+    return masks1d[:, None, :] & masks1d[:, :, None]
+
+
+# Credits: https://docs.allennlp.org/main/api/nn/util/#replace_masked_values
+def replace_masked_values(tensor: Tensor, mask: Tensor, replace_with: float):
+    """
+    Replace all masked values in tensor with `replace_with`.
+    """
+    assert tensor.dim() == mask.dim(), "tensor.dim() of {tensor.dim()} != mask.dim() of {mask.dim()}"
+    tensor.masked_fill_(~mask, replace_with)
+
+
+def prepend_cls(sentences: list[list[str]]) -> list[list[str]]:
+    """
+    Return a copy of sentences with [CLS] token prepended.
+    """
+    return [["[CLS]", *sentence] for sentence in sentences]
+
+def remove_nulls(sentences: list[list[str]]) -> list[list[str]]:
+    """
+    Return a copy of sentences with nulls removed.
+    """
+    return [[word for word in sentence if word != "#NULL"] for sentence in sentences]
+
+def add_nulls(sentences: list[list[str]], counting_mask) -> list[list[str]]:
+    """
+    Return a copy of sentences with nulls restored according to counting masks.
+    """
+    sentences_with_nulls = []
+    for sentence, counting_mask in zip(sentences, counting_mask):
+        sentence_with_nulls = []
+        for word, n_nulls_to_insert in zip(sentence, counting_mask):
+            sentence_with_nulls.append(word)
+            sentence_with_nulls.extend(["#NULL"] * n_nulls_to_insert)
+        sentences_with_nulls.append(sentence_with_nulls)
+    return sentences_with_nulls