




'''simple docstring'''




import sys
from collections.abc import Mapping
from typing import TYPE_CHECKING

import numpy as np
import pyarrow as pa

from .. import config
from ..utils.py_utils import map_nested
from .formatting import TensorFormatter


if TYPE_CHECKING:
  import tensorflow as tf




class 					a__   (     TensorFormatter[Mapping, "tf.Tensor", Mapping]				):






     def __init__(				self	:	List[Any]						,       a	:	List[str]=None						,       **a	:	Dict					):





         """simple docstring"""



         super().__init__(features=a					)
         __lowerCamelCase						=						tf_tensor_kwargs
         import tensorflow as tf  # noqa: import tf at initialization






     def      SCREAMING_SNAKE_CASE__     (				self	:	str						,       a	:	int					):





         """simple docstring"""



         import tensorflow as tf

         if isinstance(a						,       a					) and column:
             if all(
                 isinstance(a						,       tf.Tensor					) and x.shape == column[0].shape and x.dtype == column[0].dtype for x in column					):
                 return tf.stack(a					)
             elif all(
                 isinstance(a						,       (tf.Tensor, tf.RaggedTensor)					) and x.ndim == 1 and x.dtype == column[0].dtype
                 for x in column					):
                 # only rag 1-D tensors, otherwise some dimensions become ragged even though they were consolidated
                 return tf.ragged.stack(a					)

         return column






     def      SCREAMING_SNAKE_CASE__     (				self	:	Optional[int]						,       a	:	int					):





         """simple docstring"""



         import tensorflow as tf

         if value is None:
             return value

         __lowerCamelCase						=						{}

         if isinstance(a						,       (np.number, np.ndarray)					) and np.issubdtype(value.dtype						,       np.integer					):
             __lowerCamelCase						=						{'''dtype''': tf.intaa}
         elif isinstance(a						,       (np.number, np.ndarray)					) and np.issubdtype(value.dtype						,       np.floating					):
             __lowerCamelCase						=						{'''dtype''': tf.floataa}
         elif config.PIL_AVAILABLE and "PIL" in sys.modules:
             import PIL.Image

             if isinstance(a						,       PIL.Image.Image					):
                 __lowerCamelCase						=						np.asarray(a					)

         return tf.convert_to_tensor(a						,       **{**default_dtype, **self.tf_tensor_kwargs}					)






     def      SCREAMING_SNAKE_CASE__     (				self	:	int						,       a	:	str					):





         """simple docstring"""



         import tensorflow as tf

         # support for torch, tf, jax etc.
         if config.TORCH_AVAILABLE and "torch" in sys.modules:
             import torch

             if isinstance(a						,       torch.Tensor					):
                 return self._tensorize(data_struct.detach().cpu().numpy()[()]					)
         if hasattr(a						,       '''__array__'''					) and not isinstance(a						,       tf.Tensor					):
             __lowerCamelCase						=						data_struct.__array__()
         # support for nested types like struct of list of struct
         if isinstance(a						,       np.ndarray					):
             if data_struct.dtype == object:  # tf tensors cannot be instantied from an array of objects
                 return self._consolidate([self.recursive_tensorize(a					) for substruct in data_struct]					)
         elif isinstance(a						,       (list, tuple)					):
             return self._consolidate([self.recursive_tensorize(a					) for substruct in data_struct]					)
         return self._tensorize(a					)






     def      SCREAMING_SNAKE_CASE__     (				self	:	Union[str, Any]						,       a	:	dict					):





         """simple docstring"""



         return map_nested(self._recursive_tensorize						,       a						,       map_list=a					)






     def      SCREAMING_SNAKE_CASE__     (				self	:	Dict						,       a	:	pa.Table					):





         """simple docstring"""



         __lowerCamelCase						=						self.numpy_arrow_extractor().extract_row(a					)
         __lowerCamelCase						=						self.python_features_decoder.decode_row(a					)
         return self.recursive_tensorize(a					)






     def      SCREAMING_SNAKE_CASE__     (				self	:	List[str]						,       a	:	pa.Table					):





         """simple docstring"""



         __lowerCamelCase						=						self.numpy_arrow_extractor().extract_column(a					)
         __lowerCamelCase						=						self.python_features_decoder.decode_column(a						,       pa_table.column_names[0]					)
         __lowerCamelCase						=						self.recursive_tensorize(a					)
         __lowerCamelCase						=						self._consolidate(a					)
         return column






     def      SCREAMING_SNAKE_CASE__     (				self	:	Optional[Any]						,       a	:	pa.Table					):





         """simple docstring"""



         __lowerCamelCase						=						self.numpy_arrow_extractor().extract_batch(a					)
         __lowerCamelCase						=						self.python_features_decoder.decode_batch(a					)
         __lowerCamelCase						=						self.recursive_tensorize(a					)
         for column_name in batch:
             __lowerCamelCase						=						self._consolidate(batch[column_name]					)
         return batch

