| from collections.abc import Generator, Iterable |
| from dataclasses import dataclass |
| from enum import StrEnum |
| from itertools import chain |
|
|
| from nltk.corpus import wordnet |
| from nltk.metrics import edit_distance |
| import numpy as np |
| import pandas as pd |
| import torch |
| import torch.nn as nn |
| from transformers import ( |
| AutoConfig, |
| AutoModel, |
| AutoTokenizer, |
| ModernBertModel, |
| PreTrainedConfig, |
| PreTrainedModel, |
| ) |
| from transformers.modeling_outputs import SequenceClassifierOutput |
|
|
| 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 DisamBertCrossEncoder(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") |
| else: |
| self.BaseModel = ModernBertModel(config) |
| config.init_basemodel = False |
| self.loss = nn.BCEWithLogitsLoss() |
| self.post_init() |
| |
| @classmethod |
| def from_base(cls, base_id: ModelURI): |
| config = AutoConfig.from_pretrained(base_id) |
| config.init_basemodel = True |
| config.tokenizer_path = base_id |
| return cls(config) |
| |
| def forward( |
| self, |
| input_ids: torch.Tensor, |
| attention_mask: torch.Tensor, |
| labels: torch.Tensor | None = None, |
| output_hidden_states: bool = False, |
| output_attentions: bool = False, |
| ) -> SequenceClassifierOutput: |
| 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 |
| prev = -1 |
| rows = [] |
| cols = [] |
| for (i,j) in (input_ids == self.config.sep_token_id).nonzero(): |
| if i!=prev: |
| rows.append(i) |
| cols.append(j) |
| prev=i |
| gloss_vectors = token_vectors[rows,cols] |
| logits = torch.einsum("ij,ij->i",token_vectors[:,0],gloss_vectors) |
| return SequenceClassifierOutput( |
| 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 get_lemma(text: str, synset: wordnet.synset) -> wordnet.lemma: |
| best_score = 1000000 |
| best_lemma = None |
| for lemma in synset.lemmas(): |
| score = edit_distance(text, lemma.name()) |
| if score < best_score: |
| best_score = score |
| best_lemma = lemma |
| return best_lemma |
| |
| class CrossEncoderTagger: |
| def __init__(self,url:str): |
| self.model=AutoModel.from_pretrained(url, |
| device_map="auto", |
| trust_remote_code=True) |
| print(self.model) |
| self.tokenizer=AutoTokenizer.from_pretrained(url) |
| |
| def __call__(self,target:str,sentence:str,candidates:str)->str: |
| text = f"{target}::{sentence}" |
| synsets = [wordnet.synset(candidate) for candidate in candidates] |
| definitions = [f"{get_lemma(target,synset)}::{synset.definition()}" |
| for synset in synsets] |
| sentences = [text]*len(candidates) |
| with self.model.device: |
| tokens = self.tokenizer(sentences,definitions,padding=True,return_tensors="pt") |
| output = self.model(tokens.input_ids, |
| tokens.attention_mask) |
| print(dir(output)) |
| logits = output.logits |
| return candidates[logits.argmax()] |
| |
| |
| |
