"""simple docstring"""


import argparse
import json
import os

import fairseq
import torch
from fairseq.data import Dictionary

from transformers import (
    WavaVecaConfig,
    WavaVecaCTCTokenizer,
    WavaVecaFeatureExtractor,
    WavaVecaForCTC,
    WavaVecaForPreTraining,
    WavaVecaProcessor,
    logging,
)
from transformers.models.wavaveca.modeling_wavaveca import WavaVecaForSequenceClassification


logging.set_verbosity_info()
UpperCAmelCase__  		=						logging.get_logger(__name__)

UpperCAmelCase__  		=						{
    """post_extract_proj""": """feature_projection.projection""",
    """encoder.pos_conv.0""": """encoder.pos_conv_embed.conv""",
    """self_attn.k_proj""": """encoder.layers.*.attention.k_proj""",
    """self_attn.v_proj""": """encoder.layers.*.attention.v_proj""",
    """self_attn.q_proj""": """encoder.layers.*.attention.q_proj""",
    """self_attn.out_proj""": """encoder.layers.*.attention.out_proj""",
    """self_attn_layer_norm""": """encoder.layers.*.layer_norm""",
    """fc1""": """encoder.layers.*.feed_forward.intermediate_dense""",
    """fc2""": """encoder.layers.*.feed_forward.output_dense""",
    """final_layer_norm""": """encoder.layers.*.final_layer_norm""",
    """encoder.layer_norm""": """encoder.layer_norm""",
    """adapter_layer""": """encoder.layers.*.adapter_layer""",
    """w2v_model.layer_norm""": """feature_projection.layer_norm""",
    """quantizer.weight_proj""": """quantizer.weight_proj""",
    """quantizer.vars""": """quantizer.codevectors""",
    """project_q""": """project_q""",
    """final_proj""": """project_hid""",
    """w2v_encoder.proj""": """lm_head""",
    """mask_emb""": """masked_spec_embed""",
    """pooling_layer.linear""": """projector""",
    """pooling_layer.projection""": """classifier""",
}
UpperCAmelCase__  		=						[
    """lm_head""",
    """quantizer.weight_proj""",
    """quantizer.codevectors""",
    """project_q""",
    """project_hid""",
    """projector""",
    """classifier""",
]






def     __UpperCAmelCase					(			lowercase		):







  """simple docstring"""
  _UpperCAmelCase						 =      {}
  with open(lowercase   ,"""r"""		) as file:
    for line_number, line in enumerate(lowercase		):
      _UpperCAmelCase						 =      line.strip()
      if line:
        _UpperCAmelCase						 =      line.split()
        _UpperCAmelCase						 =      line_number
        _UpperCAmelCase						 =      words[0]
        _UpperCAmelCase						 =      value
  return result






def     __UpperCAmelCase					(			lowercase   ,lowercase   ,lowercase   ,lowercase   ,lowercase		):







  """simple docstring"""
  for attribute in key.split("""."""		):
    _UpperCAmelCase						 =      getattr(lowercase   ,lowercase		)

  _UpperCAmelCase						 =      None
  for param_key in PARAM_MAPPING.keys():
    if full_name.endswith(lowercase		):
      _UpperCAmelCase						 =      PARAM_MAPPING[full_name.split("""."""		)[-1]]
      _UpperCAmelCase						 =      """param"""

  if weight_type is not None and weight_type != "param":
    _UpperCAmelCase						 =      getattr(lowercase   ,lowercase		).shape
  elif weight_type is not None and weight_type == "param":
    _UpperCAmelCase						 =      hf_pointer
    for attribute in hf_param_name.split("""."""		):
      _UpperCAmelCase						 =      getattr(lowercase   ,lowercase		)
    _UpperCAmelCase						 =      shape_pointer.shape

    # let's reduce dimension
    _UpperCAmelCase						 =      value[0]
  else:
    _UpperCAmelCase						 =      hf_pointer.shape

  if hf_shape != value.shape:
    raise ValueError(
        f'''Shape of hf {key + "." + weight_type if weight_type is not None else ""} is {hf_shape}, but should be'''
        f''' {value.shape} for {full_name}'''		)

  if weight_type == "weight":
    _UpperCAmelCase						 =      value
  elif weight_type == "weight_g":
    _UpperCAmelCase						 =      value
  elif weight_type == "weight_v":
    _UpperCAmelCase						 =      value
  elif weight_type == "bias":
    _UpperCAmelCase						 =      value
  elif weight_type == "param":
    for attribute in hf_param_name.split("""."""		):
      _UpperCAmelCase						 =      getattr(lowercase   ,lowercase		)
    _UpperCAmelCase						 =      value
  else:
    _UpperCAmelCase						 =      value

  logger.info(f'''{key + "." + weight_type if weight_type is not None else ""} was initialized from {full_name}.'''		)






