Update modeling_emotion_classifier.py

modeling_emotion_classifier.py

@@ -1,20 +1,23 @@
-from transformers import PreTrainedModel, HubertModel
+from transformers import PreTrainedModel, HubertConfig, HubertModel
 import torch.nn as nn
 import torch
 from .configuration_emotion_classifier import EmotionClassifierConfig
 
-
 class EmotionClassifierHuBERT(PreTrainedModel):
     config_class = EmotionClassifierConfig
 
     def __init__(self, config):
         super().__init__(config)
-        self.hubert = HubertModel.from_pretrained("facebook/hubert-large-ls960-ft")
+        
+        # Initialize HuBERT without pre-trained weights
+        hubert_config = HubertConfig.from_pretrained("facebook/hubert-large-ls960-ft")
+        self.hubert = HubertModel(hubert_config)
+        
         self.conv1 = nn.Conv1d(in_channels=1024, out_channels=512, kernel_size=3, padding=1)
         self.conv2 = nn.Conv1d(in_channels=512, out_channels=256, kernel_size=3, padding=1)
         self.transformer_encoder = nn.TransformerEncoderLayer(d_model=256, nhead=8)
         self.bilstm = nn.LSTM(input_size=256, hidden_size=config.hidden_size_lstm, num_layers=2, batch_first=True, bidirectional=True)
-        self.fc = nn.Linear(config.hidden_size_lstm * 2, config.num_classes)  # * 2 for bidirectional
+        self.fc = nn.Linear(config.hidden_size_lstm * 2, config.num_classes)
 
     def forward(self, x):
         with torch.no_grad():
@@ -27,3 +30,130 @@ class EmotionClassifierHuBERT(PreTrainedModel):
         x, _ = self.bilstm(x)
         x = self.fc(x[:, -1, :])
         return x
