text: # 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 contextlib
  import torch
  import torch.nn as nn
  
  from fairseq import checkpoint_utils, tasks, utils
  
  from fairseq.models import (
      FairseqEncoder,
      BaseFairseqModel,
      register_model,
      register_model_architecture,
  )
  
  
  def add_common_args(parser):
      parser.add_argument("--w2v-path", help="path to wav2vec 2.0 model")
      parser.add_argument(
          "--no-pretrained-weights",
          action="store_true",
          help="if true, does not load pretrained weights",
      )
      parser.add_argument(
          "--dropout-input",
          type=float,
          metavar="D",
          help="dropout to apply to the input (after feat extr)",
      )
      parser.add_argument(
          "--final-dropout",
          type=float,
          metavar="D",
          help="dropout after transformer and before final projection",
      )
      parser.add_argument(
          "--apply-mask", action="store_true", help="apply masking during fine-tuning"
      )
      parser.add_argument(
          "--dropout",
          type=float,
          metavar="D",
          help="dropout probability inside wav2vec 2.0 model",
      )
      parser.add_argument(
          "--attention-dropout",
          type=float,
          metavar="D",
          help="dropout probability for attention weights inside wav2vec 2.0 model",
      )
      parser.add_argument(
          "--activation-dropout",
          "--relu-dropout",
          type=float,
          metavar="D",
          help="dropout probability after activation in FFN inside wav2vec 2.0 model",
      )
  
      parser.add_argument(
          "--mask-length", type=int, help="repeat the mask indices multiple times"
      )
  
      parser.add_argument(
          "--mask-prob", type=float, help="probability of replacing a token with mask"
      )
  
      parser.add_argument(
          "--mask-selection",
          type=str,
          choices=["static", "uniform", "normal", "poisson"],
          help="how to choose masks",
      )
  
      parser.add_argument(
          "--mask-other",
          type=float,
          help="stdev of the mask length in case of 'normal' selection strategy",
      )
  
      parser.add_argument(
          "--no-mask-overlap",
          action="store_true",
          help="whether to allow masks to overlap",
      )
  
      parser.add_argument(
          "--mask-channel-length", type=int, help="repeat the mask indices multiple times"
      )
  
      parser.add_argument(
          "--mask-channel-prob",
          type=float,
          help="probability of replacing a token with mask",
      )
  
      parser.add_argument(
          "--mask-channel-selection",
          type=str,
          choices=["static", "uniform", "normal", "poisson"],
          help="how to choose masks",
      )
  
      parser.add_argument(
          "--mask-channel-other",
          type=float,
          help="stdev of the mask length in case of 'normal' selection strategy",
      )
  
      parser.add_argument(
          "--no-mask-channel-overlap",
          action="store_true",
          help="whether to allow masks to overlap",
      )
  
      parser.add_argument(
          "--freeze-finetune-updates",
          default=0,
          type=int,
          help="dont finetune wav2vec for this many updates",
      )
  
      parser.add_argument(
          "--feature-grad-mult",
          default=None,
          type=float,
          help="reset feature grad mult in wav2vec 2.0 to this",
      )
  
      parser.add_argument(
          "--layerdrop",
          default=0.0,
          type=float,
          help="probability of dropping a layer in wav2vec 2.0",
      )
  
      parser.add_argument(
          "--w2v-update-rate",
          default=1.0,
          type=float,
          help="probability of update encoder",
      )
  
      parser.add_argument(
          "--decoder",
          default='ctc_decoder',
          type=str,
          help="performance decoder",
      )
  
  
  @register_model("wav2vec_ctc")
  class Wav2VecCtc(BaseFairseqModel):
      @staticmethod
      def add_args(parser):
          """Add model-specific arguments to the parser."""
          add_common_args(parser)
  
      def __init__(self, w2v_encoder, args):
          super().__init__()
          self.w2v_encoder = w2v_encoder
          self.args = args
  
      def upgrade_state_dict_named(self, state_dict, name):
          super().upgrade_state_dict_named(state_dict, name)
          return state_dict
  
      @classmethod
      def build_model(cls, args, task):
          """Build a new model instance."""
          base_architecture(args)
          w2v_encoder = Wav2VecEncoder(args, task.target_dictionary)
          return cls(w2v_encoder, args)
  
      def get_normalized_probs(self, net_output, log_probs):
          """Get normalized probabilities (or log probs) from a net's output."""
  
          logits = net_output["encoder_out"]
          if log_probs:
              res = utils.log_softmax(logits.float(), dim=-1)
          else:
              res = utils.softmax(logits.float(), dim=-1)
  
          return res
  
      def forward(self, **kwargs):
          x = self.w2v_encoder(**kwargs)
          return x
  
  
  class Wav2VecEncoder(FairseqEncoder):
      def __init__(self, args, tgt_dict=None):
          self.apply_mask = args.apply_mask
          self.update_rate = args.w2v_update_rate
          arg_overrides = {
              "dropout": args.dropout,
              "activation_dropout": args.activation_dropout,
              "dropout_input": args.dropout_input,
              "attention_dropout": args.attention_dropout,
              "mask_length": args.mask_length,
              "mask_prob": args.mask_prob,
              "mask_selection": args.mask_selection,
              "mask_other": args.mask_other,
              "no_mask_overlap": args.no_mask_overlap,
              "mask_channel_length": args.mask_channel_length,
              "mask_channel_prob": args.mask_channel_prob,
              "mask_channel_selection": args.mask_channel_selection,
              "mask_channel_other": args.mask_channel_other,
              "no_mask_channel_overlap": args.no_mask_channel_overlap,
              "encoder_layerdrop": args.layerdrop,
              "feature_grad_mult": args.feature_grad_mult,
          }
          if args.w2v_path:
              print('load Wav2VecEncoder from {}'.format(args.w2v_path))
              state = checkpoint_utils.load_checkpoint_to_cpu(
                  args.w2v_path, arg_overrides
              )
              w2v_args = state["args"]
              # assert getattr(w2v_args, "w2v_path", None) is None # w2v_path is the pretrain model which should not have w2v_path
              if getattr(w2v_args, "w2v_path", None):
                  # w2v_args.w2v_path = '../libri/wav2vec2_small.pt'
                  w2v_args.w2v_path = None
          else:
              state = None
              w2v_args = args.w2v_args
  
          assert args.normalize == w2v_args.normalize, 'Fine-tuning works best when data normalization is the same'
  
          w2v_args.data = args.data
  
          if w2v_args.arch in ('wav2vec2_semi', 'wav2vec2_v1', 'wav2vec2_v2', 'wav2vec2_v3'):
              w2v_args.arch = 'wav2vec2'
              w2v_args.task = 'audio_pretraining'
              task = tasks.setup_task(w2v_args)
              model = task.build_model(w2v_args)
              model.load_state_dict(state["model"], strict=False)
          else:
              task = tasks.setup_task(w2v_args)
              model = task.build_model(w2v_args)
  
              if state is not None and not args.no_pretrained_weights:
                  model.load_state_dict(state["model"], strict=True)
  
          model.remove_pretraining_modules()
  
          super().__init__(None)
  
          d = w2v_args.encoder_embed_dim
          self.d = d
          self.w2v_model = model
  
          self.final_dropout = nn.Dropout(args.final_dropout)
          self.freeze_finetune_updates = args.freeze_finetune_updates
          self.num_updates = 0
  
          if tgt_dict is not None:
              self.proj = Linear(d, len(tgt_dict))
          else:
              self.proj = None
  
      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, source, padding_mask, tbc=False, **kwargs):
  
          w2v_args = {
              "source": source,
              "padding_mask": padding_mask,
              "mask": self.apply_mask and self.training,
          }
  
          ft = self.freeze_finetune_updates <= self.num_updates and torch.rand(1) < self.update_rate
  
          with torch.no_grad() if not ft else contextlib.ExitStack():
              x, padding_mask = self.w2v_model.extract_features(**w2v_args)
              if tbc:
                  # B x T x C -> T x B x C
                  x = x.transpose(0, 1)
  
          encoded = self.final_dropout(x)
  
          if tbc:
              encoded = encoded * (~padding_mask.transpose(0, 1)).unsqueeze(-1)
          else:
              encoded = encoded * (~padding_mask).unsqueeze(-1)
  
          res = {
              "encoder_out": encoded,
              "encoder_padding_mask": padding_mask,  # B x T
              "padding_mask": padding_mask,
          }
  
          if self.proj:
              x = self.proj(encoded)
  
              if tbc:
                  x = x * (~padding_mask.transpose(0, 1)).unsqueeze(-1)
              else:
                  x = x * (~padding_mask).unsqueeze(-1)
  
              res = {
                  "encoder_out": x,
                  "encoded": encoded,
                  "encoder_padding_mask": padding_mask,  # B x T
                  "padding_mask": padding_mask,
              }
          return res
  
      def reorder_encoder_out(self, encoder_out, new_order):
          if encoder_out["encoder_out"] is not None:
              encoder_out["encoder_out"] = encoder_out["encoder_out"].index_select(
                  1, new_order
              )
          if encoder_out["encoder_padding_mask"] is not None:
              encoder_out["encoder_padding_mask"] = encoder_out[
                  "encoder_padding_mask"
              ].index_select(0, new_order)
  
          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 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
  
  
  @register_model_architecture("wav2vec_ctc", "wav2vec_ctc")
  def base_architecture(args):
      args.no_pretrained_weights = getattr(args, "no_pretrained_weights", False)
      args.dropout_input = getattr(args, "dropout_input", 0)
      args.final_dropout = getattr(args, "final_dropout", 0)
      args.apply_mask = getattr(args, "apply_mask", False)
      args.dropout = getattr(args, "dropout", 0)
      args.attention_dropout = getattr(args, "attention_dropout", 0)
      args.activation_dropout = getattr(args, "activation_dropout", 0)
  
      args.mask_length = getattr(args, "mask_length", 10)
      args.mask_prob = getattr(args, "mask_prob", 0.5)
      args.mask_selection = getattr(args, "mask_selection", "static")
      args.mask_other = getattr(args, "mask_other", 0)
      args.no_mask_overlap = getattr(args, "no_mask_overlap", False)
      args.mask_channel_length = getattr(args, "mask_channel_length", 10)
      args.mask_channel_prob = getattr(args, "mask_channel_prob", 0.5)
      args.mask_channel_selection = getattr(args, "mask_channel_selection", "static")
      args.mask_channel_other = getattr(args, "mask_channel_other", 0)
      args.no_mask_channel_overlap = getattr(args, "no_mask_channel_overlap", False)
  
      args.freeze_finetune_updates = getattr(args, "freeze_finetune_updates", 0)
      args.feature_grad_mult = getattr(args, "feature_grad_mult", 0)
      args.layerdrop = getattr(args, "layerdrop", 0.0)
  