Upload model

config.json

@@ -5,6 +5,9 @@
   ],
   "args_random_seed": 42,
   "attention_probs_dropout_prob": 0.1,
+  "aug_prob": [
+    0.5
+  ],
   "augmentation": "None",
   "auto_map": {
     "AutoConfig": "configuration_multiheadcrf.MultiHeadCRFConfig",
@@ -53,6 +56,6 @@
   "transformers_version": "4.40.2",
   "type_vocab_size": 1,
   "use_cache": true,
-  "version": "0.1.2",
+  "version": "0.1.3",
   "vocab_size": 50262
 }

configuration_multiheadcrf.py

@@ -14,9 +14,10 @@ class MultiHeadCRFConfig(PretrainedConfig):
         context_size = 64,
         percentage_tags = 0.2,
         p_augmentation = 0.5,
+        aug_prob = 0.5,
         crf_reduction = "mean",
         freeze = False,
-        version="0.1.2",
+        version="0.1.3",
         **kwargs,
     ):
         self.classes = classes
@@ -26,8 +27,10 @@ class MultiHeadCRFConfig(PretrainedConfig):
         self.context_size = context_size
         self.percentage_tags = percentage_tags
         self.p_augmentation = p_augmentation
+        self.aug_prob = aug_prob,
         self.crf_reduction = crf_reduction
         self.freeze=freeze
+        self.version = version
         super().__init__(**kwargs)
         
 

modeling_multiheadcrf.py

@@ -41,7 +41,15 @@ class RobertaMultiHeadCRFModel(PreTrainedModel):
 
         if self.config.freeze == True:
             self.manage_freezing()
-            
+    
+    def training_mode(self):
+        # for some reason these layers are not being correctly init
+        # probably related with the lifecycle of the hf .from_pretrained method
+        self.dense.reset_parameters()
+        self.classifier.reset_parameters()
+        self.crf.reset_parameters()
+        self.crf.mask_impossible_transitions()
+    
     def manage_freezing(self):
         for _, param in self.bert.embeddings.named_parameters():
             param.requires_grad = False