'''simple docstring'''






import logging

import torch
from torch import nn
from torch.nn import CrossEntropyLoss, MSELoss

from transformers.file_utils import add_start_docstrings, add_start_docstrings_to_model_forward
from transformers.models.bert.modeling_bert import (
    BERT_INPUTS_DOCSTRING,
    BERT_START_DOCSTRING,
    BertEncoder,
    BertModel,
    BertPreTrainedModel,
)


_UpperCAmelCase    :     Dict         =       logging.getLogger(__name__)







class lowercase_       (		_UpperCamelCase			):







   """simple docstring"""

   def 	__UpperCAmelCase   (      self      :			List[str], UpperCamelCase__      :			str, UpperCamelCase__      :			List[Any], UpperCamelCase__      :			str=None, UpperCamelCase__      :			Tuple=None     )  ->			List[str]:
        _A    		=  self.layer[current_layer](UpperCamelCase__, UpperCamelCase__, head_mask[current_layer]     )

        _A    		=  layer_outputs[0]

        return hidden_states







@add_start_docstrings(
    "The bare Bert Model transformer with PABEE outputting raw hidden-states without any specific head on top."    ,    _UpperCamelCase    ,    )
class lowercase_       (		_UpperCamelCase			):







   """simple docstring"""

   def __init__(      self      :			Tuple, UpperCamelCase__      :			Tuple     )  ->			List[Any]:
        super().__init__(UpperCamelCase__     )

        _A    		=  BertEncoderWithPabee(UpperCamelCase__     )

        self.init_weights()
        _A    		=  0
        _A    		=  0
        _A    		=  0

        _A    		=  0
   def 	__UpperCAmelCase   (      self      :			Union[str, Any], UpperCamelCase__      :			Optional[int]     )  ->			Dict:
        _A    		=  threshold
   def 	__UpperCAmelCase   (      self      :			int, UpperCamelCase__      :			List[Any]     )  ->			List[str]:
        _A    		=  patience
   def 	__UpperCAmelCase   (      self      :			Optional[Any]     )  ->			List[str]:
        _A    		=  0
        _A    		=  0
   def 	__UpperCAmelCase   (      self      :			List[Any]     )  ->			Any:
        _A    		=  self.inference_layers_num / self.inference_instances_num
        _A    		=  (
            f'*** Patience = {self.patience} Avg. Inference Layers = {avg_inf_layers:.2f} Speed Up ='
            f' {1 - avg_inf_layers / self.config.num_hidden_layers:.2f} ***'
        )
        print(UpperCamelCase__     )

   @add_start_docstrings_to_model_forward(UpperCamelCase__     )
   def 	__UpperCAmelCase   (      self      :			List[Any], UpperCamelCase__      :			Optional[Any]=None, UpperCamelCase__      :			Any=None, UpperCamelCase__      :			List[str]=None, UpperCamelCase__      :			str=None, UpperCamelCase__      :			Dict=None, UpperCamelCase__      :			Dict=None, UpperCamelCase__      :			Dict=None, UpperCamelCase__      :			Any=None, UpperCamelCase__      :			Union[str, Any]=None, UpperCamelCase__      :			List[Any]=None, UpperCamelCase__      :			Any=False, )  ->			str:

        if input_ids is not None and inputs_embeds is not None:
             raise ValueError('You cannot specify both input_ids and inputs_embeds at the same time'     )
        elif input_ids is not None:
             _A    		=  input_ids.size()
        elif inputs_embeds is not None:
             _A    		=  inputs_embeds.size()[:-1]
        else:
             raise ValueError('You have to specify either input_ids or inputs_embeds'     )

        _A    		=  input_ids.device if input_ids is not None else inputs_embeds.device

        if attention_mask is None:
             _A    		=  torch.ones(UpperCamelCase__, device=UpperCamelCase__     )
        if token_type_ids is None:
             _A    		=  torch.zeros(UpperCamelCase__, dtype=torch.long, device=UpperCamelCase__     )

        # We can provide a self-attention mask of dimensions [batch_size, from_seq_length, to_seq_length]
        # ourselves in which case we just need to make it broadcastable to all heads.
        _A    		=  self.get_extended_attention_mask(UpperCamelCase__, UpperCamelCase__, UpperCamelCase__     )

        # If a 2D ou 3D attention mask is provided for the cross-attention
        # we need to make broadcastable to [batch_size, num_heads, seq_length, seq_length]
        if self.config.is_decoder and encoder_hidden_states is not None:
             _A		,       _A		,       _A    		=  encoder_hidden_states.size()
             _A    		=  (encoder_batch_size, encoder_sequence_length)
             if encoder_attention_mask is None:
                  _A    		=  torch.ones(UpperCamelCase__, device=UpperCamelCase__     )
             _A    		=  self.invert_attention_mask(UpperCamelCase__     )
        else:
             _A    		=  None

        # Prepare head mask if needed
        # 1.0 in head_mask indicate we keep the head
        # attention_probs has shape bsz x n_heads x N x N
        # input head_mask has shape [num_heads] or [num_hidden_layers x num_heads]
        # and head_mask is converted to shape [num_hidden_layers x batch x num_heads x seq_length x seq_length]
        _A    		=  self.get_head_mask(UpperCamelCase__, self.config.num_hidden_layers     )

        _A    		=  self.embeddings(
            input_ids=UpperCamelCase__, position_ids=UpperCamelCase__, token_type_ids=UpperCamelCase__, inputs_embeds=UpperCamelCase__     )
        _A    		=  embedding_output

        if self.training:
             _A    		=  []
             for i in range(self.config.num_hidden_layers     ):
                  _A    		=  self.encoder.adaptive_forward(
                      UpperCamelCase__, current_layer=UpperCamelCase__, attention_mask=UpperCamelCase__, head_mask=UpperCamelCase__     )

                  _A    		=  self.pooler(UpperCamelCase__     )
                  _A    		=  output_layers[i](output_dropout(UpperCamelCase__     )     )
                  res.append(UpperCamelCase__     )
        elif self.patience == 0:  # Use all layers for inference
             _A    		=  self.encoder(
                 UpperCamelCase__, attention_mask=UpperCamelCase__, head_mask=UpperCamelCase__, encoder_hidden_states=UpperCamelCase__, encoder_attention_mask=UpperCamelCase__, )
             _A    		=  self.pooler(encoder_outputs[0]     )
             _A    		=  [output_layers[self.config.num_hidden_layers - 1](UpperCamelCase__     )]
        else:
             _A    		=  0
             _A    		=  None
             _A    		=  0
             for i in range(self.config.num_hidden_layers     ):
                  calculated_layer_num += 1
                  _A    		=  self.encoder.adaptive_forward(
                      UpperCamelCase__, current_layer=UpperCamelCase__, attention_mask=UpperCamelCase__, head_mask=UpperCamelCase__     )

                  _A    		=  self.pooler(UpperCamelCase__     )
                  _A    		=  output_layers[i](UpperCamelCase__     )
                  if regression:
                       _A    		=  logits.detach()
                       if patient_result is not None:
                            _A    		=  patient_result.detach()
                       if (patient_result is not None) and torch.abs(patient_result - labels     ) < self.regression_threshold:
                            patient_counter += 1
                       else:
                            _A    		=  0
                  else:
                       _A    		=  logits.detach().argmax(dim=1     )
                       if patient_result is not None:
                            _A    		=  patient_result.detach().argmax(dim=1     )
                       if (patient_result is not None) and torch.all(labels.eq(UpperCamelCase__     )     ):
                            patient_counter += 1
                       else:
                            _A    		=  0

                  _A    		=  logits
                  if patient_counter == self.patience:
                       break
             _A    		=  [patient_result]
             self.inference_layers_num += calculated_layer_num
             self.inference_instances_num += 1

        return res







