Upload DisamBert

DisamBert.py

@@ -42,12 +42,14 @@ class DisamBert(PreTrainedModel):
         if config.init_basemodel:
             self.BaseModel = AutoModel.from_pretrained(config.name_or_path, device_map="auto")
             self.classifier_head = nn.UninitializedParameter()
+            self.bias = nn.UninitializedParameter()
             self.__entities = None
         else:
             self.BaseModel = ModernBertModel(config)
             self.classifier_head = nn.Parameter(
                 torch.empty((config.vocab_size, config.hidden_size))
             )
+            self.bias = nn.Parameter(torch.empty((config.vocab_size, 1)))
             self.__entities = pd.Series(config.entities)
         config.init_basemodel = False
         self.tokenizer = AutoTokenizer.from_pretrained(config.tokenizer_path)
@@ -87,6 +89,11 @@ class DisamBert(PreTrainedModel):
             self.config.entities = entity_ids
             self.config.vocab_size = len(entity_ids)
             self.classifier_head = nn.Parameter(torch.cat(vectors, dim=0))
+            self.bias = nn.Parameter(
+                torch.nn.init.normal_(
+                    torch.empty((self.config.vocab_size, 1)), std=self.classifier_head.std().item()
+                )
+            )
 
     @property
     def entities(self) -> pd.Series:
@@ -125,11 +132,11 @@ class DisamBert(PreTrainedModel):
                 for (i, sentence_indices) in enumerate(lengths)
             ]
         )
-        logits = torch.einsum("ij,kj->ki", span_vectors, self.classifier_head)
+        logits = torch.einsum("ij,kj->ki", span_vectors, self.classifier_head) + self.bias
         logits1 = logits - logits.min()
         mask = torch.zeros_like(logits)
-        for (i,concepts) in enumerate(chain.from_iterable(candidates)):
-            mask[concepts,i] = torch.tensor(1.0)
+        for i, concepts in enumerate(chain.from_iterable(candidates)):
+            mask[concepts, i] = torch.tensor(1.0)
         logits2 = logits1 * mask
         sentence_lengths = [len(sentence_indices) for sentence_indices in lengths]
         maxlen = max(sentence_lengths)
@@ -224,7 +231,7 @@ class DisamBert(PreTrainedModel):
                 "input_ids": padded.input_ids,
                 "attention_mask": padded.attention_mask,
                 "lengths": all_indices,
-                "candidates": [example['candidates'] for example in batch]
+                "candidates": [example["candidates"] for example in batch],
             }
             if "labels" in batch[0]:
                 result["labels"] = self.pad_labels([example["labels"] for example in batch])

model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:50c403c889a37e9ed106f0912eafe6e97fd2e9bffff26a34d9af7b284643657e
-size 957523088
+oid sha256:ff4e9bebae857919d9ca236d04b7bb8aae63f405f9cd624bc7ee5ac59f2bd54f
+size 957993808