+
+    @classmethod
+    def from_pretrained(cls, pretrained_model_name_or_path, *model_args, **kwargs):
+        config = kwargs.pop("config", None)
+        state_dict = kwargs.pop("state_dict", None)
+        cache_dir = kwargs.pop("cache_dir", None)
+        from_tf = kwargs.pop("from_tf", False)
+        force_download = kwargs.pop("force_download", False)
+        resume_download = kwargs.pop("resume_download", False)
+        proxies = kwargs.pop("proxies", None)
+        output_loading_info = kwargs.pop("output_loading_info", False)
+        local_files_only = kwargs.pop("local_files_only", False)
+        use_auth_token = kwargs.pop("use_auth_token", None)
+        revision = kwargs.pop("revision", None)
+        mirror = kwargs.pop("mirror", None)
+
+        # Load config if we don't provide a configuration
+        if not isinstance(config, PretrainedConfig):
+            config_path = config if config is not None else pretrained_model_name_or_path
+            config, model_kwargs = cls.config_class.from_pretrained(
+                config_path,
+                *model_args,
+                cache_dir=cache_dir,
+                return_unused_kwargs=True,
+                force_download=force_download,
+                resume_download=resume_download,
+                proxies=proxies,
+                local_files_only=local_files_only,
+                use_auth_token=use_auth_token,
+                revision=revision,
+                **kwargs,
+            )
+        else:
+            model_kwargs = kwargs
+
+        # Load model
+        if pretrained_model_name_or_path is not None:
+            pretrained_model_name_or_path = str(pretrained_model_name_or_path)
+            if os.path.isdir(pretrained_model_name_or_path):
+                if from_tf and os.path.isfile(os.path.join(pretrained_model_name_or_path, TF_WEIGHTS_NAME + ".index")):
+                    # Load from a TF 1.0 checkpoint in priority if from_tf
+                    archive_file = os.path.join(pretrained_model_name_or_path, TF_WEIGHTS_NAME + ".index")
+                elif from_tf and os.path.isfile(os.path.join(pretrained_model_name_or_path, TF2_WEIGHTS_NAME)):
+                    # Load from a TF 2.0 checkpoint in priority if from_tf
+                    archive_file = os.path.join(pretrained_model_name_or_path, TF2_WEIGHTS_NAME)
+                elif os.path.isfile(os.path.join(pretrained_model_name_or_path, WEIGHTS_NAME)):
+                    # Load from a PyTorch checkpoint
+                    archive_file = os.path.join(pretrained_model_name_or_path, WEIGHTS_NAME)
+                else:
+                    raise EnvironmentError(
+                        f"Error no file named {[WEIGHTS_NAME, TF2_WEIGHTS_NAME, TF_WEIGHTS_NAME + '.index']} found in "
+                        f"directory {pretrained_model_name_or_path} or '{pretrained_model_name_or_path}' is not a directory."
+                    )
+            elif os.path.isfile(pretrained_model_name_or_path) or is_remote_url(pretrained_model_name_or_path):
+                archive_file = pretrained_model_name_or_path
+            else:
+                # Load from URL or cache
+                archive_file = hf_bucket_url(
+                    pretrained_model_name_or_path,
+                    filename=WEIGHTS_NAME,
+                    revision=revision,
+                    mirror=mirror,
+                )
+
+            try:
+                # Load from URL or cache
+                resolved_archive_file = cached_path(
+                    archive_file,
+                    cache_dir=cache_dir,
+                    force_download=force_download,
+                    proxies=proxies,
+                    resume_download=resume_download,
+                    local_files_only=local_files_only,
+                    use_auth_token=use_auth_token,
+                )
+            except EnvironmentError as err:
+                logger.error(err)
+                msg = (
+                    f"Can't load weights for '{pretrained_model_name_or_path}'. Make sure that:\n\n"
+                    f"- '{pretrained_model_name_or_path}' is a correct model identifier listed on 'https://huggingface.co/models'\n\n"
+                    f"- or '{pretrained_model_name_or_path}' is the correct path to a directory containing a file named one of {WEIGHTS_NAME}, {TF2_WEIGHTS_NAME}, {TF_WEIGHTS_NAME}.\n\n"
+                )
+                raise EnvironmentError(msg)
+
+            if resolved_archive_file == archive_file:
+                logger.info(f"loading weights file {archive_file}")
+            else:
+                logger.info(f"loading weights file {archive_file} from cache at {resolved_archive_file}")
+        else:
+            resolved_archive_file = None
+
+        # Initialize the model
+        model = cls(config)
+
+        if state_dict is None:
+            try:
+                state_dict = torch.load(resolved_archive_file, map_location="cpu")
+            except Exception:
+                raise OSError(
+                    f"Unable to load weights from pytorch checkpoint file for '{pretrained_model_name_or_path}' "
+                    f"at '{resolved_archive_file}'"
+                )
+
+        # Remove the prefix 'module' from the keys if present (happens when using DataParallel)
+        state_dict = {k.replace("module.", ""): v for k, v in state_dict.items()}
+
+        # Load only the custom model weights, excluding HuBERT
+        custom_state_dict = {k: v for k, v in state_dict.items() if not k.startswith('hubert.')}
+        missing_keys, unexpected_keys = model.load_state_dict(custom_state_dict, strict=False)
+
+        if len(missing_keys) > 0:
+            logger.warning(f"Some weights of {model.__class__.__name__} were not initialized from the model checkpoint at {pretrained_model_name_or_path} "
+                           f"and are newly initialized: {missing_keys}\n"
+                           f"You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.")
+        if len(unexpected_keys) > 0:
+            logger.warning(f"Some weights of the model checkpoint at {pretrained_model_name_or_path} were not used when "
+                           f"initializing {model.__class__.__name__}: {unexpected_keys}\n"
+                           f"This IS expected if you are initializing {model.__class__.__name__} from the checkpoint of a model trained on another task "
+                           f"or with another architecture (e.g. initializing a BertForSequenceClassification model from a BertForPreTraining model).\n"
+                           f"This IS NOT expected if you are initializing {model.__class__.__name__} from the checkpoint of a model that you expect to be exactly identical "
+                           f"(initializing a BertForSequenceClassification model from a BertForSequenceClassification model).")
+
+        if output_loading_info:
+            loading_info = {"missing_keys": missing_keys, "unexpected_keys": unexpected_keys}
+            return model, loading_info
+
+        return model