@add_start_docstrings(
    "Bert Model transformer with PABEE and a sequence classification/regression head on top (a linear layer on top of\n    the pooled output) e.g. for GLUE tasks. "    ,    _UpperCamelCase    ,    )
class lowercase_       (		_UpperCamelCase			):







   """simple docstring"""

   def __init__(      self      :			Dict, UpperCamelCase__      :			List[str]     )  ->			Any:
        super().__init__(UpperCamelCase__     )
        _A    		=  config.num_labels

        _A    		=  BertModelWithPabee(UpperCamelCase__     )
        _A    		=  nn.Dropout(config.hidden_dropout_prob     )
        _A    		=  nn.ModuleList(
            [nn.Linear(config.hidden_size, self.config.num_labels     ) for _ in range(config.num_hidden_layers     )]     )

        self.init_weights()

   @add_start_docstrings_to_model_forward(UpperCamelCase__     )
   def 	__UpperCAmelCase   (      self      :			Dict, UpperCamelCase__      :			int=None, UpperCamelCase__      :			str=None, UpperCamelCase__      :			Any=None, UpperCamelCase__      :			List[Any]=None, UpperCamelCase__      :			Union[str, Any]=None, UpperCamelCase__      :			Tuple=None, UpperCamelCase__      :			List[str]=None, )  ->			str:
        _A    		=  self.bert(
            input_ids=UpperCamelCase__, attention_mask=UpperCamelCase__, token_type_ids=UpperCamelCase__, position_ids=UpperCamelCase__, head_mask=UpperCamelCase__, inputs_embeds=UpperCamelCase__, output_dropout=self.dropout, output_layers=self.classifiers, regression=self.num_labels == 1, )

        _A    		=  (logits[-1],)

        if labels is not None:
             _A    		=  None
             _A    		=  0
             for ix, logits_item in enumerate(UpperCamelCase__     ):
                  if self.num_labels == 1:
                       #  We are doing regression
                       _A    		=  MSELoss()
                       _A    		=  loss_fct(logits_item.view(-1     ), labels.view(-1     )     )
                  else:
                       _A    		=  CrossEntropyLoss()
                       _A    		=  loss_fct(logits_item.view(-1, self.num_labels     ), labels.view(-1     )     )
                  if total_loss is None:
                       _A    		=  loss
                  else:
                       total_loss += loss * (ix + 1)
                  total_weights += ix + 1
             _A    		=  (total_loss / total_weights,) + outputs

        return outputs