'''simple docstring'''




import copy

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


__UpperCAmelCase  							=logging.get_logger(__name__)

__UpperCAmelCase  							={
    "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 					a__   (     UpperCAmelCase__				):
     lowerCamelCase						:  Union[str, Any]							="musicgen_decoder"
     lowerCamelCase						:  List[Any]							=["past_key_values"]






     def __init__(				self	:	List[str]						,       a	:	Dict=20_48						,       a	:	Any=20_48						,       a	:	int=24						,       a	:	Optional[Any]=40_96						,       a	:	Tuple=16						,       a	:	Optional[int]=0.0						,       a	:	Tuple=True						,       a	:	Dict="gelu"						,       a	:	Optional[Any]=10_24						,       a	:	Optional[Any]=0.1						,       a	:	Tuple=0.0						,       a	:	Optional[int]=0.0						,       a	:	Any=0.02						,       a	:	str=False						,       a	:	List[Any]=4						,       a	:	List[str]=20_48						,       a	:	Any=20_48						,       a	:	Dict=None						,       a	:	str=False						,       **a	:	Tuple						,       ):





         """simple docstring"""



         __lowerCamelCase						=						vocab_size
         __lowerCamelCase						=						max_position_embeddings
         __lowerCamelCase						=						hidden_size
         __lowerCamelCase						=						ffn_dim
         __lowerCamelCase						=						num_hidden_layers
         __lowerCamelCase						=						num_attention_heads
         __lowerCamelCase						=						dropout
         __lowerCamelCase						=						attention_dropout
         __lowerCamelCase						=						activation_dropout
         __lowerCamelCase						=						activation_function
         __lowerCamelCase						=						initializer_factor
         __lowerCamelCase						=						layerdrop
         __lowerCamelCase						=						use_cache
         __lowerCamelCase						=						scale_embedding  # scale factor will be sqrt(d_model) if True
         __lowerCamelCase						=						num_codebooks
         super().__init__(
             pad_token_id=a						,       bos_token_id=a						,       eos_token_id=a						,       tie_word_embeddings=a						,       **a						,       )




class 					a__   (     UpperCAmelCase__				):
     lowerCamelCase						:  Any							="musicgen"
     lowerCamelCase						:  str							=True






     def __init__(				self	:	str						,       **a	:	int					):





         """simple docstring"""



         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'''					)

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

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

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

         __lowerCamelCase						=						AutoConfig.for_model(a						,       **a					)
         __lowerCamelCase						=						AutoConfig.for_model(a						,       **a					)
         __lowerCamelCase						=						MusicgenDecoderConfig(**a					)
         __lowerCamelCase						=						True






     @classmethod
     def      SCREAMING_SNAKE_CASE__     (				cls	:	Dict						,       a	:	PretrainedConfig						,       a	:	PretrainedConfig						,       a	:	MusicgenDecoderConfig						,       **a	:	str						,       ):





         """simple docstring"""



         return cls(
             text_encoder=text_encoder_config.to_dict()						,       audio_encoder=audio_encoder_config.to_dict()						,       decoder=decoder_config.to_dict()						,       **a						,       )






     def      SCREAMING_SNAKE_CASE__     (				self	:	List[Any]					):





         """simple docstring"""



         __lowerCamelCase						=						copy.deepcopy(self.__dict__					)
         __lowerCamelCase						=						self.text_encoder.to_dict()
         __lowerCamelCase						=						self.audio_encoder.to_dict()
         __lowerCamelCase						=						self.decoder.to_dict()
         __lowerCamelCase						=						self.__class__.model_type
         return output
