'''simple docstring'''




import copy

from ...configuration_utils import PretrainedConfig
from ...utils import logging
from ..auto.configuration_auto import AutoConfig


__snake_case  : Optional[int]		       =    logging.get_logger(__name__)

__snake_case  : List[Any]		       =    {
    'facebook/musicgen-small': 'https://huggingface.co/facebook/musicgen-small/resolve/main/config.json',
    # See all Musicgen models at https://huggingface.co/models?filter=musicgen
}




class       __UpperCAmelCase							(							_UpperCamelCase	):





       '''simple docstring'''



       __lowercase   :   List[Any]  					=							'musicgen_decoder'
       __lowercase   :   str  					=							['past_key_values']

       def __init__( self       ,	_SCREAMING_SNAKE_CASE=2048       ,	_SCREAMING_SNAKE_CASE=2048       ,	_SCREAMING_SNAKE_CASE=24       ,	_SCREAMING_SNAKE_CASE=4096       ,	_SCREAMING_SNAKE_CASE=16       ,	_SCREAMING_SNAKE_CASE=0.0       ,	_SCREAMING_SNAKE_CASE=True       ,	_SCREAMING_SNAKE_CASE="gelu"       ,	_SCREAMING_SNAKE_CASE=1024       ,	_SCREAMING_SNAKE_CASE=0.1       ,	_SCREAMING_SNAKE_CASE=0.0       ,	_SCREAMING_SNAKE_CASE=0.0       ,	_SCREAMING_SNAKE_CASE=0.02       ,	_SCREAMING_SNAKE_CASE=False       ,	_SCREAMING_SNAKE_CASE=4       ,	_SCREAMING_SNAKE_CASE=2048       ,	_SCREAMING_SNAKE_CASE=2048       ,	_SCREAMING_SNAKE_CASE=None       ,	_SCREAMING_SNAKE_CASE=False       ,	**_SCREAMING_SNAKE_CASE       ,	)						->	Tuple:
           A_      =  vocab_size
           A_      =  max_position_embeddings
           A_      =  hidden_size
           A_      =  ffn_dim
           A_      =  num_hidden_layers
           A_      =  num_attention_heads
           A_      =  dropout
           A_      =  attention_dropout
           A_      =  activation_dropout
           A_      =  activation_function
           A_      =  initializer_factor
           A_      =  layerdrop
           A_      =  use_cache
           A_      =  scale_embedding  # scale factor will be sqrt(d_model) if True
           A_      =  num_codebooks
           super().__init__(
               pad_token_id=_SCREAMING_SNAKE_CASE       ,	bos_token_id=_SCREAMING_SNAKE_CASE       ,	eos_token_id=_SCREAMING_SNAKE_CASE       ,	tie_word_embeddings=_SCREAMING_SNAKE_CASE       ,	**_SCREAMING_SNAKE_CASE       ,	)




class       __UpperCAmelCase							(							_UpperCamelCase	):





       '''simple docstring'''



       __lowercase   :   str  					=							'musicgen'
       __lowercase   :   Union[str, Any]  					=							True

       def __init__( self       ,	**_SCREAMING_SNAKE_CASE					)						->	List[Any]:
           super().__init__(**_SCREAMING_SNAKE_CASE					)
           if "text_encoder" not in kwargs or "audio_encoder" not in kwargs or "decoder" not in kwargs:
               raise ValueError('''Config has to be initialized with text_encoder, audio_encoder and decoder config'''					)

           A_      =  kwargs.pop('''text_encoder'''					)
           A_      =  text_encoder_config.pop('''model_type'''					)

           A_      =  kwargs.pop('''audio_encoder'''					)
           A_      =  audio_encoder_config.pop('''model_type'''					)

           A_      =  kwargs.pop('''decoder'''					)

           A_      =  AutoConfig.for_model(_SCREAMING_SNAKE_CASE       ,	**_SCREAMING_SNAKE_CASE					)
           A_      =  AutoConfig.for_model(_SCREAMING_SNAKE_CASE       ,	**_SCREAMING_SNAKE_CASE					)
           A_      =  MusicgenDecoderConfig(**_SCREAMING_SNAKE_CASE					)
           A_      =  True

       @classmethod
       def    __A			( cls       ,	_SCREAMING_SNAKE_CASE       ,	_SCREAMING_SNAKE_CASE       ,	_SCREAMING_SNAKE_CASE       ,	**_SCREAMING_SNAKE_CASE       ,	)						->	Any:
           return cls(
               text_encoder=text_encoder_config.to_dict()       ,	audio_encoder=audio_encoder_config.to_dict()       ,	decoder=decoder_config.to_dict()       ,	**_SCREAMING_SNAKE_CASE       ,	)

       def    __A			( self					)						->	Tuple:
           A_      =  copy.deepcopy(self.__dict__					)
           A_      =  self.text_encoder.to_dict()
           A_      =  self.audio_encoder.to_dict()
           A_      =  self.decoder.to_dict()
           A_      =  self.__class__.model_type
           return output







'''simple docstring'''




import json
import os
import tempfile

import datasets
from utils import generate_example_dataset, get_duration


__snake_case  : Dict		       =    500_000

__snake_case       ,       __snake_case  : List[str]		       =    os.path.split(__file__)
__snake_case  : Optional[Any]		       =    os.path.join(RESULTS_BASEPATH, 'results', RESULTS_FILENAME.replace('.py', '.json'))





@get_duration
def 	_UpperCAmelCase       (				_UpperCamelCase       :      datasets.Dataset  )      -> Union[str, Any]:
    A_      =  dataset.select(range(0,  len(_UpperCamelCase  ),  2  )  )





@get_duration
def 	_UpperCAmelCase       (				_UpperCamelCase       :      datasets.Dataset  )      -> str:
    A_      =  dataset.sort('''numbers'''  )





@get_duration
def 	_UpperCAmelCase       (				_UpperCamelCase       :      datasets.Dataset  )      -> str:
    A_      =  dataset.shuffle()





@get_duration
def 	_UpperCAmelCase       (				_UpperCamelCase       :      datasets.Dataset  )      -> Optional[Any]:
    A_      =  dataset.train_test_split(0.1  )





@get_duration
def 	_UpperCAmelCase       (				_UpperCamelCase       :      datasets.Dataset,  _UpperCamelCase       :      Optional[int]=10  )      -> Dict:
    for shard_id in range(_UpperCamelCase  ):
        A_      =  dataset.shard(_UpperCamelCase,  _UpperCamelCase  )





def 	_UpperCAmelCase       (				)      -> Any:
    A_      =  {'''num examples''': SPEED_TEST_N_EXAMPLES}
    A_      =  (select, sort, shuffle, train_test_split, shard)
    with tempfile.TemporaryDirectory() as tmp_dir:
        print('''generating dataset'''  )
        A_      =  datasets.Features({'''text''': datasets.Value('''string'''  ), '''numbers''': datasets.Value('''float32'''  )}  )
        A_      =  generate_example_dataset(
            os.path.join(_UpperCamelCase,  '''dataset.arrow'''  ),  _UpperCamelCase,  num_examples=_UpperCamelCase  )
        print('''Functions'''  )
        for func in functions:
            print(func.__name__  )
            A_      =  func(_UpperCamelCase  )

    with open(_UpperCamelCase,  '''wb'''  ) as f:
        f.write(json.dumps(_UpperCamelCase  ).encode('''utf-8'''  )  )


if __name__ == "__main__":  # useful to run the profiler
  benchmark_indices_mapping()






