




'''simple docstring'''




import collections
import inspect
import unittest

from transformers import NatConfig
from transformers.testing_utils import require_natten, require_torch, require_vision, slow, torch_device
from transformers.utils import cached_property, is_torch_available, is_vision_available

from ...test_backbone_common import BackboneTesterMixin
from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, _config_zero_init, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
  import torch
  from torch import nn

  from transformers import NatBackbone, NatForImageClassification, NatModel
  from transformers.models.nat.modeling_nat import NAT_PRETRAINED_MODEL_ARCHIVE_LIST

if is_vision_available():
  from PIL import Image

  from transformers import AutoImageProcessor




class 					a__   :






     def __init__(				self	:	Union[str, Any]						,       a	:	str						,       a	:	List[Any]=13						,       a	:	Union[str, Any]=64						,       a	:	List[Any]=4						,       a	:	List[Any]=3						,       a	:	List[Any]=16						,       a	:	str=[1, 2, 1]						,       a	:	Any=[2, 4, 8]						,       a	:	Dict=3						,       a	:	int=2.0						,       a	:	List[Any]=True						,       a	:	Optional[int]=0.0						,       a	:	List[str]=0.0						,       a	:	Union[str, Any]=0.1						,       a	:	Optional[Any]="gelu"						,       a	:	int=True						,       a	:	int=0.02						,       a	:	Optional[Any]=1e-5						,       a	:	Optional[Any]=True						,       a	:	List[str]=None						,       a	:	Any=True						,       a	:	List[str]=10						,       a	:	List[Any]=["stage1", "stage2"]						,       a	:	str=[1, 2]						,       ):





         """simple docstring"""



         __lowerCamelCase						=						parent
         __lowerCamelCase						=						batch_size
         __lowerCamelCase						=						image_size
         __lowerCamelCase						=						patch_size
         __lowerCamelCase						=						num_channels
         __lowerCamelCase						=						embed_dim
         __lowerCamelCase						=						depths
         __lowerCamelCase						=						num_heads
         __lowerCamelCase						=						kernel_size
         __lowerCamelCase						=						mlp_ratio
         __lowerCamelCase						=						qkv_bias
         __lowerCamelCase						=						hidden_dropout_prob
         __lowerCamelCase						=						attention_probs_dropout_prob
         __lowerCamelCase						=						drop_path_rate
         __lowerCamelCase						=						hidden_act
         __lowerCamelCase						=						patch_norm
         __lowerCamelCase						=						layer_norm_eps
         __lowerCamelCase						=						initializer_range
         __lowerCamelCase						=						is_training
         __lowerCamelCase						=						scope
         __lowerCamelCase						=						use_labels
         __lowerCamelCase						=						num_labels
         __lowerCamelCase						=						out_features
         __lowerCamelCase						=						out_indices






     def      SCREAMING_SNAKE_CASE__     (				self	:	Optional[int]					):





         """simple docstring"""



         __lowerCamelCase						=						floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]					)

         __lowerCamelCase						=						None
         if self.use_labels:
             __lowerCamelCase						=						ids_tensor([self.batch_size]						,       self.num_labels					)

         __lowerCamelCase						=						self.get_config()

         return config, pixel_values, labels






     def      SCREAMING_SNAKE_CASE__     (				self	:	List[Any]					):





         """simple docstring"""



         return NatConfig(
             num_labels=self.num_labels						,       image_size=self.image_size						,       patch_size=self.patch_size						,       num_channels=self.num_channels						,       embed_dim=self.embed_dim						,       depths=self.depths						,       num_heads=self.num_heads						,       kernel_size=self.kernel_size						,       mlp_ratio=self.mlp_ratio						,       qkv_bias=self.qkv_bias						,       hidden_dropout_prob=self.hidden_dropout_prob						,       attention_probs_dropout_prob=self.attention_probs_dropout_prob						,       drop_path_rate=self.drop_path_rate						,       hidden_act=self.hidden_act						,       patch_norm=self.patch_norm						,       layer_norm_eps=self.layer_norm_eps						,       initializer_range=self.initializer_range						,       out_features=self.out_features						,       out_indices=self.out_indices						,       )






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





         """simple docstring"""



         __lowerCamelCase						=						NatModel(config=a					)
         model.to(a					)
         model.eval()
         __lowerCamelCase						=						model(a					)

         __lowerCamelCase						=						__lowerCamelCase						=						(config.image_size // config.patch_size) // (2 ** (len(config.depths					) - 1))
         __lowerCamelCase						=						int(config.embed_dim * 2 ** (len(config.depths					) - 1)					)

         self.parent.assertEqual(
             result.last_hidden_state.shape						,       (self.batch_size, expected_height, expected_width, expected_dim)					)






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





         """simple docstring"""



         __lowerCamelCase						=						NatForImageClassification(a					)
         model.to(a					)
         model.eval()
         __lowerCamelCase						=						model(a						,       labels=a					)
         self.parent.assertEqual(result.logits.shape						,       (self.batch_size, self.num_labels)					)

         # test greyscale images
         __lowerCamelCase						=						1
         __lowerCamelCase						=						NatForImageClassification(a					)
         model.to(a					)
         model.eval()

         __lowerCamelCase						=						floats_tensor([self.batch_size, 1, self.image_size, self.image_size]					)
         __lowerCamelCase						=						model(a					)
         self.parent.assertEqual(result.logits.shape						,       (self.batch_size, self.num_labels)					)






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





         """simple docstring"""



         __lowerCamelCase						=						NatBackbone(config=a					)
         model.to(a					)
         model.eval()
         __lowerCamelCase						=						model(a					)

         # verify hidden states
         self.parent.assertEqual(len(result.feature_maps					)						,       len(config.out_features					)					)
         self.parent.assertListEqual(list(result.feature_maps[0].shape					)						,       [self.batch_size, model.channels[0], 16, 16]					)

         # verify channels
         self.parent.assertEqual(len(model.channels					)						,       len(config.out_features					)					)

         # verify backbone works with out_features=None
         __lowerCamelCase						=						None
         __lowerCamelCase						=						NatBackbone(config=a					)
         model.to(a					)
         model.eval()
         __lowerCamelCase						=						model(a					)

         # verify feature maps
         self.parent.assertEqual(len(result.feature_maps					)						,       1					)
         self.parent.assertListEqual(list(result.feature_maps[0].shape					)						,       [self.batch_size, model.channels[-1], 4, 4]					)

         # verify channels
         self.parent.assertEqual(len(model.channels					)						,       1					)






     def      SCREAMING_SNAKE_CASE__     (				self	:	List[str]					):





         """simple docstring"""



         __lowerCamelCase						=						self.prepare_config_and_inputs()
         __lowerCamelCase			,						__lowerCamelCase			,						__lowerCamelCase						=						config_and_inputs
         __lowerCamelCase						=						{'''pixel_values''': pixel_values}
         return config, inputs_dict




@require_natten
@require_torch
class 					a__   (     snake_case_							,     snake_case_							,     unittest.TestCase				):
     lowerCamelCase						:  Optional[Any]							=(
         (
             NatModel,
             NatForImageClassification,
             NatBackbone,
         )
         if is_torch_available()
         else ()
     )
     lowerCamelCase						:  Optional[Any]							=(
         {"feature-extraction": NatModel, "image-classification": NatForImageClassification}
         if is_torch_available()
         else {}
     )
     lowerCamelCase						:  Any							=False

     lowerCamelCase						:  Dict							=False
     lowerCamelCase						:  List[str]							=False
     lowerCamelCase						:  Tuple							=False
     lowerCamelCase						:  List[str]							=False






     def      SCREAMING_SNAKE_CASE__     (				self	:	int					):





         """simple docstring"""



         __lowerCamelCase						=						NatModelTester(self					)
         __lowerCamelCase						=						ConfigTester(self						,       config_class=a						,       embed_dim=37					)






     def      SCREAMING_SNAKE_CASE__     (				self	:	List[str]					):





         """simple docstring"""



         self.create_and_test_config_common_properties()
         self.config_tester.create_and_test_config_to_json_string()
         self.config_tester.create_and_test_config_to_json_file()
         self.config_tester.create_and_test_config_from_and_save_pretrained()
         self.config_tester.create_and_test_config_with_num_labels()
         self.config_tester.check_config_can_be_init_without_params()
         self.config_tester.check_config_arguments_init()






     def      SCREAMING_SNAKE_CASE__     (				self	:	str					):





         """simple docstring"""



         return






     def      SCREAMING_SNAKE_CASE__     (				self	:	Dict					):





         """simple docstring"""



         __lowerCamelCase						=						self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_model(*a					)






     def      SCREAMING_SNAKE_CASE__     (				self	:	int					):





         """simple docstring"""



         __lowerCamelCase						=						self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_for_image_classification(*a					)






     def      SCREAMING_SNAKE_CASE__     (				self	:	Any					):





         """simple docstring"""



         __lowerCamelCase						=						self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_backbone(*a					)






     @unittest.skip(reason='''Nat does not use inputs_embeds'''					)
     def      SCREAMING_SNAKE_CASE__     (				self	:	List[str]					):





         """simple docstring"""



         pass






     @unittest.skip(reason='''Nat does not use feedforward chunking'''					)
     def      SCREAMING_SNAKE_CASE__     (				self	:	str					):





         """simple docstring"""



         pass






     def      SCREAMING_SNAKE_CASE__     (				self	:	int					):





         """simple docstring"""



         __lowerCamelCase			,						__lowerCamelCase						=						self.model_tester.prepare_config_and_inputs_for_common()

         for model_class in self.all_model_classes:
             __lowerCamelCase						=						model_class(a					)
             self.assertIsInstance(model.get_input_embeddings()						,       (nn.Module)					)
             __lowerCamelCase						=						model.get_output_embeddings()
             self.assertTrue(x is None or isinstance(a						,       nn.Linear					)					)






     def      SCREAMING_SNAKE_CASE__     (				self	:	int					):





         """simple docstring"""



         __lowerCamelCase			,						__lowerCamelCase						=						self.model_tester.prepare_config_and_inputs_for_common()

         for model_class in self.all_model_classes:
             __lowerCamelCase						=						model_class(a					)
             __lowerCamelCase						=						inspect.signature(model.forward					)
             # signature.parameters is an OrderedDict => so arg_names order is deterministic
             __lowerCamelCase						=						[*signature.parameters.keys()]

             __lowerCamelCase						=						['''pixel_values''']
             self.assertListEqual(arg_names[:1]						,       a					)






     def      SCREAMING_SNAKE_CASE__     (				self	:	Any					):





         """simple docstring"""



         self.skipTest('''Nat\'s attention operation is handled entirely by NATTEN.'''					)






     def      SCREAMING_SNAKE_CASE__     (				self	:	Any						,       a	:	Optional[int]						,       a	:	Union[str, Any]						,       a	:	Optional[int]						,       a	:	List[Any]					):





         """simple docstring"""



         __lowerCamelCase						=						model_class(a					)
         model.to(a					)
         model.eval()

         with torch.no_grad():
             __lowerCamelCase						=						model(**self._prepare_for_class(a						,       a					)					)

         __lowerCamelCase						=						outputs.hidden_states

         __lowerCamelCase						=						getattr(
             self.model_tester						,       '''expected_num_hidden_layers'''						,       len(self.model_tester.depths					) + 1					)
         self.assertEqual(len(a					)						,       a					)

         # Nat has a different seq_length
         __lowerCamelCase						=						(
             config.patch_size
             if isinstance(config.patch_size						,       collections.abc.Iterable					)
             else (config.patch_size, config.patch_size)
         )

         __lowerCamelCase						=						image_size[0] // patch_size[0]
         __lowerCamelCase						=						image_size[1] // patch_size[1]

         self.assertListEqual(
             list(hidden_states[0].shape[-3:]					)						,       [height, width, self.model_tester.embed_dim]						,       )

         if model_class.__name__ != "NatBackbone":
             __lowerCamelCase						=						outputs.reshaped_hidden_states
             self.assertEqual(len(a					)						,       a					)

             __lowerCamelCase			,						__lowerCamelCase			,						__lowerCamelCase			,						__lowerCamelCase						=						reshaped_hidden_states[0].shape
             __lowerCamelCase						=						(
                 reshaped_hidden_states[0].view(a						,       a						,       a						,       a					).permute(0						,       2						,       3						,       1					)
             )
             self.assertListEqual(
                 list(reshaped_hidden_states.shape[-3:]					)						,       [height, width, self.model_tester.embed_dim]						,       )






     def      SCREAMING_SNAKE_CASE__     (				self	:	List[str]					):





         """simple docstring"""



         __lowerCamelCase			,						__lowerCamelCase						=						self.model_tester.prepare_config_and_inputs_for_common()

         __lowerCamelCase						=						(
             self.model_tester.image_size
             if isinstance(self.model_tester.image_size						,       collections.abc.Iterable					)
             else (self.model_tester.image_size, self.model_tester.image_size)
         )

         for model_class in self.all_model_classes:
             __lowerCamelCase						=						True
             self.check_hidden_states_output(a						,       a						,       a						,       a					)

             # check that output_hidden_states also work using config
             del inputs_dict["output_hidden_states"]
             __lowerCamelCase						=						True

             self.check_hidden_states_output(a						,       a						,       a						,       a					)






     @slow
     def      SCREAMING_SNAKE_CASE__     (				self	:	Any					):





         """simple docstring"""



         for model_name in NAT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
             __lowerCamelCase						=						NatModel.from_pretrained(a					)
             self.assertIsNotNone(a					)






     def      SCREAMING_SNAKE_CASE__     (				self	:	Optional[Any]					):





         """simple docstring"""



         __lowerCamelCase			,						__lowerCamelCase						=						self.model_tester.prepare_config_and_inputs_for_common()

         __lowerCamelCase						=						_config_zero_init(a					)
         for model_class in self.all_model_classes:
             __lowerCamelCase						=						model_class(config=a					)
             for name, param in model.named_parameters():
                 if "embeddings" not in name and param.requires_grad:
                     self.assertIn(
                         ((param.data.mean() * 1e9).round() / 1e9).item()						,       [0.0, 1.0]						,       msg=f"""Parameter {name} of model {model_class} seems not properly initialized"""						,       )




@require_natten
@require_vision
@require_torch
class 					a__   (     unittest.TestCase				):






     @cached_property
     def      SCREAMING_SNAKE_CASE__     (				self	:	str					):





         """simple docstring"""



         return AutoImageProcessor.from_pretrained('''shi-labs/nat-mini-in1k-224'''					) if is_vision_available() else None






     @slow
     def      SCREAMING_SNAKE_CASE__     (				self	:	Union[str, Any]					):





         """simple docstring"""



         __lowerCamelCase						=						NatForImageClassification.from_pretrained('''shi-labs/nat-mini-in1k-224'''					).to(a					)
         __lowerCamelCase						=						self.default_image_processor

         __lowerCamelCase						=						Image.open('''./tests/fixtures/tests_samples/COCO/000000039769.png'''					)
         __lowerCamelCase						=						image_processor(images=a						,       return_tensors='''pt'''					).to(a					)

         # forward pass
         with torch.no_grad():
             __lowerCamelCase						=						model(**a					)

         # verify the logits
         __lowerCamelCase						=						torch.Size((1, 10_00)					)
         self.assertEqual(outputs.logits.shape						,       a					)
         __lowerCamelCase						=						torch.tensor([0.38_05, -0.86_76, -0.39_12]					).to(a					)
         self.assertTrue(torch.allclose(outputs.logits[0, :3]						,       a						,       atol=1e-4					)					)




@require_torch
@require_natten
class 					a__   (     unittest.TestCase							,     snake_case_				):
     lowerCamelCase						:  Optional[Any]							=(NatBackbone,) if is_torch_available() else ()
     lowerCamelCase						:  Dict							=NatConfig






     def      SCREAMING_SNAKE_CASE__     (				self	:	Dict					):





         """simple docstring"""



         __lowerCamelCase						=						NatModelTester(self					)

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