Upload custom_modeling.py with huggingface_hub

custom_modeling.py

@@ -1,4 +1,4 @@
-from transformers import PreTrainedModel, AutoModel, AutoConfig
+from transformers import AutoModelForTokenClassification, AutoModel, AutoConfig
 from transformers.modeling_outputs import TokenClassifierOutput
 import torch
 import torch.nn as nn
@@ -7,33 +7,19 @@ from typing import Optional, Union, Tuple, List
 import os
 import json
 
-class BertCRFPreTrainedModel(PreTrainedModel):
-    """An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained models."""
-    config_class = AutoConfig
-    base_model_prefix = "bert"
-
-    def _init_weights(self, module):
-        """Initialize the weights"""
-        if isinstance(module, nn.Linear):
-            module.weight.data.normal_(mean=0.0, std=self.config.initializer_range if hasattr(self.config, 'initializer_range') else 0.02)
-            if module.bias is not None:
-                module.bias.data.zero_()
-        elif isinstance(module, nn.Embedding):
-            module.weight.data.normal_(mean=0.0, std=self.config.initializer_range if hasattr(self.config, 'initializer_range') else 0.02)
-            if module.padding_idx is not None:
-                module.weight.data[module.padding_idx].zero_()
-        elif isinstance(module, nn.LayerNorm):
-            module.bias.data.zero_()
-            module.weight.data.fill_(1.0)
-
-class BertCRFForTokenClassification(BertCRFPreTrainedModel):
+
+class TransformerCRFForTokenClassification(AutoModelForTokenClassification):
     def __init__(self, config):
         super().__init__(config)
         self.num_labels = config.num_labels
-        self.bert = AutoModel.from_config(config)
+
+        self.base_model = AutoModel.from_config(config=config, use_safetensors=True)
+        hidden_size = config.hidden_size if hasattr(config, 'hidden_size') else 768
+
         self.dropout = nn.Dropout(config.hidden_dropout_prob if hasattr(config, 'hidden_dropout_prob') else 0.1)
-        self.classifier = nn.Linear(config.hidden_size, config.num_labels)
+        self.classifier = nn.Linear(hidden_size, config.num_labels)
         
+
         self.use_crf = config.use_crf if hasattr(config, 'use_crf') else False
         if self.use_crf:
             self.crf = CRF(num_tags=self.num_labels, batch_first=True)
@@ -43,28 +29,27 @@ class BertCRFForTokenClassification(BertCRFPreTrainedModel):
 
         # Initialize weights and apply final processing
         self.post_init()
+        
 
     def forward(
-        self,
-        input_ids: Optional[torch.Tensor] = None,
-        attention_mask: Optional[torch.Tensor] = None,
-        token_type_ids: Optional[torch.Tensor] = None,
-        position_ids: Optional[torch.Tensor] = None,
-        head_mask: Optional[torch.Tensor] = None,
-        inputs_embeds: Optional[torch.Tensor] = None,
-        labels: Optional[torch.Tensor] = None,
-        output_attentions: Optional[bool] = None,
-        output_hidden_states: Optional[bool] = None,
-        return_dict: Optional[bool] = None,
+            self,
+            input_ids: Optional[torch.Tensor] = None,
+            attention_mask: Optional[torch.Tensor] = None,
+            token_type_ids: Optional[torch.Tensor] = None,
+            position_ids: Optional[torch.Tensor] = None,
+            head_mask: Optional[torch.Tensor] = None,
+            inputs_embeds: Optional[torch.Tensor] = None,
+            labels: Optional[torch.Tensor] = None,
+            output_attentions: Optional[bool] = None,
+            output_hidden_states: Optional[bool] = None,
+            return_dict: Optional[bool] = None,
     ) -> Union[Tuple[torch.Tensor], TokenClassifierOutput]:
         return_dict = return_dict if return_dict is not None else self.config.use_return_dict
 
-        outputs = self.bert(
+        outputs = self.base_model(
             input_ids,
             attention_mask=attention_mask,
-            token_type_ids=token_type_ids,
             position_ids=position_ids,
-            head_mask=head_mask,
             inputs_embeds=inputs_embeds,
             output_attentions=output_attentions,
             output_hidden_states=output_hidden_states,
@@ -74,7 +59,7 @@ class BertCRFForTokenClassification(BertCRFPreTrainedModel):
         sequence_output = outputs[0]
         sequence_output = self.dropout(sequence_output)
         logits = self.classifier(sequence_output)
-
+        
         loss = None
         if labels is not None:
             if self.crf is not None:
@@ -100,11 +85,11 @@ class BertCRFForTokenClassification(BertCRFPreTrainedModel):
         """Save model with custom CRF layer"""
         # Save the config
         self.config.use_crf = self.use_crf
-        self.config.save_pretrained(save_directory)
-        
+        self.config.save_pretrained(save_directory, safe_serialization=True)
+
         # Save the model weights
-        super().save_pretrained(save_directory, **kwargs)
-        
+        super().save_pretrained(save_directory, safe_serialization=True, **kwargs)
+
         if self.crf is not None:
             crf_path = os.path.join(save_directory, "crf.pt")
             torch.save(self.crf.state_dict(), crf_path)
@@ -112,11 +97,25 @@ class BertCRFForTokenClassification(BertCRFPreTrainedModel):
     @classmethod
     def from_pretrained(cls, pretrained_model_name_or_path: str, *model_args, **kwargs):
         """Load model with custom CRF layer"""
-        model = super().from_pretrained(pretrained_model_name_or_path, *model_args, **kwargs)
-        
-        # Load CRF if it exists
-        crf_path = os.path.join(pretrained_model_name_or_path, "crf.pt")
-        if os.path.exists(crf_path) and model.use_crf:
-            model.crf.load_state_dict(torch.load(crf_path))
+        if 'config' in kwargs:
+            config = kwargs.pop('config')
+        else:
+            config = AutoConfig.from_pretrained(pretrained_model_name_or_path)
         
-        return model 
+        # Ensure use_crf is set in the configuration
+        if not hasattr(config, 'use_crf'):
+            config.use_crf = False  # or True, depending on your default
+
+        # Load the model
+        model = super().from_pretrained(pretrained_model_name_or_path, config=config, use_safetensors=True, *model_args, **kwargs)
+
+        # Initialize CRF if needed
+        if config.use_crf:
+            model.crf = CRF(num_tags=config.num_labels, batch_first=True)
+            crf_path = os.path.join(pretrained_model_name_or_path, "crf.pt")
+            if os.path.exists(crf_path):
+                model.crf.load_state_dict(torch.load(crf_path))
+        else:
+            model.crf = None
+
+        return model