Spaces:
Runtime error
Runtime error
| # Copyright (c) Facebook, Inc. and its affiliates. | |
| # | |
| # This source code is licensed under the MIT license found in the | |
| # LICENSE file in the root directory of this source tree. | |
| import logging | |
| import contextlib | |
| from argparse import Namespace | |
| from typing import Any, Optional | |
| import torch | |
| import torch.nn as nn | |
| from dataclasses import dataclass, field | |
| from fairseq import checkpoint_utils, tasks, utils | |
| from fairseq.dataclass import FairseqDataclass | |
| from fairseq.dataclass.utils import convert_namespace_to_omegaconf | |
| from fairseq.models import BaseFairseqModel, FairseqEncoder, register_model | |
| from fairseq.models.hubert.hubert import MASKING_DISTRIBUTION_CHOICES | |
| from fairseq.tasks import FairseqTask | |
| from omegaconf import II, MISSING | |
| from .hubert_asr import HubertAsrConfig | |
| from fairseq.models.transformer import TransformerConfig | |
| logger = logging.getLogger(__name__) | |
| class HubertMTConfig(HubertAsrConfig): | |
| load_pretrained_mbart_from: Optional[str] = field( | |
| default=None, | |
| metadata={ | |
| "help": "model to take text encoder decoder weights from (for initialization)" | |
| }, | |
| ) | |
| use_rel_pos_enc: bool = field( | |
| default=True, | |
| metadata={"help": "whether to use relative positional encoding"}, | |
| ) | |
| text_transformer_encoder_layers: int = field( | |
| default=12, | |
| metadata={"help": "reset text_transformer_encoder_layers"}, | |
| ) | |
| class HubertMT(BaseFairseqModel): | |
| def __init__(self, cfg: HubertMTConfig, w2v_encoder: BaseFairseqModel): | |
| super().__init__() | |
| self.cfg = cfg | |
| self.w2v_encoder = w2v_encoder | |
| def upgrade_state_dict_named(self, state_dict, name): | |
| super().upgrade_state_dict_named(state_dict, name) | |
| return state_dict | |
| def build_model(cls, cfg: HubertMTConfig, task: FairseqTask): | |
| """Build a new model instance.""" | |
| w2v_encoder = HubertEncoder(cfg, task.target_dictionary) | |
| return cls(cfg, w2v_encoder) | |
| def get_normalized_probs(self, net_output, log_probs, sample=None): | |
| """Get normalized probabilities (or log probs) from a net's output.""" | |
| if "decoder_out" in net_output: | |
| return self.w2v_encoder.get_normalized_probs_decoder(net_output["decoder_out"], log_probs, sample) | |
| assert "encoder_out" not in net_output | |
| if "encoder_out" not in net_output: | |
| return self.w2v_encoder.get_normalized_probs_decoder(net_output, log_probs, sample) | |
| def get_logits(self, net_output): | |
| logits = net_output["encoder_out"] | |
| padding = net_output["encoder_padding_mask"] | |
| if padding is not None and padding.any(): | |
| padding = padding.T | |
| logits[padding][..., 0] = 0 | |
| logits[padding][..., 1:] = float("-inf") | |
| return logits | |
| def forward(self, **kwargs): | |
| x = self.w2v_encoder(**kwargs) | |
| return x | |
| def encoder(self): | |
| return self.w2v_encoder | |
| def reorder_encoder_out(self, encoder_out, new_order): | |
| return self.encoder.reorder_encoder_out(encoder_out, new_order) | |
| def decoder(self): | |
| return self.w2v_encoder.w2v_model.decoder | |
| class HubertEncoder(FairseqEncoder): | |
| def __init__(self, cfg: HubertMTConfig, tgt_dict=None): | |
| self.apply_mask = cfg.apply_mask | |
| arg_overrides = { | |
| "dropout": cfg.dropout, | |
| "activation_dropout": cfg.activation_dropout, | |
| "dropout_input": cfg.dropout_input, | |
| "attention_dropout": cfg.attention_dropout, | |
| "mask_length": cfg.mask_length, | |
| "mask_prob": cfg.mask_prob, | |
| "mask_selection": cfg.mask_selection, | |
| "mask_other": cfg.mask_other, | |
| "no_mask_overlap": cfg.no_mask_overlap, | |
| "mask_channel_length": cfg.mask_channel_length, | |
| "mask_channel_prob": cfg.mask_channel_prob, | |
| "mask_channel_selection": cfg.mask_channel_selection, | |
| "mask_channel_other": cfg.mask_channel_other, | |
| "no_mask_channel_overlap": cfg.no_mask_channel_overlap, | |
| "encoder_layerdrop": cfg.layerdrop, | |
| "decoder_layerdrop": cfg.decoder_layerdrop, | |
| "feature_grad_mult": cfg.feature_grad_mult, | |
| "decoder_dict_size": -1, | |
| "add_text_modality": True, | |
| "add_text_encoder": True, | |
| "load_pretrained_mbart_from": None, | |
| "load_pretrained_w2v_from": None, | |
| "text_transformer": { | |
| "encoder":{ | |
| "layers": cfg.text_transformer_encoder_layers, | |
| "layerdrop": cfg.layerdrop, | |
| }, | |
| 'dropout': cfg.dropout, | |
| 'attention_dropout': cfg.attention_dropout, | |
| 'activation_dropout': cfg.activation_dropout, | |
| } | |
| } | |
| if cfg.no_pretrained_weights: | |
| arg_overrides["use_rel_pos_enc"] = cfg.use_rel_pos_enc | |
| if cfg.w2v_args is None: | |
| state = checkpoint_utils.load_checkpoint_to_cpu( | |
| cfg.w2v_path, arg_overrides | |
| ) | |
| w2v_args = state.get("cfg", None) | |
| if w2v_args is None: | |
| w2v_args = convert_namespace_to_omegaconf(state["args"]) | |
| cfg.w2v_args = w2v_args | |
| else: | |
| state = None | |
| w2v_args = cfg.w2v_args | |
| if isinstance(w2v_args, Namespace): | |
| cfg.w2v_args = w2v_args = convert_namespace_to_omegaconf(w2v_args) | |
| # logger.info("---------------------state.keys()-------------------------------------------") | |
| # logger.info(state.keys()) | |
| # logger.info("---------------------w2v_args.task-------------------------------------------") | |
| # logger.info(w2v_args.task) | |
| # logger.info("---------------------w2v_args.model-------------------------------------------") | |
| # logger.info(w2v_args.model) | |
| # logger.info("----------------------------------------------------------------") | |
| w2v_args.task.data = cfg.data | |
| w2v_args.task.text_cfg.text_data = cfg.data | |
| w2v_args.task.text_cfg.data_config = None | |
| task = tasks.setup_task(w2v_args.task) | |
| if state is not None and "task_state" in state: | |
| # This will load the stored "dictionaries" object | |
| task.load_state_dict(state["task_state"]) | |
| model = task.build_model(w2v_args.model) | |
| ### load mbart if specificed | |
| if cfg.load_pretrained_mbart_from is not None and cfg.no_pretrained_weights: | |
| logger.info("Loading mbart....") | |
| mbart_model_state = model.load_checkpoint(cfg.load_pretrained_mbart_from) | |
| model.text_encoder = model.load_pretrained_component_from_model( | |
| component=model.text_encoder, state=mbart_model_state | |
| ) | |
| model.decoder = model.load_pretrained_component_from_model( | |
| component=model.decoder, state=mbart_model_state | |
| ) | |
| if state is not None and not cfg.no_pretrained_weights: | |
| logger.info("Loading pre-trained models....") | |
| model.load_state_dict(state["model"], strict=True) | |
| ### remove_pretraining_modules model.remove_pretraining_modules() | |
| model.target_glu = None | |
| model.final_proj = None | |
| model.feature_extractor = None | |
| model.post_extract_proj = None | |
| model.encoder = None | |
| dropout_keys = [ n for n in w2v_args.model.text_transformer if n.find("drop") >= 0 ] | |
| for key in dropout_keys: | |
| logger.info(f"{key}: {w2v_args.model.text_transformer[key]}") | |
| super().__init__(task.source_dictionary) | |
| d = w2v_args.model.encoder_embed_dim | |
| self.w2v_model = model | |
| self.final_dropout = nn.Dropout(cfg.final_dropout) | |
| self.freeze_finetune_updates = cfg.freeze_finetune_updates | |
| self.freeze_decoder_updates = cfg.freeze_decoder_updates | |
| self.num_updates = 0 | |
| def set_num_updates(self, num_updates): | |
| """Set the number of parameters updates.""" | |
| super().set_num_updates(num_updates) | |
| self.num_updates = num_updates | |
| def forward(self, src_tokens, src_lengths, prev_output_tokens, tbc=True, **kwargs): | |
| # ft = self.freeze_finetune_updates <= self.num_updates | |
| w2v_args = { | |
| "src_tokens": src_tokens, | |
| "src_lengths": src_lengths, | |
| "mask": self.apply_mask and self.training, | |
| "prev_output_tokens": prev_output_tokens, | |
| } | |
| results = self.w2v_model(**w2v_args) | |
| return results | |
| def get_normalized_probs_decoder(self, net_output, log_probs, sample=None): | |
| # net_output['encoder_out'] is a (B, T, D) tensor | |
| return self.w2v_model.get_normalized_probs(net_output, log_probs, sample) | |
| def reorder_encoder_out(self, encoder_out, new_order): | |
| if encoder_out["encoder_out"] is not None: | |
| if isinstance(encoder_out["encoder_out"], list): | |
| encoder_out["encoder_out"] = ( | |
| [] if len(encoder_out["encoder_out"]) == 0 | |
| else [x.index_select(1, new_order) for x in encoder_out["encoder_out"]] | |
| ) | |
| else: | |
| encoder_out["encoder_out"] = encoder_out[ | |
| "encoder_out" | |
| ].index_select(1, new_order) | |
| if encoder_out["encoder_padding_mask"] is not None: | |
| if isinstance(encoder_out["encoder_padding_mask"], list): | |
| encoder_out["encoder_padding_mask"] = ( | |
| [] if len(encoder_out["encoder_padding_mask"]) == 0 | |
| else [x.index_select(0, new_order) for x in encoder_out["encoder_padding_mask"]] | |
| ) | |
| else: | |
| encoder_out["encoder_padding_mask"] = encoder_out[ | |
| "encoder_padding_mask" | |
| ].index_select(0, new_order) | |
| if "decoder_out" in encoder_out and encoder_out["decoder_out"] is not None: | |
| if isinstance(encoder_out["decoder_out"], list): | |
| encoder_out["decoder_out"] = ( | |
| [] if len(encoder_out["decoder_out"]) == 0 | |
| else [x.index_select(0, new_order) for x in encoder_out["decoder_out"]] | |
| ) | |
| else: | |
| encoder_out["decoder_out"] = encoder_out[ | |
| "decoder_out" | |
| ].index_select(0, new_order) | |
| if "encoder_out_for_ctc" in encoder_out and encoder_out["encoder_out_for_ctc"] is not None: | |
| if isinstance(encoder_out["encoder_out_for_ctc"], list): | |
| encoder_out["encoder_out_for_ctc"] = ( | |
| [] if len(encoder_out["encoder_out_for_ctc"]) == 0 | |
| else [x.index_select(1, new_order) for x in encoder_out["encoder_out_for_ctc"]] | |
| ) | |
| else: | |
| encoder_out["encoder_out_for_ctc"] = encoder_out[ | |
| "encoder_out_for_ctc" | |
| ].index_select(1, new_order) | |
| return encoder_out | |
| def forward_torchscript(self, net_input): | |
| """A TorchScript-compatible version of forward. | |
| Encoders which use additional arguments may want to override | |
| this method for TorchScript compatibility. | |
| """ | |
| encoder_out = self.w2v_model.forward_torchscript(net_input) | |
| if "encoder_out_for_ctc" in encoder_out: | |
| del encoder_out['encoder_out_for_ctc'] | |
| return encoder_out | |
| def max_positions(self): | |
| """Maximum input length supported by the encoder.""" | |
| return None | |
| def upgrade_state_dict_named(self, state_dict, name): | |
| return state_dict | |
| def Embedding(num_embeddings, embedding_dim, padding_idx): | |
| m = nn.Embedding(num_embeddings, embedding_dim, padding_idx=padding_idx) | |
| nn.init.normal_(m.weight, mean=0, std=embedding_dim ** -0.5) | |
| nn.init.constant_(m.weight[padding_idx], 0) | |
| return m | |
| def Linear(in_features, out_features, bias=True): | |
| m = nn.Linear(in_features, out_features, bias) | |
| nn.init.xavier_uniform_(m.weight) | |
| if bias: | |
| nn.init.constant_(m.bias, 0.0) | |
| return m | |