Update modeling_emotion_classifier.py

modeling_emotion_classifier.py

@@ -1,35 +1,20 @@
-import os
-import torch
+from transformers import PreTrainedModel, HubertModel
 import torch.nn as nn
-from transformers import PreTrainedModel, HubertConfig, HubertModel,PretrainedConfig
-from transformers.file_utils import (
-    WEIGHTS_NAME,
-    TF2_WEIGHTS_NAME,
-    TF_WEIGHTS_NAME,
-    cached_path,
-    hf_bucket_url,
-    is_remote_url,
-)
-from transformers.utils import logging
+import torch
 from .configuration_emotion_classifier import EmotionClassifierConfig
 
-logger = logging.get_logger(__name__)
 
 class EmotionClassifierHuBERT(PreTrainedModel):
     config_class = EmotionClassifierConfig
 
     def __init__(self, config):
         super().__init__(config)
-        
-        # Initialize HuBERT without pre-trained weights
-        hubert_config = HubertConfig.from_pretrained("facebook/hubert-large-ls960-ft")
-        self.hubert = HubertModel(hubert_config)
-        
+        self.hubert = HubertModel.from_pretrained("facebook/hubert-large-ls960-ft")
         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)
+        self.fc = nn.Linear(config.hidden_size_lstm * 2, config.num_classes)  # * 2 for bidirectional
 
     def forward(self, x):
         with torch.no_grad():
@@ -41,131 +26,4 @@ class EmotionClassifierHuBERT(PreTrainedModel):
         x = self.transformer_encoder(x)
         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, EmotionClassifierConfig):
-            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
+        return x