def     __UpperCAmelCase					(			lowercase   ,lowercase   ,lowercase   ,lowercase   ,lowercase		):







  """simple docstring"""
  _UpperCAmelCase						 =      None
  for param_key in PARAM_MAPPING.keys():
    if full_name.endswith(lowercase		):
      _UpperCAmelCase						 =      PARAM_MAPPING[full_name.split("""."""		)[-1]]
      _UpperCAmelCase						 =      """param"""

  if weight_type is not None and weight_type != "param":
    _UpperCAmelCase						 =      """.""".join([key, weight_type]		)
  elif weight_type is not None and weight_type == "param":
    _UpperCAmelCase						 =      """.""".join([key, hf_param_name]		)
  else:
    _UpperCAmelCase						 =      key

  _UpperCAmelCase						 =      value if """lm_head""" in full_key else value[0]


UpperCAmelCase__  		=						{
    """W_a""": """linear_1.weight""",
    """W_b""": """linear_2.weight""",
    """b_a""": """linear_1.bias""",
    """b_b""": """linear_2.bias""",
    """ln_W""": """norm.weight""",
    """ln_b""": """norm.bias""",
}






def     __UpperCAmelCase					(			lowercase   ,lowercase   ,lowercase=None   ,lowercase=None		):







  """simple docstring"""
  _UpperCAmelCase						 =      False
  for key, mapped_key in MAPPING.items():
    _UpperCAmelCase						 =      """wav2vec2.""" + mapped_key if mapped_key not in TOP_LEVEL_KEYS else mapped_key
    if key in name or key.split("""w2v_model."""		)[-1] == name.split("""."""		)[0]:
      _UpperCAmelCase						 =      True
      if "*" in mapped_key:
        _UpperCAmelCase						 =      name.split(lowercase		)[0].split("""."""		)[-2]
        _UpperCAmelCase						 =      mapped_key.replace("""*"""   ,lowercase		)
      if "weight_g" in name:
        _UpperCAmelCase						 =      """weight_g"""
      elif "weight_v" in name:
        _UpperCAmelCase						 =      """weight_v"""
      elif "bias" in name:
        _UpperCAmelCase						 =      """bias"""
      elif "weight" in name:
        # TODO: don't match quantizer.weight_proj
        _UpperCAmelCase						 =      """weight"""
      else:
        _UpperCAmelCase						 =      None
      if hf_dict is not None:
        rename_dict(lowercase   ,lowercase   ,lowercase   ,lowercase   ,lowercase		)
      else:
        set_recursively(lowercase   ,lowercase   ,lowercase   ,lowercase   ,lowercase		)
      return is_used
  return is_used






def     __UpperCAmelCase					(			lowercase   ,lowercase   ,lowercase		):







  """simple docstring"""
  _UpperCAmelCase						 =      []
  _UpperCAmelCase						 =      fairseq_model.state_dict()

  _UpperCAmelCase						 =      hf_model.wavaveca.feature_extractor

  for name, value in fairseq_dict.items():
    _UpperCAmelCase						 =      False
    if "conv_layers" in name:
      load_conv_layer(
          lowercase   ,lowercase   ,lowercase   ,lowercase   ,hf_model.config.feat_extract_norm == """group"""   ,)
      _UpperCAmelCase						 =      True
    else:
      _UpperCAmelCase						 =      load_wavaveca_layer(lowercase   ,lowercase   ,lowercase		)
    if not is_used:
      unused_weights.append(lowercase		)

  logger.warning(f'''Unused weights: {unused_weights}'''		)






