

import inspect
import unittest

from transformers import ResNetConfig, is_flax_available
from transformers.testing_utils import require_flax, slow
from transformers.utils import cached_property, is_vision_available

from ...test_configuration_common import ConfigTester
from ...test_modeling_flax_common import FlaxModelTesterMixin, floats_tensor


if is_flax_available():
       import jax
       import jax.numpy as jnp

       from transformers.models.resnet.modeling_flax_resnet import FlaxResNetForImageClassification, FlaxResNetModel

if is_vision_available():
       from PIL import Image

       from transformers import AutoImageProcessor







class 						__magic_name__      (							unittest.TestCase							):

   '''simple docstring'''




   def __init__(    self,   lowercase_,   lowercase_=3,   lowercase_=32,   lowercase_=3,   lowercase_=10,   lowercase_=[10, 20, 30, 40],   lowercase_=[1, 1, 2, 1],   lowercase_=True,   lowercase_=True,   lowercase_="relu",   lowercase_=3,   lowercase_=None,   )							->   Optional[int]:






     """simple docstring"""






     a__  =parent
     a__  =batch_size
     a__  =image_size
     a__  =num_channels
     a__  =embeddings_size
     a__  =hidden_sizes
     a__  =depths
     a__  =is_training
     a__  =use_labels
     a__  =hidden_act
     a__  =num_labels
     a__  =scope
     a__  =len(lowercase_	)



   def 			_UpperCAmelCase			(    self	)							->   Dict:






     """simple docstring"""






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

     a__  =self.get_config()

     return config, pixel_values



   def 			_UpperCAmelCase			(    self	)							->   List[Any]:






     """simple docstring"""






     return ResNetConfig(
         num_channels=self.num_channels,   embeddings_size=self.embeddings_size,   hidden_sizes=self.hidden_sizes,   depths=self.depths,   hidden_act=self.hidden_act,   num_labels=self.num_labels,   image_size=self.image_size,   )



   def 			_UpperCAmelCase			(    self,   lowercase_,   lowercase_	)							->   str:






     """simple docstring"""






     a__  =FlaxResNetModel(config=lowercase_	)
     a__  =model(lowercase_	)

     # Output shape (b, c, h, w)
     self.parent.assertEqual(
         result.last_hidden_state.shape,   (self.batch_size, self.hidden_sizes[-1], self.image_size // 32, self.image_size // 32),   )



   def 			_UpperCAmelCase			(    self,   lowercase_,   lowercase_	)							->   str:






     """simple docstring"""






     a__  =self.num_labels
     a__  =FlaxResNetForImageClassification(config=lowercase_	)
     a__  =model(lowercase_	)
     self.parent.assertEqual(result.logits.shape,   (self.batch_size, self.num_labels)	)



   def 			_UpperCAmelCase			(    self	)							->   List[Any]:






     """simple docstring"""






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







@require_flax
class 						__magic_name__      (							lowerCamelCase__		,	unittest.TestCase							):

   '''simple docstring'''




   lowerCamelCase__			:     Optional[int]    =       (FlaxResNetModel, FlaxResNetForImageClassification) if is_flax_available() else ()

   lowerCamelCase__			:     List[Any]    =       False
   lowerCamelCase__			:     str    =       False
   lowerCamelCase__			:     int    =       False



   def 			_UpperCAmelCase			(    self	)							->   None:






     """simple docstring"""






     a__  =FlaxResNetModelTester(self	)
     a__  =ConfigTester(self,   config_class=lowercase_,   has_text_modality=lowercase_	)



   def 			_UpperCAmelCase			(    self	)							->   Tuple:






     """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 			_UpperCAmelCase			(    self	)							->   Dict:






     """simple docstring"""






     return



   def 			_UpperCAmelCase			(    self	)							->   Any:






     """simple docstring"""






     a__  =self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_model(*lowercase_	)



   def 			_UpperCAmelCase			(    self	)							->   List[Any]:






     """simple docstring"""






     a__  =self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_for_image_classification(*lowercase_	)



   @unittest.skip(reason='''ResNet does not use inputs_embeds'''	)
   def 			_UpperCAmelCase			(    self	)							->   str:






     """simple docstring"""






     pass



   @unittest.skip(reason='''ResNet does not support input and output embeddings'''	)
   def 			_UpperCAmelCase			(    self	)							->   List[str]:






     """simple docstring"""






     pass



   def 			_UpperCAmelCase			(    self	)							->   Optional[int]:






     """simple docstring"""






     a__,				a__  =self.model_tester.prepare_config_and_inputs_for_common()

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

       a__  =['''pixel_values''']
       self.assertListEqual(arg_names[:1],   lowercase_	)



   def 			_UpperCAmelCase			(    self	)							->   Any:






     """simple docstring"""






     def check_hidden_states_output(lowercase_,   lowercase_,   lowercase_	):
       a__  =model_class(lowercase_	)

       a__  =model(**self._prepare_for_class(lowercase_,   lowercase_	)	)

       a__  =outputs.encoder_hidden_states if config.is_encoder_decoder else outputs.hidden_states

       a__  =self.model_tester.num_stages
       self.assertEqual(len(lowercase_	),   expected_num_stages + 1	)



   @unittest.skip(reason='''ResNet does not use feedforward chunking'''	)
   def 			_UpperCAmelCase			(    self	)							->   int:






     """simple docstring"""






     pass



   def 			_UpperCAmelCase			(    self	)							->   List[Any]:






     """simple docstring"""






     a__,				a__  =self.model_tester.prepare_config_and_inputs_for_common()

     for model_class in self.all_model_classes:
       with self.subTest(model_class.__name__	):
         a__  =self._prepare_for_class(lowercase_,   lowercase_	)
         a__  =model_class(lowercase_	)

         @jax.jit
         def model_jitted(lowercase_,   **lowercase_	):
           return model(pixel_values=lowercase_,   **lowercase_	)

         with self.subTest('''JIT Enabled'''	):
           a__  =model_jitted(**lowercase_	).to_tuple()

         with self.subTest('''JIT Disabled'''	):
           with jax.disable_jit():
             a__  =model_jitted(**lowercase_	).to_tuple()

         self.assertEqual(len(lowercase_	),   len(lowercase_	)	)
         for jitted_output, output in zip(lowercase_,   lowercase_	):
           self.assertEqual(jitted_output.shape,   output.shape	)




def       UpperCAmelCase__    (      ):





  '''simple docstring'''







  a__  =Image.open('''./tests/fixtures/tests_samples/COCO/000000039769.png'''							)
  return image







@require_flax
class 						__magic_name__      (							unittest.TestCase							):

   '''simple docstring'''




   @cached_property
   def 			_UpperCAmelCase			(    self	)							->   Optional[Any]:






     """simple docstring"""






     return AutoImageProcessor.from_pretrained('''microsoft/resnet-50'''	) if is_vision_available() else None



   @slow
   def 			_UpperCAmelCase			(    self	)							->   List[Any]:






     """simple docstring"""






     a__  =FlaxResNetForImageClassification.from_pretrained('''microsoft/resnet-50'''	)

     a__  =self.default_image_processor
     a__  =prepare_img()
     a__  =image_processor(images=lowercase_,   return_tensors='''np'''	)

     a__  =model(**lowercase_	)

     # verify the logits
     a__  =(1, 1000)
     self.assertEqual(outputs.logits.shape,   lowercase_	)

     a__  =jnp.array([-11.1069, -9.7877, -8.3777]	)

     self.assertTrue(jnp.allclose(outputs.logits[0, :3],   lowercase_,   atol=1E-4	)	)



from collections import OrderedDict
from typing import Any, Mapping, Optional

from ... import PreTrainedTokenizer
from ...configuration_utils import PretrainedConfig
from ...file_utils import TensorType, is_torch_available
from ...onnx import OnnxConfig, OnnxConfigWithPast, OnnxSeqaSeqConfigWithPast
from ...onnx.utils import compute_effective_axis_dimension
from ...utils import logging


lowerCamelCase									= logging.get_logger(__name__)

lowerCamelCase									= {
    '''facebook/blenderbot-3B''': '''https://huggingface.co/facebook/blenderbot-3B/resolve/main/config.json''',
    # See all Blenderbot models at https://huggingface.co/models?filter=blenderbot
}







class 						__magic_name__      (							lowerCamelCase__							):

   '''simple docstring'''




   lowerCamelCase__			:     List[Any]    =       'blenderbot'
   lowerCamelCase__			:     int    =       ['past_key_values']
   lowerCamelCase__			:     Optional[Any]    =       {'num_attention_heads': 'encoder_attention_heads', 'hidden_size': 'd_model'}



   def __init__(    self,   lowercase_=8008,   lowercase_=128,   lowercase_=2,   lowercase_=10240,   lowercase_=32,   lowercase_=24,   lowercase_=10240,   lowercase_=32,   lowercase_=0.0,   lowercase_=0.0,   lowercase_=True,   lowercase_=True,   lowercase_="gelu",   lowercase_=2560,   lowercase_=0.1,   lowercase_=0.0,   lowercase_=0.0,   lowercase_=0.02,   lowercase_=1,   lowercase_=False,   lowercase_=0,   lowercase_=1,   lowercase_=2,   lowercase_=3,   lowercase_=2,   **lowercase_,   )							->   Any:






     """simple docstring"""






     a__  =vocab_size
     a__  =max_position_embeddings
     a__  =d_model
     a__  =encoder_ffn_dim
     a__  =encoder_layers
     a__  =encoder_attention_heads
     a__  =decoder_ffn_dim
     a__  =decoder_layers
     a__  =decoder_attention_heads
     a__  =dropout
     a__  =attention_dropout
     a__  =activation_dropout
     a__  =activation_function
     a__  =init_std
     a__  =encoder_layerdrop
     a__  =decoder_layerdrop
     a__  =use_cache
     a__  =encoder_layers
     a__  =scale_embedding  # scale factor will be sqrt(d_model) if True

     super().__init__(
         pad_token_id=lowercase_,   bos_token_id=lowercase_,   eos_token_id=lowercase_,   is_encoder_decoder=lowercase_,   decoder_start_token_id=lowercase_,   encoder_no_repeat_ngram_size=lowercase_,   forced_eos_token_id=lowercase_,   **lowercase_,   )







class 						__magic_name__      (							lowerCamelCase__							):

   '''simple docstring'''




   @property
   def 			_UpperCAmelCase			(    self	)							->   Mapping[str, Mapping[int, str]]:






     """simple docstring"""






     if self.task in ["default", "seq2seq-lm"]:
       a__  =OrderedDict(
           [
               ('''input_ids''', {0: '''batch''', 1: '''encoder_sequence'''}),
               ('''attention_mask''', {0: '''batch''', 1: '''encoder_sequence'''}),
           ]	)
       if self.use_past:
         a__  ={0: '''batch'''}
         a__  ={0: '''batch''', 1: '''past_decoder_sequence + sequence'''}
       else:
         a__  ={0: '''batch''', 1: '''decoder_sequence'''}
         a__  ={0: '''batch''', 1: '''decoder_sequence'''}
       if self.use_past:
         self.fill_with_past_key_values_(lowercase_,   direction='''inputs'''	)
     elif self.task == "causal-lm":
       a__  =OrderedDict(
           [
               ('''input_ids''', {0: '''batch''', 1: '''encoder_sequence'''}),
               ('''attention_mask''', {0: '''batch''', 1: '''encoder_sequence'''}),
           ]	)
       if self.use_past:
         a__,				a__  =self.num_layers
         for i in range(lowercase_	):
           a__  ={0: '''batch''', 2: '''past_sequence + sequence'''}
           a__  ={0: '''batch''', 2: '''past_sequence + sequence'''}
     else:
       a__  =OrderedDict(
           [
               ('''input_ids''', {0: '''batch''', 1: '''encoder_sequence'''}),
               ('''attention_mask''', {0: '''batch''', 1: '''encoder_sequence'''}),
               ('''decoder_input_ids''', {0: '''batch''', 1: '''decoder_sequence'''}),
               ('''decoder_attention_mask''', {0: '''batch''', 1: '''decoder_sequence'''}),
           ]	)

     return common_inputs



   @property
   # Copied from transformers.models.bart.configuration_bart.BartOnnxConfig.outputs
   def 			_UpperCAmelCase			(    self	)							->   Mapping[str, Mapping[int, str]]:






     """simple docstring"""






     if self.task in ["default", "seq2seq-lm"]:
       a__  =super().outputs
     else:
       a__  =super(lowercase_,   self	).outputs
       if self.use_past:
         a__,				a__  =self.num_layers
         for i in range(lowercase_	):
           a__  ={0: '''batch''', 2: '''past_sequence + sequence'''}
           a__  ={0: '''batch''', 2: '''past_sequence + sequence'''}
     return common_outputs



   def 			_UpperCAmelCase			(    self,   lowercase_,   lowercase_ = -1,   lowercase_ = -1,   lowercase_ = False,   lowercase_ = None,   )							->   Mapping[str, Any]:






     """simple docstring"""






     a__  =self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
         lowercase_,   lowercase_,   lowercase_,   lowercase_,   lowercase_	)
     # Generate decoder inputs
     a__  =seq_length if not self.use_past else 1
     a__  =self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
         lowercase_,   lowercase_,   lowercase_,   lowercase_,   lowercase_	)
     a__  ={F"""decoder_{name}""": tensor for name, tensor in decoder_inputs.items()}
     a__  =dict(**lowercase_,   **lowercase_	)

     if self.use_past:
       if not is_torch_available():
         raise ValueError('''Cannot generate dummy past_keys inputs without PyTorch installed.'''	)
       else:
         import torch
       a__,				a__  =common_inputs['''input_ids'''].shape
       a__  =common_inputs['''decoder_input_ids'''].shape[1]
       a__,				a__  =self.num_attention_heads
       a__  =(
           batch,
           num_encoder_attention_heads,
           encoder_seq_length,
           self._config.hidden_size // num_encoder_attention_heads,
       )
       a__  =decoder_seq_length
       a__  =(
           batch,
           num_decoder_attention_heads,
           decoder_past_length,
           self._config.hidden_size // num_decoder_attention_heads,
       )
       a__  =torch.cat(
           [common_inputs['''decoder_attention_mask'''], torch.ones(lowercase_,   lowercase_	)],   dim=1	)
       a__  =[]
       a__,				a__  =self.num_layers

       for _ in range(lowercase_	):
         common_inputs["past_key_values"].append(
             (
                 torch.zeros(lowercase_	),
                 torch.zeros(lowercase_	),
                 torch.zeros(lowercase_	),
                 torch.zeros(lowercase_	),
             )	)
     return common_inputs



   def 			_UpperCAmelCase			(    self,   lowercase_,   lowercase_ = -1,   lowercase_ = -1,   lowercase_ = False,   lowercase_ = None,   )							->   Mapping[str, Any]:






     """simple docstring"""






     a__  =self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
         lowercase_,   lowercase_,   lowercase_,   lowercase_,   lowercase_	)

     if self.use_past:
       if not is_torch_available():
         raise ValueError('''Cannot generate dummy past_keys inputs without PyTorch installed.'''	)
       else:
         import torch
       a__,				a__  =common_inputs['''input_ids'''].shape
       a__  =seqlen
       a__,				a__  =self.num_layers
       a__,				a__  =self.num_attention_heads
       a__  =(
           batch,
           num_encoder_attention_heads,
           past_key_values_length,
           self._config.hidden_size // num_encoder_attention_heads,
       )
       a__  =common_inputs['''attention_mask'''].dtype
       a__  =torch.cat(
           [common_inputs['''attention_mask'''], torch.ones(lowercase_,   lowercase_,   dtype=lowercase_	)],   dim=1	)
       a__  =[
           (torch.zeros(lowercase_	), torch.zeros(lowercase_	)) for _ in range(lowercase_	)
       ]
     return common_inputs



   def 			_UpperCAmelCase			(    self,   lowercase_,   lowercase_ = -1,   lowercase_ = -1,   lowercase_ = False,   lowercase_ = None,   )							->   Mapping[str, Any]:






     """simple docstring"""






     a__  =compute_effective_axis_dimension(
         lowercase_,   fixed_dimension=OnnxConfig.default_fixed_batch,   num_token_to_add=0	)

     # If dynamic axis (-1) we forward with a fixed dimension of 8 tokens to avoid optimizations made by ONNX
     a__  =tokenizer.num_special_tokens_to_add(lowercase_	)
     a__  =compute_effective_axis_dimension(
         lowercase_,   fixed_dimension=OnnxConfig.default_fixed_sequence,   num_token_to_add=lowercase_	)

     # Generate dummy inputs according to compute batch and sequence
     a__  =[''' '''.join([tokenizer.unk_token]	) * seq_length] * batch_size
     a__  =dict(tokenizer(lowercase_,   return_tensors=lowercase_	)	)
     return common_inputs



   def 			_UpperCAmelCase			(    self,   lowercase_,   lowercase_ = -1,   lowercase_ = -1,   lowercase_ = False,   lowercase_ = None,   )							->   Mapping[str, Any]:






     """simple docstring"""






     if self.task in ["default", "seq2seq-lm"]:
       a__  =self._generate_dummy_inputs_for_default_and_seqaseq_lm(
           lowercase_,   batch_size=lowercase_,   seq_length=lowercase_,   is_pair=lowercase_,   framework=lowercase_	)

     elif self.task == "causal-lm":
       a__  =self._generate_dummy_inputs_for_causal_lm(
           lowercase_,   batch_size=lowercase_,   seq_length=lowercase_,   is_pair=lowercase_,   framework=lowercase_	)
     else:
       a__  =self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
           lowercase_,   batch_size=lowercase_,   seq_length=lowercase_,   is_pair=lowercase_,   framework=lowercase_	)

     return common_inputs



   def 			_UpperCAmelCase			(    self,   lowercase_,   lowercase_,   lowercase_,   lowercase_	)							->   Optional[int]:






     """simple docstring"""






     if self.task in ["default", "seq2seq-lm"]:
       a__  =super()._flatten_past_key_values_(lowercase_,   lowercase_,   lowercase_,   lowercase_	)
     else:
       a__  =super(lowercase_,   self	)._flatten_past_key_values_(
           lowercase_,   lowercase_,   lowercase_,   lowercase_	)



   def 			_UpperCAmelCase			(    self,   lowercase_,   lowercase_	)							->   Tuple:






     """simple docstring"""






     if direction not in ["inputs", "outputs"]:
       raise ValueError(F"""direction must either be \"inputs\" or \"outputs\", but {direction} was given"""	)

     a__  ='''past_key_values''' if direction == '''inputs''' else '''present'''
     a__,				a__  =self.num_layers

     a__  ='''past_encoder_sequence'''
     a__  ='''past_decoder_sequence''' if direction == '''inputs''' else '''past_decoder_sequence + sequence'''

     for i in range(lowercase_	):
       a__  ={0: '''batch''', 2: decoder_sequence}
       a__  ={0: '''batch''', 2: decoder_sequence}
       a__  ={0: '''batch''', 2: encoder_sequence}
       a__  ={0: '''batch''', 2: encoder_sequence}