import argparse
from pathlib import Path
from typing import Dict, OrderedDict, Tuple

import torch
from audiocraft.models import MusicGen

from transformers import (
    AutoFeatureExtractor,
    AutoTokenizer,
    EncodecModel,
    MusicgenDecoderConfig,
    MusicgenForConditionalGeneration,
    MusicgenProcessor,
    TaEncoderModel,
)
from transformers.models.musicgen.modeling_musicgen import MusicgenForCausalLM
from transformers.utils import logging


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


UpperCamelCase    							=      ['model.decoder.embed_positions.weights']


def 				__lowerCamelCase     (							__lowerCAmelCase							:					Optional[int]							)     ->  Union[str, Any]:
   if "emb" in name:
      __UpperCamelCase	:			Any					=					name.replace("""emb"""					,							"""model.decoder.embed_tokens"""							)
   if "transformer" in name:
      __UpperCamelCase	:			str					=					name.replace("""transformer"""					,							"""model.decoder"""							)
   if "cross_attention" in name:
      __UpperCamelCase	:			List[str]					=					name.replace("""cross_attention"""					,							"""encoder_attn"""							)
   if "linear1" in name:
      __UpperCamelCase	:			Optional[int]					=					name.replace("""linear1"""					,							"""fc1"""							)
   if "linear2" in name:
      __UpperCamelCase	:			Optional[Any]					=					name.replace("""linear2"""					,							"""fc2"""							)
   if "norm1" in name:
      __UpperCamelCase	:			Tuple					=					name.replace("""norm1"""					,							"""self_attn_layer_norm"""							)
   if "norm_cross" in name:
      __UpperCamelCase	:			List[str]					=					name.replace("""norm_cross"""					,							"""encoder_attn_layer_norm"""							)
   if "norm2" in name:
      __UpperCamelCase	:			Tuple					=					name.replace("""norm2"""					,							"""final_layer_norm"""							)
   if "out_norm" in name:
      __UpperCamelCase	:			Union[str, Any]					=					name.replace("""out_norm"""					,							"""model.decoder.layer_norm"""							)
   if "linears" in name:
      __UpperCamelCase	:			Optional[int]					=					name.replace("""linears"""					,							"""lm_heads"""							)
   if "condition_provider.conditioners.description.output_proj" in name:
      __UpperCamelCase	:			Tuple					=					name.replace("""condition_provider.conditioners.description.output_proj"""					,							"""enc_to_dec_proj"""							)
   return name


