End of training

ConSec.py

@@ -0,0 +1,176 @@
+from collections.abc import Generator, Iterable
+from dataclasses import dataclass
+from enum import StrEnum
+
+from nltk.tokenize import TreebankWordDetokenizer
+
+import torch
+import torch.nn as nn
+from transformers import (
+    AutoConfig,
+    AutoModel,
+    BatchEncoding,
+    DebertaV2Model,
+    PreTrainedConfig,
+    PreTrainedModel,
+    PreTrainedTokenizer,
+)
+from transformers.modeling_outputs import TokenClassifierOutput
+
+class ModelURI(StrEnum):
+    BASE = "microsoft/deberta-v3-base"
+    LARGE = "microsoft/deberta-v3-large"
+
+class ConSec(PreTrainedModel):
+    def __init__(self, config: PreTrainedConfig):
+        super().__init__(config)
+        if config.init_basemodel:
+            self.BaseModel = AutoModel.from_pretrained(config.name_or_path, 
+                                                       device_map="auto",
+                                                       dtype=torch.bfloat16)
+            self.config.vocab_size += 2
+            self.BaseModel.resize_token_embeddings(self.config.vocab_size)
+        else:
+            self.BaseModel = DebertaV2Model(config)
+        config.init_basemodel = False
+
+        self.loss = nn.CrossEntropyLoss()
+        self.post_init()
+
+    @classmethod
+    def from_base(cls, base_id: ModelURI):
+        config = AutoConfig.from_pretrained(base_id)
+        config.init_basemodel = True
+        return cls(config)
+    
+    def add_special_tokens(self, start: int, end: int, gloss: int):
+        self.config.start_token = start
+        self.config.end_token = end
+        self.config.gloss_token = gloss
+    
+    def forward(self,
+                input_ids: torch.Tensor | None = None, 
+                attention_mask: torch.Tensor | None = None, 
+                token_type_ids: torch.Tensor | None = None, 
+                position_ids: torch.Tensor | None = None, 
+                inputs_embeds: torch.Tensor | None = None, 
+                labels: torch.Tensor | None = None, 
+                output_attentions: bool | None = None, 
+                output_hidden_states: bool | None = None, 
+                return_dict: bool | None = None,
+                **kwargs)->TokenClassifierOutput:
+        base_model_output = self.BaseModel(input_ids=input_ids,
+                                           attention_mask=attention_mask,
+                                           token_type_ids=token_type_ids,
+                                           position_ids=position_ids,
+                                           inputs_embeds=inputs_embeds,
+                                           output_attentions=output_attentions,
+                                           output_hidden_states=output_hidden_states,
+                                           **kwargs)
+        token_vectors = base_model_output.last_hidden_state
+        selection = torch.zeros_like(input_ids, dtype=token_vectors.dtype)
+        starts = (input_ids == self.config.start_token).nonzero()
+        ends = (input_ids == self.config.end_token).nonzero()
+        for startpos, endpos in zip(starts, ends, strict=True):
+            selection[startpos[0], startpos[1] : endpos[1] + 1] = 1.0
+        entity_vectors = torch.einsum("ijk,ij->ik", token_vectors, selection)
+        gloss_vectors = self.gloss_vectors(
+            input_ids, starts, position_ids, token_vectors
+        )
+        logits = torch.einsum("ij,ikj->ik", entity_vectors, gloss_vectors)
+
+        return TokenClassifierOutput(
+            logits=logits,
+            loss=self.loss(logits, labels) if labels is not None else None,
+            hidden_states=base_model_output.hidden_states if output_hidden_states else None,
+            attentions=base_model_output.attentions if output_attentions else None,
+        )
+    
+    def gloss_vectors(self,input_ids: torch.Tensor,
+                      starts: torch.Tensor,
+                      position_ids: torch.Tensor,
+                      token_vectors: torch.Tensor)->torch.Tensor:
+        with self.device:
+            vectors = [token_vectors[i,((position_ids[i]==position_ids[i,j])&(input_ids[i]==self.config.gloss_token))]
+                       for (i,j) in starts]
+            maxlen  = max(vector.shape[0] for vector in vectors)
+            return torch.stack([torch.cat([vector,torch.zeros((maxlen-vector.shape[0],vector.shape[1]),
+                                                              dtype=torch.bfloat16)])
+                                for vector in vectors])
+    
+def json_sequencer(sentence:list[dict])->Generator[tuple[list[str], list[str], int]]:
+    for site in sorted([{"span":i,
+                        "n_candidates":len(chunk["candidates"])}
+                        for (i,chunk) in enumerate(sentence)
+                        if "candidates" in chunk],
+                       key = lambda x: x["n_candidates"]):
+        words = [word for chunk in sentence[:site["span"]]
+                 for word in chunk["words"]]
+        words.append("[START]")
+        words.extend(sentence[site["span"]]["words"])
+        words.append("[END]")
+        words.extend([word for chunk in sentence[site["span"]+1:]
+                      for word in chunk["words"]])
+        yield (words,
+               sentence[site["span"]]["candidates"],
+               site["span"])
+        
+def json_labeller(sentence,tags):
+    for tag in tags:
+        sentence[tag["index"]]["label"]=tag["label"]
+    return sentence
+    
+class ConSecTagger:
+    def __init__(self,model,
+                 tokenizer,
+                 ontology,
+                 sequencer=json_sequencer,
+                 labeller=json_labeller):
+        self.model = model
+        self.tokenizer = tokenizer
+        special_tokens = self.tokenizer.get_added_tokens()
+        self.start_token = special_tokens["[START]"]
+        self.gloss_token = special_tokens["[GLOSS]"]
+        self.sequencer = sequencer
+        self.detokenizer = TreebankWordDetokenizer()
+        self.glosses = {synset.concept:synset.definition
+                        for synset in ontology}
+        self.label=labeller
+        
+        
+    def __call__(self,sentence):
+        already_tagged = []
+        for (words,candidates,index) in self.sequencer(sentence):
+            text = self.detokenizer.detokenize(words)
+            glosses = ['']
+            glosses.extend([self.glosses[candidate] for candidate in candidates])
+            glosses.extend([self.glosses[previous["label"]] for previous in already_tagged])
+            with self.model.device:
+                tokens = self.tokenizer(text,"[GLOSS] ".join(glosses),
+                                        return_tensors="pt")
+                length = tokens.input_ids.shape[0]
+                positions = torch.arange(length)
+                place = (tokens.input_ids==self.start_token).nonzero(as_tuple=True)[0]
+                wordpos = tokens.token_to_word(place+1)
+                gloss_positions = list((tokens.input_ids==self.gloss_token).nonzero(as_tuple=True))
+                gloss_positions.append(length)
+                n_candidates = len(candidates)
+                for (i,position) in enumerate(gloss_positions[:-1]):
+                    if i<n_candidates:
+                        end = place + gloss_positions[i+1]-position
+                        positions [position:gloss_positions[i+1]] = torch.arange(place,end)
+                    else:
+                        known = already_tagged[i-n_candidates]
+                        start = tokens.word_to_tokens(known["place"]).start
+                        end = start + gloss_positions[i+1] - position
+                        positions[position:gloss_positions[i+1]] = torch.arange(start,end)
+                prediction = self.model(token_ids=tokens.token_ids,
+                                    attention_mask=tokens.attention_mask,
+                                    token_type_ids=tokens.token_type_ids,
+                                    position_ids=positions)
+                label = candidates[prediction.logits.argmax()]
+                already_tagged.append({"label":label,
+                                       "place":wordpos,
+                                       "index":index})
+        return(self.label(sentence,already_tagged))
+                

