


import itertools
from dataclasses import dataclass
from typing import Optional

import pyarrow as pa

import datasets
from datasets.table import table_cast


lowerCAmelCase__       :	List[str]     				=datasets.utils.logging.get_logger(__name__)







@dataclass
class      __lowercase  (datasets.BuilderConfig   ):





							"""simple docstring"""







							_UpperCAmelCase     =		None







class      __lowercase  (datasets.ArrowBasedBuilder   ):





							"""simple docstring"""







							_UpperCAmelCase     =		ArrowConfig






							def  UpperCamelCase__			(	self    ):





										"""simple docstring"""




										return datasets.DatasetInfo(features=self.config.features    )






							def  UpperCamelCase__			(	self       ,		lowerCAmelCase__    ):





										"""simple docstring"""




										if not self.config.data_files:
													raise ValueError(F'''At least one data file must be specified, but got data_files={self.config.data_files}'''    )
										SCREAMING_SNAKE_CASE_						:			Dict       						= dl_manager.download_and_extract(self.config.data_files    )
										if isinstance(lowerCAmelCase__       ,		(str, list, tuple)    ):
													SCREAMING_SNAKE_CASE_						:			Union[str, Any]       						= data_files
													if isinstance(lowerCAmelCase__       ,		lowerCAmelCase__    ):
																SCREAMING_SNAKE_CASE_						:			Optional[int]       						= [files]
													# Use `dl_manager.iter_files` to skip hidden files in an extracted archive
													SCREAMING_SNAKE_CASE_						:			str       						= [dl_manager.iter_files(lowerCAmelCase__    ) for file in files]
													return [datasets.SplitGenerator(name=datasets.Split.TRAIN       ,		gen_kwargs={'files': files}    )]
										SCREAMING_SNAKE_CASE_						:			Union[str, Any]       						= []
										for split_name, files in data_files.items():
													if isinstance(lowerCAmelCase__       ,		lowerCAmelCase__    ):
																SCREAMING_SNAKE_CASE_						:			Optional[Any]       						= [files]
													# Use `dl_manager.iter_files` to skip hidden files in an extracted archive
													SCREAMING_SNAKE_CASE_						:			Optional[int]       						= [dl_manager.iter_files(lowerCAmelCase__    ) for file in files]
													# Infer features is they are stoed in the arrow schema
													if self.info.features is None:
																for file in itertools.chain.from_iterable(lowerCAmelCase__    ):
																			with open(lowerCAmelCase__       ,		'rb'    ) as f:
																						SCREAMING_SNAKE_CASE_						:			Union[str, Any]       						= datasets.Features.from_arrow_schema(pa.ipc.open_stream(lowerCAmelCase__    ).schema    )
																			break
													splits.append(datasets.SplitGenerator(name=lowerCAmelCase__       ,		gen_kwargs={'files': files}    )    )
										return splits






							def  UpperCamelCase__			(	self       ,		lowerCAmelCase__    ):





										"""simple docstring"""




										if self.info.features is not None:
													# more expensive cast to support nested features with keys in a different order
													# allows str <-> int/float or str to Audio for example
													SCREAMING_SNAKE_CASE_						:			Any       						= table_cast(lowerCAmelCase__       ,		self.info.features.arrow_schema    )
										return pa_table






							def  UpperCamelCase__			(	self       ,		lowerCAmelCase__    ):





										"""simple docstring"""




										for file_idx, file in enumerate(itertools.chain.from_iterable(lowerCAmelCase__    )    ):
													with open(lowerCAmelCase__       ,		'rb'    ) as f:
																try:
																			for batch_idx, record_batch in enumerate(pa.ipc.open_stream(lowerCAmelCase__    )    ):
																						SCREAMING_SNAKE_CASE_						:			int       						= pa.Table.from_batches([record_batch]    )
																						# Uncomment for debugging (will print the Arrow table size and elements)
																						# logger.warning(f"pa_table: {pa_table} num rows: {pa_table.num_rows}")
																						# logger.warning('\n'.join(str(pa_table.slice(i, 1).to_pydict()) for i in range(pa_table.num_rows)))
																						yield F'''{file_idx}_{batch_idx}''', self._cast_table(lowerCAmelCase__    )
																except ValueError as e:
																			logger.error(F'''Failed to read file \'{file}\' with error {type(lowerCAmelCase__    )}: {e}'''    )
																			raise








import copy

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


UpperCamelCase__     :					Optional[Any]								=			logging.get_logger(__name__)

UpperCamelCase__     :					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 					__snake_case     (			lowerCAmelCase__				):
  __lowerCAmelCase   :		Any         =     'musicgen_decoder'
  __lowerCAmelCase   :		Optional[Any]         =     ['past_key_values']

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






class 					__snake_case     (			lowerCAmelCase__				):
  __lowerCAmelCase   :		int         =     'musicgen'
  __lowerCAmelCase   :		Optional[Any]         =     True

  def __init__(    self  ,						**_A):
      super().__init__(**_A)
      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')

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

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

      SCREAMING_SNAKE_CASE_													=  kwargs.pop('decoder')

      SCREAMING_SNAKE_CASE_													=  AutoConfig.for_model(_A  ,						**_A)
      SCREAMING_SNAKE_CASE_													=  AutoConfig.for_model(_A  ,						**_A)
      SCREAMING_SNAKE_CASE_													=  MusicgenDecoderConfig(**_A)
      SCREAMING_SNAKE_CASE_													=  True

  @classmethod
  def   lowerCAmelCase__   (    cls  ,						_A  ,						_A  ,						_A  ,						**_A  ,						):
      return cls(
          text_encoder=text_encoder_config.to_dict()  ,						audio_encoder=audio_encoder_config.to_dict()  ,						decoder=decoder_config.to_dict()  ,						**_A  ,						)


  def   lowerCAmelCase__   (    self):
      SCREAMING_SNAKE_CASE_													=  copy.deepcopy(self.__dict__)
      SCREAMING_SNAKE_CASE_													=  self.text_encoder.to_dict()
      SCREAMING_SNAKE_CASE_													=  self.audio_encoder.to_dict()
      SCREAMING_SNAKE_CASE_													=  self.decoder.to_dict()
      SCREAMING_SNAKE_CASE_													=  self.__class__.model_type
      return output


