Upload model

config.json

@@ -48,7 +48,6 @@
   "num_attention_heads": 12,
   "num_hidden_layers": 12,
   "number_of_layer_per_head": 3,
-  "p_augmentation": 0.5,
   "pad_token_id": 1,
   "percentage_tags": 0.25,
   "position_embedding_type": "absolute",

configuration_multiheadcrf.py

@@ -13,7 +13,6 @@ class MultiHeadCRFConfig(PretrainedConfig):
         augmentation = "random",
         context_size = 64,
         percentage_tags = 0.2,
-        p_augmentation = 0.5,
         aug_prob = 0.5,
         crf_reduction = "mean",
         freeze = False,
@@ -26,7 +25,6 @@ class MultiHeadCRFConfig(PretrainedConfig):
         self.augmentation = augmentation
         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

model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:41aa5c6b4abc89c294564cf0f3ff950a424c0a05ab4423966cab474c96e0326e
-size 531721208
+oid sha256:840cc2a3c88f98c414223f2b99ddf20023d16b8ba654a1052627e926acf3bcb6
+size 531721800

modeling_multiheadcrf.py

@@ -12,9 +12,10 @@ NUM_PER_LAYER = 16
 
 class RobertaMultiHeadCRFModel(PreTrainedModel):
     config_class = MultiHeadCRFConfig
-    transformer_backbone_class = RobertaModel
+    transformers_backbone_name = "roberta"
+    transformers_backbone_class = RobertaModel
     _keys_to_ignore_on_load_unexpected = [r"pooler"]
-
+    
     def __init__(self, config):
         super().__init__(config)
         self.num_labels = config.num_labels
@@ -24,7 +25,10 @@ class RobertaMultiHeadCRFModel(PreTrainedModel):
         self.heads = config.classes #expected an array of classes we are predicting
         
         # this can be BERT ROBERTA and other BERT-variants
-        self.bert = self.transformer_backbone_class(config, add_pooling_layer=False)
+        # THIS IS BC HF needs to have "roberta" for roberta models and "bert" for BERT models as var so tha I can load
+        # check https://github.com/huggingface/transformers/blob/b487096b02307cd6e0f132b676cdcc7255fe8e74/src/transformers/models/roberta/modeling_roberta.py#L1170C16-L1170C20
+        setattr(self, self.transformers_backbone_name, self.transformers_backbone_class(config, add_pooling_layer=False))
+        #self.roberta = self.transformer_backbone_class(config, add_pooling_layer=False)
                     #AutoModel(config, add_pooling_layer=False)
                     #AutoModel.from_pretrained(config._name_or_path, config=config, add_pooling_layer=False)
         self.dropout = nn.Dropout(config.hidden_dropout_prob)
@@ -43,20 +47,23 @@ class RobertaMultiHeadCRFModel(PreTrainedModel):
             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()
+        for ent in  self.heads:
+            for i in range(self.number_of_layer_per_head):
+                getattr(self, f"{ent}_dense_{i}").reset_parameters()
+            getattr(self, f"{ent}_classifier").reset_parameters()
+            getattr(self, f"{ent}_crf").reset_parameters()
+            getattr(self, f"{ent}_crf").mask_impossible_transitions()
     
     def manage_freezing(self):
-        for _, param in self.bert.embeddings.named_parameters():
+        for _, param in getattr(self, self.transformers_backbone_name).embeddings.named_parameters():
             param.requires_grad = False
         
         num_encoders_to_freeze = self.config.num_frozen_encoder
         if num_encoders_to_freeze > 0:
-            for _, param in islice(self.bert.encoder.named_parameters(), num_encoders_to_freeze*NUM_PER_LAYER):
+            for _, param in islice(getattr(self, self.transformers_backbone_name).encoder.named_parameters(), num_encoders_to_freeze*NUM_PER_LAYER):
                 param.requires_grad = False
     
     
@@ -75,7 +82,7 @@ class RobertaMultiHeadCRFModel(PreTrainedModel):
         # Default `model.config.use_return_dict´ is `True´
         return_dict = return_dict if return_dict is not None else self.config.use_return_dict
         
-        outputs = self.bert(input_ids,
+        outputs = getattr(self, self.transformers_backbone_name)(input_ids,
                             attention_mask=attention_mask,
                             token_type_ids=token_type_ids,
                             position_ids=position_ids,
@@ -119,7 +126,8 @@ class RobertaMultiHeadCRFModel(PreTrainedModel):
 
 class BertMultiHeadCRFModel(RobertaMultiHeadCRFModel):
     config_class = MultiHeadCRFConfig
-    transformer_backbone_class = BertModel
+    transformers_backbone_name = "bert"
+    transformers_backbone_class = BertModel
     _keys_to_ignore_on_load_unexpected = [r"pooler"]
 
 # Taken from https://github.com/kmkurn/pytorch-crf/blob/master/torchcrf/__init__.py and fixed got uint8 warning