DisamBertSingleSense.py

@@ -0,0 +1,172 @@
+from collections.abc import Generator, Iterable
+from dataclasses import dataclass
+from enum import StrEnum
+
+import pprint
+
+import torch
+import torch.nn as nn
+from transformers import (
+    AutoConfig,
+    AutoModel,
+    BatchEncoding,
+    ModernBertModel,
+    PreTrainedConfig,
+    PreTrainedModel,
+    PreTrainedTokenizer,
+)
+from transformers.modeling_outputs import TokenClassifierOutput
+
+BATCH_SIZE = 16
+
+
+class ModelURI(StrEnum):
+    BASE = "answerdotai/ModernBERT-base"
+    LARGE = "answerdotai/ModernBERT-large"
+
+
+@dataclass(slots=True, frozen=True)
+class LexicalExample:
+    concept: str
+    definition: str
+
+
+@dataclass(slots=True, frozen=True)
+class PaddedBatch:
+    input_ids: torch.Tensor
+    attention_mask: torch.Tensor
+
+
+class DisamBertSingleSense(PreTrainedModel):
+    def __init__(self, config: PreTrainedConfig):
+        super().__init__(config)
+        if config.init_basemodel:
+            self.BaseModel = AutoModel.from_pretrained(config.name_or_path, 
+                                                       device_map="auto")
+            self.config.vocab_size += 2
+            self.BaseModel.resize_token_embeddings(self.config.vocab_size)
+        else:
+            self.BaseModel = ModernBertModel(config)
+        config.init_basemodel = False
+
+        self.loss = nn.CrossEntropyLoss()
+        self.post_init()
+
+    @classmethod
+    def from_base(cls, base_id: ModelURI):
+        config = AutoConfig.from_pretrained(base_id)
+        config.init_basemodel = True
+        return cls(config)
+
+    def add_special_tokens(self, start: int, end: int):
+        self.config.start_token = start
+        self.config.end_token = end
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        attention_mask: torch.Tensor,
+        candidate_tokens: torch.Tensor,
+        candidate_attention_masks: torch.Tensor,
+        candidate_mapping: torch.Tensor,
+        labels: Iterable[int] | None = None,
+        output_hidden_states: bool = False,
+        output_attentions: bool = False,
+    ) -> TokenClassifierOutput:
+        base_model_output = self.BaseModel(
+            input_ids,
+            attention_mask,
+            output_hidden_states=output_hidden_states,
+            output_attentions=output_attentions,
+        )
+        token_vectors = base_model_output.last_hidden_state
+        selection = torch.zeros_like(input_ids, dtype=token_vectors.dtype)
+        starts = (input_ids == self.config.start_token).nonzero()
+        ends = (input_ids == self.config.end_token).nonzero()
+        for startpos, endpos in zip(starts, ends, strict=True):
+            selection[startpos[0], startpos[1] : endpos[1] + 1] = 1.0
+        entity_vectors = torch.einsum("ijk,ij->ik", token_vectors, selection)
+        gloss_vectors = self.gloss_vectors(
+            candidate_tokens, candidate_attention_masks, candidate_mapping
+        )
+        logits = torch.einsum("ij,ikj->ik", entity_vectors, gloss_vectors)
+
+        return TokenClassifierOutput(
+            logits=logits,
+            loss=self.loss(logits, labels) if labels is not None else None,
+            hidden_states=base_model_output.hidden_states if output_hidden_states else None,
+            attentions=base_model_output.attentions if output_attentions else None,
+        )
+
+    def gloss_vectors(self, candidates, candidate_attention_masks, candidate_mapping):
+        with self.device:
+            vectors = self.BaseModel(candidates, candidate_attention_masks).last_hidden_state[:, 0]
+            chunks = [
+                torch.squeeze(vectors[(candidate_mapping == sentence_index).nonzero()], dim=1)
+                for sentence_index in torch.unique(candidate_mapping)
+            ]
+            maxlen = max(chunk.shape[0] for chunk in chunks)
+            return torch.stack(
+                [
+                    torch.cat(
+                        [chunk, torch.zeros((maxlen - chunk.shape[0], self.config.hidden_size))]
+                    )
+                    for chunk in chunks
+                ]
+            )
+
+
+class CandidateLabeller:
+    def __init__(
+        self,
+        tokenizer: PreTrainedTokenizer,
+        ontology: Generator[LexicalExample],
+        device: torch.device,
+        retain_candidates: bool = False,
+    ):
+        self.tokenizer = tokenizer
+        self.device = device
+        self.gloss_tokens = {
+            example.concept: self.tokenizer(example.definition, padding=True)
+            for example in ontology
+        }
+        self.retain_candidates = retain_candidates
+
+    def __call__(self, batch: dict) -> dict:
+        with self.device:
+            encoded = [
+                BatchEncoding(
+                    {"input_ids": example["input_ids"], "attention_mask": example["attention_mask"]}
+                )
+                for example in batch
+            ]
+            tokens = self.tokenizer.pad(encoded, padding=True, return_tensors="pt")
+            candidate_tokens = self.tokenizer.pad(
+                [
+                    self.gloss_tokens[concept]
+                    for example in batch
+                    for concept in example["candidates"]
+                ],
+                padding=True,
+                return_attention_mask=True,
+                return_tensors="pt",
+            )
+            result = {
+                "input_ids": tokens.input_ids,
+                "attention_mask": tokens.attention_mask,
+                "candidate_tokens": candidate_tokens.input_ids,
+                "candidate_attention_masks": candidate_tokens.attention_mask,
+                "candidate_mapping": torch.cat(
+                    [
+                        torch.tensor([i] * len(example["candidates"]))
+                        for (i, example) in enumerate(batch)
+                    ]
+                ),
+            }
+            if "label" in batch[0]:
+                result["labels"] = torch.tensor(
+                    [example["candidates"].index(example["label"]) for example in batch]
+                )
+            if self.retain_candidates:
+                result["candidates"] = [example["candidates"] for example in batch]
+            return result