def 				__lowerCamelCase     (							__lowerCAmelCase							:					OrderedDict					,							__lowerCAmelCase							:					int							)     ->  Tuple[Dict, Dict]:
   __UpperCamelCase	:			Tuple					=					list(state_dict.keys()							)
   __UpperCamelCase	:			List[Any]					=					{}
   for key in keys:
      __UpperCamelCase	:			Optional[Any]					=					state_dict.pop(__lowerCAmelCase							)
      __UpperCamelCase	:			Dict					=					rename_keys(__lowerCAmelCase							)
      if "in_proj_weight" in key:
         # split fused qkv proj
         __UpperCamelCase	:			Optional[Any]					=					val[:hidden_size, :]
         __UpperCamelCase	:			Optional[Any]					=					val[hidden_size : 2 * hidden_size, :]
         __UpperCamelCase	:			List[str]					=					val[-hidden_size:, :]
      elif "enc_to_dec_proj" in key:
         __UpperCamelCase	:			Dict					=					val
      else:
         __UpperCamelCase	:			Any					=					val
   return state_dict, enc_dec_proj_state_dict


def 				__lowerCamelCase     (							__lowerCAmelCase							:					str							)     ->  MusicgenDecoderConfig:
   if checkpoint == "small":
      # default config values
      __UpperCamelCase	:			int					=					1024
      __UpperCamelCase	:			Union[str, Any]					=					24
      __UpperCamelCase	:			int					=					16
   elif checkpoint == "medium":
      __UpperCamelCase	:			List[Any]					=					1536
      __UpperCamelCase	:			Dict					=					48
      __UpperCamelCase	:			Dict					=					24
   elif checkpoint == "large":
      __UpperCamelCase	:			List[Any]					=					2048
      __UpperCamelCase	:			str					=					48
      __UpperCamelCase	:			Optional[int]					=					32
   else:
      raise ValueError(f'Checkpoint should be one of `[\'small\', \'medium\', \'large\']`, got {checkpoint}.'							)
   __UpperCamelCase	:			Any					=					MusicgenDecoderConfig(
       hidden_size=__lowerCAmelCase					,							ffn_dim=hidden_size * 4					,							num_hidden_layers=__lowerCAmelCase					,							num_attention_heads=__lowerCAmelCase					,							)
   return config


