






from __future__ import annotations

import copy
import tempfile
import unittest

from transformers import CONFIG_MAPPING, AutoConfig, BertConfig, GPTaConfig, TaConfig, TapasConfig, is_tf_available
from transformers.testing_utils import (
    DUMMY_UNKNOWN_IDENTIFIER,
    SMALL_MODEL_IDENTIFIER,
    RequestCounter,
    require_tensorflow_probability,
    require_tf,
    slow,
)

from ..bert.test_modeling_bert import BertModelTester


if is_tf_available():
       from transformers import (
           TFAutoModel,
           TFAutoModelForCausalLM,
           TFAutoModelForMaskedLM,
           TFAutoModelForPreTraining,
           TFAutoModelForQuestionAnswering,
           TFAutoModelForSeqaSeqLM,
           TFAutoModelForSequenceClassification,
           TFAutoModelForTableQuestionAnswering,
           TFAutoModelForTokenClassification,
           TFAutoModelWithLMHead,
           TFBertForMaskedLM,
           TFBertForPreTraining,
           TFBertForQuestionAnswering,
           TFBertForSequenceClassification,
           TFBertModel,
           TFFunnelBaseModel,
           TFFunnelModel,
           TFGPTaLMHeadModel,
           TFRobertaForMaskedLM,
           TFTaForConditionalGeneration,
           TFTapasForQuestionAnswering,
       )
       from transformers.models.auto.modeling_tf_auto import (
           TF_MODEL_FOR_CAUSAL_LM_MAPPING,
           TF_MODEL_FOR_MASKED_LM_MAPPING,
           TF_MODEL_FOR_PRETRAINING_MAPPING,
           TF_MODEL_FOR_QUESTION_ANSWERING_MAPPING,
           TF_MODEL_FOR_SEQUENCE_CLASSIFICATION_MAPPING,
           TF_MODEL_FOR_TOKEN_CLASSIFICATION_MAPPING,
           TF_MODEL_MAPPING,
       )
       from transformers.models.bert.modeling_tf_bert import TF_BERT_PRETRAINED_MODEL_ARCHIVE_LIST
       from transformers.models.gpta.modeling_tf_gpta import TF_GPT2_PRETRAINED_MODEL_ARCHIVE_LIST
       from transformers.models.ta.modeling_tf_ta import TF_T5_PRETRAINED_MODEL_ARCHIVE_LIST
       from transformers.models.tapas.modeling_tf_tapas import TF_TAPAS_PRETRAINED_MODEL_ARCHIVE_LIST







class 		_A						( UpperCAmelCase_ ):
  lowercase_						:			Optional[Any]						=					'''new-model'''


if is_tf_available():







       class 		_A						( UpperCAmelCase_ ):
         lowercase_						:			Dict						=					NewModelConfig







