import os import torch import torch.nn as nn from transformers import PreTrainedModel, HubertConfig, HubertModel 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 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.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) def forward(self, x): with torch.no_grad(): features = self.hubert(x).last_hidden_state features = features.transpose(1, 2) x = torch.relu(self.conv1(features)) x = torch.relu(self.conv2(x)) x = x.transpose(1, 2) 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