@torch.no_grad()
def 				__lowerCamelCase     (							__lowerCAmelCase							:					List[str]					,							__lowerCAmelCase							:					Optional[Any]=None					,							__lowerCAmelCase							:					Tuple=None					,							__lowerCAmelCase							:					str="cpu"							)     ->  Optional[int]:
   __UpperCamelCase	:			str					=					MusicGen.get_pretrained(__lowerCAmelCase					,							device=__lowerCAmelCase							)
   __UpperCamelCase	:			int					=					decoder_config_from_checkpoint(__lowerCAmelCase							)

   __UpperCamelCase	:			Optional[int]					=					fairseq_model.lm.state_dict()
   __UpperCamelCase	,     __UpperCamelCase	:			Union[str, Any]					=					rename_state_dict(
       __lowerCAmelCase					,							hidden_size=decoder_config.hidden_size							)

   __UpperCamelCase	:			List[Any]					=					TaEncoderModel.from_pretrained("""t5-base"""							)
   __UpperCamelCase	:			Dict					=					EncodecModel.from_pretrained("""facebook/encodec_32khz"""							)
   __UpperCamelCase	:			List[str]					=					MusicgenForCausalLM(__lowerCAmelCase							).eval()

   # load all decoder weights - expect that we'll be missing embeddings and enc-dec projection
   __UpperCamelCase	,     __UpperCamelCase	:			Tuple					=					decoder.load_state_dict(__lowerCAmelCase					,							strict=__lowerCAmelCase							)

   for key in missing_keys.copy():
      if key.startswith(("""text_encoder""", """audio_encoder""")							) or key in EXPECTED_MISSING_KEYS:
         missing_keys.remove(__lowerCAmelCase							)

   if len(__lowerCAmelCase							) > 0:
      raise ValueError(f'Missing key(s) in state_dict: {missing_keys}'							)

   if len(__lowerCAmelCase							) > 0:
      raise ValueError(f'Unexpected key(s) in state_dict: {unexpected_keys}'							)

   # init the composite model
   __UpperCamelCase	:			Tuple					=					MusicgenForConditionalGeneration(text_encoder=__lowerCAmelCase					,							audio_encoder=__lowerCAmelCase					,							decoder=__lowerCAmelCase							)

   # load the pre-trained enc-dec projection (from the decoder state dict)
   model.enc_to_dec_proj.load_state_dict(__lowerCAmelCase							)

   # check we can do a forward pass
   __UpperCamelCase	:			int					=					torch.arange(0					,							8					,							dtype=torch.long							).reshape(2					,							-1							)
   __UpperCamelCase	:			List[Any]					=					input_ids.reshape(2 * 4					,							-1							)

   with torch.no_grad():
      __UpperCamelCase	:			int					=					model(input_ids=__lowerCAmelCase					,							decoder_input_ids=__lowerCAmelCase							).logits

   if logits.shape != (8, 1, 2048):
      raise ValueError("""Incorrect shape for logits"""							)

   # now construct the processor
   __UpperCamelCase	:			Union[str, Any]					=					AutoTokenizer.from_pretrained("""t5-base"""							)
   __UpperCamelCase	:			List[Any]					=					AutoFeatureExtractor.from_pretrained("""facebook/encodec_32khz"""					,							padding_side="""left"""							)

   __UpperCamelCase	:			Any					=					MusicgenProcessor(feature_extractor=__lowerCAmelCase					,							tokenizer=__lowerCAmelCase							)

   # set the appropriate bos/pad token ids
   __UpperCamelCase	:			Tuple					=					2048
   __UpperCamelCase	:			int					=					2048

   # set other default generation config params
   __UpperCamelCase	:			str					=					int(30 * audio_encoder.config.frame_rate							)
   __UpperCamelCase	:			List[str]					=					True
   __UpperCamelCase	:			Tuple					=					3.0

   if pytorch_dump_folder is not None:
      Path(__lowerCAmelCase							).mkdir(exist_ok=__lowerCAmelCase							)
      logger.info(f'Saving model {checkpoint} to {pytorch_dump_folder}'							)
      model.save_pretrained(__lowerCAmelCase							)
      processor.save_pretrained(__lowerCAmelCase							)

   if repo_id:
      logger.info(f'Pushing model {checkpoint} to {repo_id}'							)
      model.push_to_hub(__lowerCAmelCase							)
      processor.push_to_hub(__lowerCAmelCase							)


if __name__ == "__main__":
       UpperCamelCase    							=      argparse.ArgumentParser()
       # Required parameters
       parser.add_argument(
           '--checkpoint',
           default='small',
           type=str,
           help='Checkpoint size of the MusicGen model you\'d like to convert. Can be one of: `[\'small\', \'medium\', \'large\']`.',
       )
       parser.add_argument(
           '--pytorch_dump_folder',
           required=True,
           default=None,
           type=str,
           help='Path to the output PyTorch model directory.',
       )
       parser.add_argument(
           '--push_to_hub', default=None, type=str, help='Where to upload the converted model on the 🤗 hub.'
       )
       parser.add_argument(
           '--device', default='cpu', type=str, help='Torch device to run the conversion, either cpu or cuda.'
       )

       UpperCamelCase    							=      parser.parse_args()
       convert_musicgen_checkpoint(args.checkpoint, args.pytorch_dump_folder, args.push_to_hub)