def     __UpperCAmelCase					(			lowercase   ,lowercase   ,lowercase   ,lowercase   ,lowercase		):







  """simple docstring"""
  _UpperCAmelCase						 =      full_name.split("""conv_layers."""		)[-1]
  _UpperCAmelCase						 =      name.split("""."""		)
  _UpperCAmelCase						 =      int(items[0]		)
  _UpperCAmelCase						 =      int(items[1]		)

  if type_id == 0:
    if "bias" in name:
      if value.shape != feature_extractor.conv_layers[layer_id].conv.bias.data.shape:
        raise ValueError(
            f'''{full_name} has size {value.shape}, but'''
            f''' {feature_extractor.conv_layers[layer_id].conv.bias.data.shape} was found.'''		)
      _UpperCAmelCase						 =      value
      logger.info(f'''Feat extract conv layer {layer_id} was initialized from {full_name}.'''		)
    elif "weight" in name:
      if value.shape != feature_extractor.conv_layers[layer_id].conv.weight.data.shape:
        raise ValueError(
            f'''{full_name} has size {value.shape}, but'''
            f''' {feature_extractor.conv_layers[layer_id].conv.weight.data.shape} was found.'''		)
      _UpperCAmelCase						 =      value
      logger.info(f'''Feat extract conv layer {layer_id} was initialized from {full_name}.'''		)
  elif (type_id == 2 and not use_group_norm) or (type_id == 2 and layer_id == 0 and use_group_norm):
    if "bias" in name:
      if value.shape != feature_extractor.conv_layers[layer_id].layer_norm.bias.data.shape:
        raise ValueError(
            f'''{full_name} has size {value.shape}, but'''
            f''' {feature_extractor.conv_layers[layer_id].layer_norm.bias.data.shape} was found.'''		)
      _UpperCAmelCase						 =      value
      logger.info(f'''Feat extract layer norm weight of layer {layer_id} was initialized from {full_name}.'''		)
    elif "weight" in name:
      if value.shape != feature_extractor.conv_layers[layer_id].layer_norm.weight.data.shape:
        raise ValueError(
            f'''{full_name} has size {value.shape}, but'''
            f''' {feature_extractor.conv_layers[layer_id].layer_norm.weight.data.shape} was found.'''		)
      _UpperCAmelCase						 =      value
      logger.info(f'''Feat extract layer norm weight of layer {layer_id} was initialized from {full_name}.'''		)
  else:
    unused_weights.append(lowercase		)






