Upload model

config.json

@@ -60,6 +60,6 @@
   "type_vocab_size": 2,
   "update_vocab": 28899,
   "use_cache": true,
-  "version": "0.1.0",
+  "version": "0.1.1",
   "vocab_size": 28899
 }

configuration_bionextextractor.py

@@ -11,13 +11,17 @@ class BioNExtExtractorConfig(PretrainedConfig):
         arch_type = "mha",
         index_type = "both",
         novel = True,
-        version="0.1.0",
+        tokenizer_special_tokens = ['[s1]','[e1]', '[s2]','[e2]' ],
+        update_vocab = None,
+        version="0.1.1",
         **kwargs,
     ):
         self.version = version
         self.arch_type = arch_type
         self.index_type = index_type
         self.novel = novel
+        self.tokenizer_special_tokens = tokenizer_special_tokens
+        self.update_vocab = update_vocab
         super().__init__(**kwargs)
         
 

model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:5614db1403a5339630fef087a18cc985693d4b7188cf6c81aa9f15eff71fe520
-size 1350790260
+oid sha256:32f371a5688163ffd745b58918b63752337769ef7223c9ad3702e5af33d06bd1
+size 1350787852

modeling_bionextextractor.py

@@ -47,9 +47,13 @@ class RelationClassifierBase(PreTrainedModel, RelationLossMixin):
     def __init__(self, config):
         super().__init__(config)
         self.num_labels = config.num_labels
-        
+        self.config = config
         #print(config)
         self.bert = BertModel(config, add_pooling_layer=False)
+        
+    def training_mode(self):
+        if self.config.update_vocab is not None:
+            self.bert.resize_token_embeddings(self.config.update_vocab)
     
     def group_embeddings_by_index(self, embeddings, indexes):
         assert len(embeddings.shape)==3
@@ -126,6 +130,11 @@ class RelationClassifierBiLSTM(RelationClassifierBase):
         self.lstm = torch.nn.LSTM(config.hidden_size, (config.hidden_size) // 2, self.num_lstm_layers, batch_first=True, bidirectional=True)
         self.fc = torch.nn.Linear(config.hidden_size, self.num_labels)  # 2 for bidirection
 
+    def training_mode(self):
+        super().training_mode()
+        self.lstm.reset_parameters()
+        self.fc.reset_parameters()
+        
     def classifier_representation(self, embeddings, mask=None):
         out, _ = self.lstm(embeddings)
         return out[:, -1, :]
@@ -139,6 +148,10 @@ class RelationAndNovelClassifierBiLSTM(RelationClassifierBiLSTM, RelationAndNove
         super().__init__(config)
         self.fc_novel = torch.nn.Linear(config.hidden_size, 2)  # 2 for bidirection
 
+    def training_mode(self):
+        super().training_mode()
+        self.fc_novel.reset_parameters()
+    
     def classifier(self, class_representation):
         return super().classifier(class_representation), self.fc_novel(class_representation)
 
@@ -155,6 +168,13 @@ class RelationClassifierMHAttention(RelationClassifierBase):
         self.fc1_activation = torch.nn.GELU(approximate='none')
         self.fc2 = torch.nn.Linear(config.hidden_size//2, self.num_labels)  # 2 for bidirection
 
+    def training_mode(self):
+        super().training_mode()
+        torch.nn.init.kaiming_uniform_(self.weight, a=math.sqrt(5))
+        self.MHattention_layer._reset_parameters()
+        self.fc1.reset_parameters()
+        self.fc2.reset_parameters()
+    
     def classifier_representation(self, embeddings, mask=None):
         batch_size = embeddings.shape[0]
         weight = self.weight.repeat(batch_size, 1, 1)
@@ -185,6 +205,11 @@ class RelationAndNovelClassifierMHAttention(RelationClassifierMHAttention, Relat
         self.fc1_novel_activation = torch.nn.GELU(approximate='none')
         self.fc2_novel = torch.nn.Linear(config.hidden_size//2, 2)  # 2 for bidirection
 
+    def training_mode(self):
+        super().training_mode()
+        self.fc1_novel.reset_parameters()
+        self.fc2_novel.reset_parameters()
+    
     def classifier(self, class_representation, relation_mask=None):
         x = self.fc1_novel(class_representation)
         x = self.fc1_novel_activation(x)
@@ -196,17 +221,9 @@ ARCH_MAPPING = {"mhawNovelty": RelationAndNovelClassifierMHAttention,
                 "bilstmwNovelty" : RelationAndNovelClassifierBiLSTM,
                 "bilstm": RelationClassifierBiLSTM}
 
-class BioNExtExtractorModel(PreTrainedModel):
+## Changing the name to be compatible with HF API
+
+class BioNExtExtractorModel(RelationAndNovelClassifierMHAttention):
     config_class=BioNExtExtractorConfig
     
-    def __init__(self, config):
-        super().__init__(config)
-        
-        if config.novel:
-            self.model = ARCH_MAPPING[f"{config.arch_type}wNovelty"](config)
-        else:
-            self.model = ARCH_MAPPING[config.arch_type](config)
-        
-    def forward(self, *args, **kwargs):
-        return self.model(*args, **kwargs)