| 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_vocab() |
| 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[1] |
| positions = torch.arange(length) |
| place = (tokens.input_ids==self.start_token).nonzero(as_tuple=True)[1].item() |
| wordpos = tokens.token_to_word(place) |
| gloss_positions = [index.item() |
| for index in (tokens.input_ids==self.gloss_token).nonzero(as_tuple=True)[1]] |
| 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(input_ids=tokens.input_ids, |
| attention_mask=tokens.attention_mask, |
| token_type_ids=tokens.token_type_ids, |
| position_ids=positions.reshape((1,length))) |
| try: |
| label = candidates[prediction.logits.argmax()] |
| except IndexError: |
| print(text) |
| print(gloss_positions) |
| print([positions[pos].item() for pos in gloss_positions[:-1]]) |
| print(already_tagged) |
| print(candidates) |
| print(prediction.logits) |
| print(prediction.logits.argmax()) |
| raise |
| already_tagged.append({"label":label, |
| "place":wordpos, |
| "index":index}) |
| return(self.label(sentence,already_tagged)) |
| |