@torch.no_grad()
def     __UpperCAmelCase					(			lowercase   ,lowercase   ,lowercase=None   ,lowercase=None   ,lowercase=True   ,lowercase=False		):







  """simple docstring"""
  if config_path is not None:
    _UpperCAmelCase						 =      WavaVecaConfig.from_pretrained(lowercase		)
  else:
    _UpperCAmelCase						 =      WavaVecaConfig()

  if is_seq_class:
    _UpperCAmelCase						 =      read_txt_into_dict(lowercase		)
    _UpperCAmelCase						 =      idalabel
    _UpperCAmelCase						 =      WavaVecaForSequenceClassification(lowercase		)
    _UpperCAmelCase						 =      WavaVecaFeatureExtractor(
        feature_size=1   ,sampling_rate=1_60_00   ,padding_value=0   ,do_normalize=lowercase   ,return_attention_mask=lowercase   ,)
    feature_extractor.save_pretrained(lowercase		)

  elif is_finetuned:
    if dict_path:
      _UpperCAmelCase						 =      Dictionary.load(lowercase		)

      # important change bos & pad token id since CTC symbol is <pad> and
      # not <s> as in fairseq
      _UpperCAmelCase						 =      target_dict.pad_index
      _UpperCAmelCase						 =      target_dict.bos_index
      _UpperCAmelCase						 =      target_dict.eos_index
      _UpperCAmelCase						 =      len(target_dict.symbols		)
      _UpperCAmelCase						 =      os.path.join(lowercase   ,"""vocab.json"""		)
      if not os.path.isdir(lowercase		):
        logger.error("""--pytorch_dump_folder_path ({}) should be a directory""".format(lowercase		)		)
        return
      os.makedirs(lowercase   ,exist_ok=lowercase		)
      _UpperCAmelCase						 =      target_dict.indices

      # fairseq has the <pad> and <s> switched
      _UpperCAmelCase						 =      0
      _UpperCAmelCase						 =      1
      with open(lowercase   ,"""w"""   ,encoding="""utf-8"""		) as vocab_handle:
        json.dump(lowercase   ,lowercase		)
      _UpperCAmelCase						 =      WavaVecaCTCTokenizer(
          lowercase   ,unk_token=target_dict.unk_word   ,pad_token=target_dict.pad_word   ,bos_token=target_dict.bos_word   ,eos_token=target_dict.eos_word   ,word_delimiter_token="""|"""   ,do_lower_case=lowercase   ,)
      _UpperCAmelCase						 =      True if config.feat_extract_norm == """layer""" else False
      _UpperCAmelCase						 =      WavaVecaFeatureExtractor(
          feature_size=1   ,sampling_rate=1_60_00   ,padding_value=0   ,do_normalize=lowercase   ,return_attention_mask=lowercase   ,)
      _UpperCAmelCase						 =      WavaVecaProcessor(feature_extractor=lowercase   ,tokenizer=lowercase		)
      processor.save_pretrained(lowercase		)

    _UpperCAmelCase						 =      WavaVecaForCTC(lowercase		)
  else:
    _UpperCAmelCase						 =      WavaVecaForPreTraining(lowercase		)

  if is_finetuned or is_seq_class:
    _UpperCAmelCase ,				_UpperCAmelCase ,				_UpperCAmelCase						 =      fairseq.checkpoint_utils.load_model_ensemble_and_task(
        [checkpoint_path]   ,arg_overrides={"""data""": """/""".join(dict_path.split("""/"""		)[:-1]		)}		)
  else:
    _UpperCAmelCase						 =      argparse.Namespace(task="""audio_pretraining"""		)
    _UpperCAmelCase						 =      fairseq.tasks.setup_task(lowercase		)

    _UpperCAmelCase ,				_UpperCAmelCase ,				_UpperCAmelCase						 =      fairseq.checkpoint_utils.load_model_ensemble_and_task([checkpoint_path]   ,task=lowercase		)

  _UpperCAmelCase						 =      model[0].eval()

  recursively_load_weights(lowercase   ,lowercase   ,not is_finetuned		)

  hf_wavavec.save_pretrained(lowercase		)


if __name__ == "__main__":
    UpperCAmelCase__  		=						argparse.ArgumentParser()
    parser.add_argument("""--pytorch_dump_folder_path""", default=None, type=str, help="""Path to the output PyTorch model.""")
    parser.add_argument("""--checkpoint_path""", default=None, type=str, help="""Path to fairseq checkpoint""")
    parser.add_argument("""--dict_path""", default=None, type=str, help="""Path to dict of fine-tuned model""")
    parser.add_argument("""--config_path""", default=None, type=str, help="""Path to hf config.json of model to convert""")
    parser.add_argument(
        """--not_finetuned""", action="""store_true""", help="""Whether the model to convert is a fine-tuned model or not"""
    )
    parser.add_argument(
        """--is_seq_class""",
        action="""store_true""",
        help="""Whether the model to convert is a fine-tuned sequence classification model or not""",
    )
    UpperCAmelCase__  		=						parser.parse_args()

    UpperCAmelCase__  		=						not args.not_finetuned and not args.is_seq_class
    convert_wavaveca_checkpoint(
        args.checkpoint_path,
        args.pytorch_dump_folder_path,
        args.config_path,
        args.dict_path,
        is_finetuned,
        args.is_seq_class,
    )

"""simple docstring"""


import torch
from torch import nn