@require_tf
class 		_A						( unittest.TestCase ):

  @slow
  def    a					(    self						:	str			):
     """simple docstring"""

     __UpperCamelCase	:			List[str]					=					"""bert-base-cased"""
     __UpperCamelCase	:			Dict					=					AutoConfig.from_pretrained(lowerCamelCase__			)
     self.assertIsNotNone(lowerCamelCase__			)
     self.assertIsInstance(lowerCamelCase__   ,				lowerCamelCase__			)

     __UpperCamelCase	:			str					=					TFAutoModel.from_pretrained(lowerCamelCase__			)
     self.assertIsNotNone(lowerCamelCase__			)
     self.assertIsInstance(lowerCamelCase__   ,				lowerCamelCase__			)

  @slow
  def    a					(    self						:	Dict			):
     """simple docstring"""

     __UpperCamelCase	:			Optional[Any]					=					"""bert-base-cased"""
     __UpperCamelCase	:			Dict					=					AutoConfig.from_pretrained(lowerCamelCase__			)
     self.assertIsNotNone(lowerCamelCase__			)
     self.assertIsInstance(lowerCamelCase__   ,				lowerCamelCase__			)

     __UpperCamelCase	:			str					=					TFAutoModelForPreTraining.from_pretrained(lowerCamelCase__			)
     self.assertIsNotNone(lowerCamelCase__			)
     self.assertIsInstance(lowerCamelCase__   ,				lowerCamelCase__			)

  @slow
  def    a					(    self						:	List[Any]			):
     """simple docstring"""

     for model_name in TF_GPT2_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
        __UpperCamelCase	:			List[Any]					=					AutoConfig.from_pretrained(lowerCamelCase__			)
        self.assertIsNotNone(lowerCamelCase__			)
        self.assertIsInstance(lowerCamelCase__   ,				lowerCamelCase__			)

        __UpperCamelCase	:			Optional[Any]					=					TFAutoModelForCausalLM.from_pretrained(lowerCamelCase__			)
        __UpperCamelCase	,     __UpperCamelCase	:			List[Any]					=					TFAutoModelForCausalLM.from_pretrained(lowerCamelCase__   ,				output_loading_info=lowerCamelCase__			)
        self.assertIsNotNone(lowerCamelCase__			)
        self.assertIsInstance(lowerCamelCase__   ,				lowerCamelCase__			)

  @slow
  def    a					(    self						:	Optional[int]			):
     """simple docstring"""

     for model_name in TF_BERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
        __UpperCamelCase	:			Optional[int]					=					AutoConfig.from_pretrained(lowerCamelCase__			)
        self.assertIsNotNone(lowerCamelCase__			)
        self.assertIsInstance(lowerCamelCase__   ,				lowerCamelCase__			)

        __UpperCamelCase	:			int					=					TFAutoModelWithLMHead.from_pretrained(lowerCamelCase__			)
        self.assertIsNotNone(lowerCamelCase__			)
        self.assertIsInstance(lowerCamelCase__   ,				lowerCamelCase__			)

  @slow
  def    a					(    self						:	Optional[int]			):
     """simple docstring"""

     for model_name in TF_BERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
        __UpperCamelCase	:			List[Any]					=					AutoConfig.from_pretrained(lowerCamelCase__			)
        self.assertIsNotNone(lowerCamelCase__			)
        self.assertIsInstance(lowerCamelCase__   ,				lowerCamelCase__			)

        __UpperCamelCase	:			Any					=					TFAutoModelForMaskedLM.from_pretrained(lowerCamelCase__			)
        __UpperCamelCase	,     __UpperCamelCase	:			List[str]					=					TFAutoModelForMaskedLM.from_pretrained(lowerCamelCase__   ,				output_loading_info=lowerCamelCase__			)
        self.assertIsNotNone(lowerCamelCase__			)
        self.assertIsInstance(lowerCamelCase__   ,				lowerCamelCase__			)

  @slow
  def    a					(    self						:	Union[str, Any]			):
     """simple docstring"""

     for model_name in TF_T5_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
        __UpperCamelCase	:			Union[str, Any]					=					AutoConfig.from_pretrained(lowerCamelCase__			)
        self.assertIsNotNone(lowerCamelCase__			)
        self.assertIsInstance(lowerCamelCase__   ,				lowerCamelCase__			)

        __UpperCamelCase	:			List[Any]					=					TFAutoModelForSeqaSeqLM.from_pretrained(lowerCamelCase__			)
        __UpperCamelCase	,     __UpperCamelCase	:			Optional[Any]					=					TFAutoModelForSeqaSeqLM.from_pretrained(lowerCamelCase__   ,				output_loading_info=lowerCamelCase__			)
        self.assertIsNotNone(lowerCamelCase__			)
        self.assertIsInstance(lowerCamelCase__   ,				lowerCamelCase__			)

  @slow
  def    a					(    self						:	int			):
     """simple docstring"""

     for model_name in ["bert-base-uncased"]:
        __UpperCamelCase	:			List[Any]					=					AutoConfig.from_pretrained(lowerCamelCase__			)
        self.assertIsNotNone(lowerCamelCase__			)
        self.assertIsInstance(lowerCamelCase__   ,				lowerCamelCase__			)

        __UpperCamelCase	:			Dict					=					TFAutoModelForSequenceClassification.from_pretrained(lowerCamelCase__			)
        self.assertIsNotNone(lowerCamelCase__			)
        self.assertIsInstance(lowerCamelCase__   ,				lowerCamelCase__			)

  @slow
  def    a					(    self						:	Any			):
     """simple docstring"""

     for model_name in ["bert-base-uncased"]:
        __UpperCamelCase	:			Union[str, Any]					=					AutoConfig.from_pretrained(lowerCamelCase__			)
        self.assertIsNotNone(lowerCamelCase__			)
        self.assertIsInstance(lowerCamelCase__   ,				lowerCamelCase__			)

        __UpperCamelCase	:			List[Any]					=					TFAutoModelForQuestionAnswering.from_pretrained(lowerCamelCase__			)
        self.assertIsNotNone(lowerCamelCase__			)
        self.assertIsInstance(lowerCamelCase__   ,				lowerCamelCase__			)

  @slow
  @require_tensorflow_probability
  def    a					(    self						:	List[Any]			):
     """simple docstring"""

     for model_name in TF_TAPAS_PRETRAINED_MODEL_ARCHIVE_LIST[5:6]:
        __UpperCamelCase	:			Any					=					AutoConfig.from_pretrained(lowerCamelCase__			)
        self.assertIsNotNone(lowerCamelCase__			)
        self.assertIsInstance(lowerCamelCase__   ,				lowerCamelCase__			)

        __UpperCamelCase	:			Tuple					=					TFAutoModelForTableQuestionAnswering.from_pretrained(lowerCamelCase__			)
        __UpperCamelCase	,     __UpperCamelCase	:			Optional[int]					=					TFAutoModelForTableQuestionAnswering.from_pretrained(
            lowerCamelCase__   ,				output_loading_info=lowerCamelCase__			)
        self.assertIsNotNone(lowerCamelCase__			)
        self.assertIsInstance(lowerCamelCase__   ,				lowerCamelCase__			)

  def    a					(    self						:	Optional[int]			):
     """simple docstring"""

     __UpperCamelCase	:			Any					=					TFAutoModelWithLMHead.from_pretrained(lowerCamelCase__			)
     self.assertIsInstance(lowerCamelCase__   ,				lowerCamelCase__			)
     self.assertEqual(model.num_parameters()   ,				1_44_10			)
     self.assertEqual(model.num_parameters(only_trainable=lowerCamelCase__			)   ,				1_44_10			)

  def    a					(    self						:	int			):
     """simple docstring"""

     __UpperCamelCase	:			Dict					=					TFAutoModelWithLMHead.from_pretrained(lowerCamelCase__			)
     self.assertIsInstance(lowerCamelCase__   ,				lowerCamelCase__			)
     self.assertEqual(model.num_parameters()   ,				1_44_10			)
     self.assertEqual(model.num_parameters(only_trainable=lowerCamelCase__			)   ,				1_44_10			)

  def    a					(    self						:	Union[str, Any]			):
     """simple docstring"""

     __UpperCamelCase	:			str					=					TFAutoModel.from_pretrained("""sgugger/funnel-random-tiny"""			)
     self.assertIsInstance(lowerCamelCase__   ,				lowerCamelCase__			)

     __UpperCamelCase	:			Union[str, Any]					=					copy.deepcopy(model.config			)
     __UpperCamelCase	:			List[Any]					=					["""FunnelBaseModel"""]
     __UpperCamelCase	:			Optional[int]					=					TFAutoModel.from_config(lowerCamelCase__			)
     self.assertIsInstance(lowerCamelCase__   ,				lowerCamelCase__			)

     with tempfile.TemporaryDirectory() as tmp_dir:
        model.save_pretrained(lowerCamelCase__			)
        __UpperCamelCase	:			str					=					TFAutoModel.from_pretrained(lowerCamelCase__			)
        self.assertIsInstance(lowerCamelCase__   ,				lowerCamelCase__			)

  def    a					(    self						:	Dict			):
     """simple docstring"""

     try:
        AutoConfig.register("""new-model"""   ,				lowerCamelCase__			)

        __UpperCamelCase	:			List[str]					=					[
            TFAutoModel,
            TFAutoModelForCausalLM,
            TFAutoModelForMaskedLM,
            TFAutoModelForPreTraining,
            TFAutoModelForQuestionAnswering,
            TFAutoModelForSequenceClassification,
            TFAutoModelForTokenClassification,
        ]

        for auto_class in auto_classes:
           with self.subTest(auto_class.__name__			):
              # Wrong config class will raise an error
              with self.assertRaises(lowerCamelCase__			):
                 auto_class.register(lowerCamelCase__   ,				lowerCamelCase__			)
              auto_class.register(lowerCamelCase__   ,				lowerCamelCase__			)
              # Trying to register something existing in the Transformers library will raise an error
              with self.assertRaises(lowerCamelCase__			):
                 auto_class.register(lowerCamelCase__   ,				lowerCamelCase__			)

              # Now that the config is registered, it can be used as any other config with the auto-API
              __UpperCamelCase	:			Dict					=					BertModelTester(self			).get_config()
              __UpperCamelCase	:			Optional[Any]					=					NewModelConfig(**tiny_config.to_dict()			)
              __UpperCamelCase	:			Optional[int]					=					auto_class.from_config(lowerCamelCase__			)
              self.assertIsInstance(lowerCamelCase__   ,				lowerCamelCase__			)

              with tempfile.TemporaryDirectory() as tmp_dir:
                 model.save_pretrained(lowerCamelCase__			)
                 __UpperCamelCase	:			Union[str, Any]					=					auto_class.from_pretrained(lowerCamelCase__			)
                 self.assertIsInstance(lowerCamelCase__   ,				lowerCamelCase__			)

     finally:
        if "new-model" in CONFIG_MAPPING._extra_content:
           del CONFIG_MAPPING._extra_content["new-model"]
        for mapping in (
            TF_MODEL_MAPPING,
            TF_MODEL_FOR_PRETRAINING_MAPPING,
            TF_MODEL_FOR_QUESTION_ANSWERING_MAPPING,
            TF_MODEL_FOR_SEQUENCE_CLASSIFICATION_MAPPING,
            TF_MODEL_FOR_TOKEN_CLASSIFICATION_MAPPING,
            TF_MODEL_FOR_CAUSAL_LM_MAPPING,
            TF_MODEL_FOR_MASKED_LM_MAPPING,
        ):
           if NewModelConfig in mapping._extra_content:
              del mapping._extra_content[NewModelConfig]

  def    a					(    self						:	Tuple			):
     """simple docstring"""

     with self.assertRaisesRegex(
         lowerCamelCase__   ,				"""bert-base is not a local folder and is not a valid model identifier"""			):
        __UpperCamelCase	:			Optional[int]					=					TFAutoModel.from_pretrained("""bert-base"""			)

  def    a					(    self						:	List[Any]			):
     """simple docstring"""

     with self.assertRaisesRegex(
         lowerCamelCase__   ,				R"""aaaaaa is not a valid git identifier \(branch name, tag name or commit id\)"""			):
        __UpperCamelCase	:			Optional[Any]					=					TFAutoModel.from_pretrained(lowerCamelCase__   ,				revision="""aaaaaa"""			)

  def    a					(    self						:	Tuple			):
     """simple docstring"""

     with self.assertRaisesRegex(
         lowerCamelCase__   ,				"""hf-internal-testing/config-no-model does not appear to have a file named pytorch_model.bin"""   ,				):
        __UpperCamelCase	:			List[Any]					=					TFAutoModel.from_pretrained("""hf-internal-testing/config-no-model"""			)

  def    a					(    self						:	List[Any]			):
     """simple docstring"""

     with self.assertRaisesRegex(lowerCamelCase__   ,				"""Use `from_pt=True` to load this model"""			):
        __UpperCamelCase	:			Union[str, Any]					=					TFAutoModel.from_pretrained("""hf-internal-testing/tiny-bert-pt-only"""			)

  def    a					(    self						:	List[Any]			):
     """simple docstring"""

     __UpperCamelCase	:			str					=					TFAutoModel.from_pretrained("""hf-internal-testing/tiny-random-bert"""			)
     with RequestCounter() as counter:
        __UpperCamelCase	:			Optional[int]					=					TFAutoModel.from_pretrained("""hf-internal-testing/tiny-random-bert"""			)
        self.assertEqual(counter.get_request_count   ,				0			)
        self.assertEqual(counter.head_request_count   ,				1			)
        self.assertEqual(counter.other_request_count   ,				0			)

     # With a sharded checkpoint
     __UpperCamelCase	:			Union[str, Any]					=					TFAutoModel.from_pretrained("""ArthurZ/tiny-random-bert-sharded"""			)
     with RequestCounter() as counter:
        __UpperCamelCase	:			List[Any]					=					TFAutoModel.from_pretrained("""ArthurZ/tiny-random-bert-sharded"""			)
        self.assertEqual(counter.get_request_count   ,				0			)
        self.assertEqual(counter.head_request_count   ,				1			)
        self.assertEqual(counter.other_request_count   ,				0			)











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)



