Upload DisamBert

DisamBert.py

@@ -5,7 +5,7 @@ from enum import StrEnum
 import pandas as pd
 import torch
 import torch.nn as nn
-from transformers import AutoConfig, AutoModel, AutoTokenizer, PreTrainedModel, PreTrainedConfig
+from transformers import AutoConfig, AutoModel, AutoTokenizer, ModernBertModel, PreTrainedConfig, PreTrainedModel
 
 BATCH_SIZE = 64
 
@@ -25,29 +25,30 @@ class LexicalExample:
 class PaddedBatch:
     input_ids: torch.Tensor
     attention_mask: torch.Tensor
-    
+
 
 class DisamBert(PreTrainedModel):
-    def __init__(self, config:PreTrainedConfig):
+    def __init__(self, config: PreTrainedConfig):
         super().__init__(config)
         if config.init_basemodel:
-            self.BaseModel = AutoModel.from_pretrained(config.name_or_path,device_map="auto")
-            with self.BaseModel.device:
-                self.classifier_head = nn.UninitializedParameter()
+            self.BaseModel = AutoModel.from_pretrained(config.name_or_path, device_map="auto")
+            self.classifier_head = nn.UninitializedParameter()
             self.__entities = None
-            config.init_basemodel = False
+        else:
+            self.BaseModel = ModernBertModel(config)
+            self.classifier_head = nn.Parameter(torch.empty((config.vocab_size,config.hidden_size)))
+            self._entities
+        config.init_basemodel = False
         self.tokenizer = AutoTokenizer.from_pretrained(config.tokenizer_path)
         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 init_classifier(self, entities: Generator[LexicalExample]) -> None:
         entity_ids = []
         vectors = []
@@ -57,25 +58,22 @@ class DisamBert(PreTrainedModel):
             for entity in entities:
                 entity_ids.append(entity.concept)
                 batch.append(entity.definition)
-    
+
                 n += 1
                 if n == BATCH_SIZE:
                     tokens = self.tokenizer(batch, padding=True, return_tensors="pt")
-                    encoding = self.BaseModel(
-                        tokens["input_ids"], tokens["attention_mask"]
-                    )
+                    encoding = self.BaseModel(tokens["input_ids"], tokens["attention_mask"])
                     vectors.append(encoding.last_hidden_state.detach()[:, 0])
                     n = 0
                     batch = []
             if n > 0:
                 tokens = self.tokenizer(batch, padding=True, return_tensors="pt")
-                encoding = self.BaseModel(
-                    tokens["input_ids"], tokens["attention_mask"]
-                )
+                encoding = self.BaseModel(tokens["input_ids"], tokens["attention_mask"])
                 vectors.append(encoding.last_hidden_state.detach()[:, 0])
-    
+
             self.__entities = pd.Series(entity_ids)
             self.config.entities = entity_ids
+            self.config.vocab_size = len(entity_ids)
             self.classifier_head = nn.Parameter(torch.cat(vectors, dim=0))
 
     @property
@@ -147,12 +145,7 @@ class DisamBert(PreTrainedModel):
             ]
         )
         attention_mask = torch.vstack(
-            [
-                torch.cat(
-                    (torch.ones(length), torch.zeros(maxlen - length))
-                )
-                for length in lengths
-            ]
+            [torch.cat((torch.ones(length), torch.zeros(maxlen - length))) for length in lengths]
         )
         return PaddedBatch(input_ids, attention_mask)
 

config.json

@@ -117741,5 +117741,5 @@
   "tie_word_embeddings": true,
   "tokenizer_path": "answerdotai/ModernBERT-base",
   "transformers_version": "5.0.0",
-  "vocab_size": 50368
+  "vocab_size": 117660
 }

model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:883bdf7c5d683a50af02c8e5d68d43000411f13b69a66c20dbb2e348977085e5
+oid sha256:79d0851573b5002b29d196af74a0b87c06e774b30889fe729bd17f323af7fc2f
 size 957523088