class        a				(		nn.Module   ):


    def __init__(					self	:  List[Any]   ,				__lowerCAmelCase	:  Optional[Any]   ,				__lowerCAmelCase	:  str   ,				__lowerCAmelCase	:  List[Any]   ,				__lowerCAmelCase	:  int   ,				__lowerCAmelCase	:  List[Any]=1   ,				__lowerCAmelCase	:  Tuple=False					):
      super().__init__()

      _UpperCAmelCase						 =      n_token
      _UpperCAmelCase						 =      d_embed
      _UpperCAmelCase						 =      d_proj

      _UpperCAmelCase						 =      cutoffs + [n_token]
      _UpperCAmelCase						 =      [0] + self.cutoffs
      _UpperCAmelCase						 =      div_val

      _UpperCAmelCase						 =      self.cutoffs[0]
      _UpperCAmelCase						 =      len(self.cutoffs					) - 1
      _UpperCAmelCase						 =      self.shortlist_size + self.n_clusters

      if self.n_clusters > 0:
        _UpperCAmelCase						 =      nn.Parameter(torch.zeros(self.n_clusters   ,				self.d_embed					)					)
        _UpperCAmelCase						 =      nn.Parameter(torch.zeros(self.n_clusters					)					)

      _UpperCAmelCase						 =      nn.ModuleList()
      _UpperCAmelCase						 =      nn.ParameterList()

      if div_val == 1:
        for i in range(len(self.cutoffs					)					):
          if d_proj != d_embed:
            self.out_projs.append(nn.Parameter(torch.FloatTensor(__lowerCAmelCase   ,				__lowerCAmelCase					)					)					)
          else:
            self.out_projs.append(__lowerCAmelCase					)

        self.out_layers.append(nn.Linear(__lowerCAmelCase   ,				__lowerCAmelCase					)					)
      else:
        for i in range(len(self.cutoffs					)					):
          _UpperCAmelCase ,				_UpperCAmelCase						 =      self.cutoff_ends[i], self.cutoff_ends[i + 1]
          _UpperCAmelCase						 =      d_embed // (div_val**i)

          self.out_projs.append(nn.Parameter(torch.FloatTensor(__lowerCAmelCase   ,				__lowerCAmelCase					)					)					)

          self.out_layers.append(nn.Linear(__lowerCAmelCase   ,				r_idx - l_idx					)					)

      _UpperCAmelCase						 =      keep_order


    def      lowerCAmelCase_  (					self	:  Optional[Any]   ,				__lowerCAmelCase	:  Optional[Any]   ,				__lowerCAmelCase	:  Dict   ,				__lowerCAmelCase	:  int   ,				__lowerCAmelCase	:  Tuple					):
      if proj is None:
        _UpperCAmelCase						 =      nn.functional.linear(__lowerCAmelCase   ,				__lowerCAmelCase   ,				bias=__lowerCAmelCase					)
      else:
        # if CUDA_MAJOR <= 9 and CUDA_MINOR <= 1:
        _UpperCAmelCase						 =      nn.functional.linear(__lowerCAmelCase   ,				proj.t().contiguous()					)
        _UpperCAmelCase						 =      nn.functional.linear(__lowerCAmelCase   ,				__lowerCAmelCase   ,				bias=__lowerCAmelCase					)
        # else:
        #     logit = torch.einsum('bd,de,ev->bv', (hidden, proj, weight.t()))
        #     if bias is not None:
        #         logit = logit + bias

      return logit


    def      lowerCAmelCase_  (					self	:  List[Any]   ,				__lowerCAmelCase	:  Optional[int]   ,				__lowerCAmelCase	:  Optional[Any]=None   ,				__lowerCAmelCase	:  Union[str, Any]=False					):

      if labels is not None:
        # Shift so that tokens < n predict n
        _UpperCAmelCase						 =      hidden[..., :-1, :].contiguous()
        _UpperCAmelCase						 =      labels[..., 1:].contiguous()
        _UpperCAmelCase						 =      hidden.view(-1   ,				hidden.size(-1					)					)
        _UpperCAmelCase						 =      labels.view(-1					)
        if hidden.size(0					) != labels.size(0					):
          raise RuntimeError("""Input and labels should have the same size in the batch dimension."""					)
      else:
        _UpperCAmelCase						 =      hidden.view(-1   ,				hidden.size(-1					)					)

      if self.n_clusters == 0:
        _UpperCAmelCase						 =      self._compute_logit(__lowerCAmelCase   ,				self.out_layers[0].weight   ,				self.out_layers[0].bias   ,				self.out_projs[0]					)
        if labels is not None:
          _UpperCAmelCase						 =      labels != -100
          _UpperCAmelCase						 =      torch.zeros_like(__lowerCAmelCase   ,				dtype=hidden.dtype   ,				device=hidden.device					)
          _UpperCAmelCase						 =      (
              -nn.functional.log_softmax(__lowerCAmelCase   ,				dim=-1					)[mask].gather(1   ,				labels[mask].unsqueeze(1					)					).squeeze(1					)
          )
        else:
          _UpperCAmelCase						 =      nn.functional.log_softmax(__lowerCAmelCase   ,				dim=-1					)
      else:
        # construct weights and biases
        _UpperCAmelCase ,				_UpperCAmelCase						 =      [], []
        for i in range(len(self.cutoffs					)					):
          if self.div_val == 1:
            _UpperCAmelCase ,				_UpperCAmelCase						 =      self.cutoff_ends[i], self.cutoff_ends[i + 1]
            _UpperCAmelCase						 =      self.out_layers[0].weight[l_idx:r_idx]
            _UpperCAmelCase						 =      self.out_layers[0].bias[l_idx:r_idx]
          else:
            _UpperCAmelCase						 =      self.out_layers[i].weight
            _UpperCAmelCase						 =      self.out_layers[i].bias

          if i == 0:
            _UpperCAmelCase						 =      torch.cat([weight_i, self.cluster_weight]   ,				dim=0					)
            _UpperCAmelCase						 =      torch.cat([bias_i, self.cluster_bias]   ,				dim=0					)

          weights.append(__lowerCAmelCase					)
          biases.append(__lowerCAmelCase					)

        _UpperCAmelCase ,				_UpperCAmelCase ,				_UpperCAmelCase						 =      weights[0], biases[0], self.out_projs[0]

        _UpperCAmelCase						 =      self._compute_logit(__lowerCAmelCase   ,				__lowerCAmelCase   ,				__lowerCAmelCase   ,				__lowerCAmelCase					)
        _UpperCAmelCase						 =      nn.functional.log_softmax(__lowerCAmelCase   ,				dim=1					)

        if labels is None:
          _UpperCAmelCase						 =      hidden.new_empty((head_logit.size(0					), self.n_token)					)
        else:
          _UpperCAmelCase						 =      torch.zeros_like(__lowerCAmelCase   ,				dtype=hidden.dtype   ,				device=hidden.device					)

        _UpperCAmelCase						 =      0
        _UpperCAmelCase						 =      [0] + self.cutoffs
        for i in range(len(__lowerCAmelCase					) - 1					):
          _UpperCAmelCase ,				_UpperCAmelCase						 =      cutoff_values[i], cutoff_values[i + 1]

          if labels is not None:
            _UpperCAmelCase						 =      (labels >= l_idx) & (labels < r_idx)
            _UpperCAmelCase						 =      mask_i.nonzero().squeeze()

            if indices_i.numel() == 0:
              continue

            _UpperCAmelCase						 =      labels.index_select(0   ,				__lowerCAmelCase					) - l_idx
            _UpperCAmelCase						 =      head_logprob.index_select(0   ,				__lowerCAmelCase					)
            _UpperCAmelCase						 =      hidden.index_select(0   ,				__lowerCAmelCase					)
          else:
            _UpperCAmelCase						 =      hidden

          if i == 0:
            if labels is not None:
              _UpperCAmelCase						 =      head_logprob_i.gather(1   ,				target_i[:, None]					).squeeze(1					)
            else:
              _UpperCAmelCase						 =      head_logprob[:, : self.cutoffs[0]]
          else:
            _UpperCAmelCase ,				_UpperCAmelCase ,				_UpperCAmelCase						 =      weights[i], biases[i], self.out_projs[i]

            _UpperCAmelCase						 =      self._compute_logit(__lowerCAmelCase   ,				__lowerCAmelCase   ,				__lowerCAmelCase   ,				__lowerCAmelCase					)
            _UpperCAmelCase						 =      nn.functional.log_softmax(__lowerCAmelCase   ,				dim=1					)
            _UpperCAmelCase						 =      self.cutoffs[0] + i - 1  # No probability for the head cluster
            if labels is not None:
              _UpperCAmelCase						 =      head_logprob_i[:, cluster_prob_idx] + tail_logprob_i.gather(
                  1   ,				target_i[:, None]					).squeeze(1					)
            else:
              _UpperCAmelCase						 =      head_logprob[:, cluster_prob_idx, None] + tail_logprob_i
              _UpperCAmelCase						 =      logprob_i

          if labels is not None:
            if (hasattr(self   ,				"""keep_order"""					) and self.keep_order) or keep_order:
              out.index_copy_(0   ,				__lowerCAmelCase   ,				-logprob_i					)
            else:
              out[offset : offset + logprob_i.size(0					)].copy_(-logprob_i					)
            offset += logprob_i.size(0					)

      return out


    def      lowerCAmelCase_  (					self	:  Union[str, Any]   ,				__lowerCAmelCase	:  List[Any]					):
      if self.n_clusters == 0:
        _UpperCAmelCase						 =      self._compute_logit(__lowerCAmelCase   ,				self.out_layers[0].weight   ,				self.out_layers[0].bias   ,				self.out_projs[0]					)
        return nn.functional.log_softmax(__lowerCAmelCase   ,				dim=-1					)
      else:
        # construct weights and biases
        _UpperCAmelCase ,				_UpperCAmelCase						 =      [], []
        for i in range(len(self.cutoffs					)					):
          if self.div_val == 1:
            _UpperCAmelCase ,				_UpperCAmelCase						 =      self.cutoff_ends[i], self.cutoff_ends[i + 1]
            _UpperCAmelCase						 =      self.out_layers[0].weight[l_idx:r_idx]
            _UpperCAmelCase						 =      self.out_layers[0].bias[l_idx:r_idx]
          else:
            _UpperCAmelCase						 =      self.out_layers[i].weight
            _UpperCAmelCase						 =      self.out_layers[i].bias

          if i == 0:
            _UpperCAmelCase						 =      torch.cat([weight_i, self.cluster_weight]   ,				dim=0					)
            _UpperCAmelCase						 =      torch.cat([bias_i, self.cluster_bias]   ,				dim=0					)

          weights.append(__lowerCAmelCase					)
          biases.append(__lowerCAmelCase					)

        _UpperCAmelCase ,				_UpperCAmelCase ,				_UpperCAmelCase						 =      weights[0], biases[0], self.out_projs[0]
        _UpperCAmelCase						 =      self._compute_logit(__lowerCAmelCase   ,				__lowerCAmelCase   ,				__lowerCAmelCase   ,				__lowerCAmelCase					)

        _UpperCAmelCase						 =      hidden.new_empty((head_logit.size(0					), self.n_token)					)
        _UpperCAmelCase						 =      nn.functional.log_softmax(__lowerCAmelCase   ,				dim=1					)

        _UpperCAmelCase						 =      [0] + self.cutoffs
        for i in range(len(__lowerCAmelCase					) - 1					):
          _UpperCAmelCase ,				_UpperCAmelCase						 =      cutoff_values[i], cutoff_values[i + 1]

          if i == 0:
            _UpperCAmelCase						 =      head_logprob[:, : self.cutoffs[0]]
          else:
            _UpperCAmelCase ,				_UpperCAmelCase ,				_UpperCAmelCase						 =      weights[i], biases[i], self.out_projs[i]

            _UpperCAmelCase						 =      self._compute_logit(__lowerCAmelCase   ,				__lowerCAmelCase   ,				__lowerCAmelCase   ,				__lowerCAmelCase					)
            _UpperCAmelCase						 =      nn.functional.log_softmax(__lowerCAmelCase   ,				dim=1					)

            _UpperCAmelCase						 =      head_logprob[:, -i] + tail_logprob_i
            _UpperCAmelCase						 =      logprob_i

        return out