README.md

@@ -0,0 +1,74 @@
+---
+library_name: transformers
+language:
+- en
+license: mit
+base_model: microsoft/deberta-v3-base
+tags:
+- generated_from_trainer
+metrics:
+- precision
+- recall
+- f1
+model-index:
+- name: ConSec
+  results: []
+---
+
+<!-- This model card has been generated automatically according to the information the Trainer had access to. You
+should probably proofread and complete it, then remove this comment. -->
+
+# ConSec
+
+This model is a fine-tuned version of [microsoft/deberta-v3-base](https://huggingface.co/microsoft/deberta-v3-base) on the None dataset.
+It achieves the following results on the evaluation set:
+- Loss: 3.9109
+- Precision: 0.5353
+- Recall: 0.5517
+- F1: 0.5434
+- Matthews: 0.5509
+
+## Model description
+
+More information needed
+
+## Intended uses & limitations
+
+More information needed
+
+## Training and evaluation data
+
+More information needed
+
+## Training procedure
+
+### Training hyperparameters
+
+The following hyperparameters were used during training:
+- learning_rate: 0.0001
+- train_batch_size: 4
+- eval_batch_size: 4
+- seed: 42
+- optimizer: Use OptimizerNames.ADAMW_TORCH_FUSED with betas=(0.9,0.999) and epsilon=1e-08 and optimizer_args=No additional optimizer arguments
+- lr_scheduler_type: inverse_sqrt
+- lr_scheduler_warmup_steps: 1000
+- num_epochs: 5
+
+### Training results
+
+| Training Loss | Epoch | Step   | Validation Loss | Precision | Recall | F1     | Matthews |
+|:-------------:|:-----:|:------:|:---------------:|:---------:|:------:|:------:|:--------:|
+| No log        | 0     | 0      | 365.7523        | 0.4494    | 0.3098 | 0.3668 | 0.3092   |
+| 4.1345        | 1.0   | 56179  | 3.8725          | 0.5383    | 0.5535 | 0.5458 | 0.5527   |
+| 4.1174        | 2.0   | 112358 | 3.9544          | 0.5360    | 0.5517 | 0.5437 | 0.5509   |
+| 3.9667        | 3.0   | 168537 | 3.9244          | 0.5348    | 0.5517 | 0.5431 | 0.5509   |
+| 4.2556        | 4.0   | 224716 | 3.9431          | 0.5348    | 0.5508 | 0.5427 | 0.5501   |
+| 4.2113        | 5.0   | 280895 | 3.9109          | 0.5353    | 0.5517 | 0.5434 | 0.5509   |
+
+
+### Framework versions
+
+- Transformers 5.3.0
+- Pytorch 2.10.0+cu128
+- Datasets 4.5.0
+- Tokenizers 0.22.2

config.json

@@ -0,0 +1,46 @@
+{
+  "architectures": [
+    "ConSec"
+  ],
+  "attention_probs_dropout_prob": 0.1,
+  "auto_map": {
+    "AutoModel": "ConSec.ConSec"
+  },
+  "bos_token_id": 1,
+  "dtype": "bfloat16",
+  "end_token": 128002,
+  "eos_token_id": 2,
+  "gloss_token": 128003,
+  "hidden_act": "gelu",
+  "hidden_dropout_prob": 0.1,
+  "hidden_size": 768,
+  "init_basemodel": false,
+  "initializer_range": 0.02,
+  "intermediate_size": 3072,
+  "layer_norm_eps": 1e-07,
+  "legacy": true,
+  "max_position_embeddings": 512,
+  "max_relative_positions": -1,
+  "model_type": "deberta-v2",
+  "norm_rel_ebd": "layer_norm",
+  "num_attention_heads": 12,
+  "num_hidden_layers": 12,
+  "pad_token_id": 0,
+  "pooler_dropout": 0,
+  "pooler_hidden_act": "gelu",
+  "pooler_hidden_size": 768,
+  "pos_att_type": [
+    "p2c",
+    "c2p"
+  ],
+  "position_biased_input": false,
+  "position_buckets": 256,
+  "relative_attention": true,
+  "share_att_key": true,
+  "start_token": 128001,
+  "tie_word_embeddings": true,
+  "transformers_version": "5.3.0",
+  "type_vocab_size": 0,
+  "use_cache": false,
+  "vocab_size": 128102
+}

model.safetensors

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

tokenizer_config.json

@@ -0,0 +1,23 @@
+{
+  "add_prefix_space": true,
+  "backend": "tokenizers",
+  "bos_token": "[CLS]",
+  "cls_token": "[CLS]",
+  "do_lower_case": false,
+  "eos_token": "[SEP]",
+  "extra_special_tokens": [
+    "[START]",
+    "[END]",
+    "[GLOSS]"
+  ],
+  "is_local": false,
+  "mask_token": "[MASK]",
+  "model_max_length": 1000000000000000019884624838656,
+  "pad_token": "[PAD]",
+  "sep_token": "[SEP]",
+  "split_by_punct": false,
+  "tokenizer_class": "DebertaV2Tokenizer",
+  "unk_id": 3,
+  "unk_token": "[UNK]",
+  "vocab_type": "spm"
+}

training_args.bin

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