




import tempfile
import unittest

import numpy as np

import transformers
from transformers import GPTaConfig, GPTaTokenizer, is_flax_available, is_torch_available
from transformers.testing_utils import is_pt_flax_cross_test, require_flax, slow

from ...generation.test_flax_utils import FlaxGenerationTesterMixin
from ...test_modeling_flax_common import FlaxModelTesterMixin, floats_tensor, ids_tensor, random_attention_mask


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

    from transformers.modeling_flax_pytorch_utils import (
        convert_pytorch_state_dict_to_flax,
        load_flax_weights_in_pytorch_model,
    )
    from transformers.models.gpta.modeling_flax_gpta import FlaxGPTaLMHeadModel, FlaxGPTaModel

if is_torch_available():
    import torch


class        UpperCAmelCase__				:



     '''simple docstring'''
     def __init__(				self ,    _lowerCAmelCase ,    _lowerCAmelCase=14 ,    _lowerCAmelCase=7 ,    _lowerCAmelCase=True ,    _lowerCAmelCase=True ,    _lowerCAmelCase=False ,    _lowerCAmelCase=True ,    _lowerCAmelCase=99 ,    _lowerCAmelCase=32 ,    _lowerCAmelCase=5 ,    _lowerCAmelCase=4 ,    _lowerCAmelCase=37 ,    _lowerCAmelCase="gelu" ,    _lowerCAmelCase=0.1 ,    _lowerCAmelCase=0.1 ,    _lowerCAmelCase=512 ,    _lowerCAmelCase=0.02 ,    ):
            a			   =parent
            a			   =batch_size
            a			   =seq_length
            a			   =is_training
            a			   =use_input_mask
            a			   =use_token_type_ids
            a			   =use_labels
            a			   =vocab_size
            a			   =hidden_size
            a			   =num_hidden_layers
            a			   =num_attention_heads
            a			   =intermediate_size
            a			   =hidden_act
            a			   =hidden_dropout_prob
            a			   =attention_probs_dropout_prob
            a			   =max_position_embeddings
            a			   =initializer_range
            a			   =None
            a			   =vocab_size - 1
            a			   =vocab_size - 1
            a			   =vocab_size - 1
     def        lowerCAmelCase__  (				self	):
            a			   =ids_tensor([self.batch_size, self.seq_length] ,    self.vocab_size	)

            a			   =None
            if self.use_input_mask:
                   a			   =random_attention_mask([self.batch_size, self.seq_length]	)

            a			   =GPTaConfig(
                vocab_size=self.vocab_size ,    n_embd=self.hidden_size ,    n_layer=self.num_hidden_layers ,    n_head=self.num_attention_heads ,    n_positions=self.max_position_embeddings ,    use_cache=_lowerCAmelCase ,    bos_token_id=self.bos_token_id ,    eos_token_id=self.eos_token_id ,    pad_token_id=self.pad_token_id ,    )

            return (config, input_ids, input_mask)
     def        lowerCAmelCase__  (				self	):
            a			   =self.prepare_config_and_inputs()
            a   ,      a   ,      a			   =config_and_inputs
            a			   ={"""input_ids""": input_ids, """attention_mask""": attention_mask}
            return config, inputs_dict
     def        lowerCAmelCase__  (				self	):
            a   ,      a   ,      a			   =self.prepare_config_and_inputs()

            a			   =floats_tensor([self.batch_size, self.seq_length, self.hidden_size]	)
            a			   =ids_tensor([self.batch_size, self.seq_length] ,    vocab_size=2	)

            return (
                config,
                input_ids,
                attention_mask,
                encoder_hidden_states,
                encoder_attention_mask,
            )
     def        lowerCAmelCase__  (				self ,    _lowerCAmelCase ,    _lowerCAmelCase ,    _lowerCAmelCase ,    _lowerCAmelCase	):
            a			   =20
            a			   =model_class_name(_lowerCAmelCase	)

            a			   =model.init_cache(input_ids.shape[0] ,    _lowerCAmelCase	)
            a			   =jnp.ones((input_ids.shape[0], max_decoder_length) ,    dtype="""i4"""	)

            a			   =jnp.broadcast_to(
                jnp.arange(input_ids.shape[-1] - 1	)[None, :] ,    (input_ids.shape[0], input_ids.shape[-1] - 1)	)
            a			   =model(
                input_ids[:, :-1] ,    attention_mask=_lowerCAmelCase ,    past_key_values=_lowerCAmelCase ,    position_ids=_lowerCAmelCase ,    )

            a			   =jnp.array(input_ids.shape[0] * [[input_ids.shape[-1] - 1]] ,    dtype="""i4"""	)
            a			   =model(
                input_ids[:, -1:] ,    attention_mask=_lowerCAmelCase ,    past_key_values=outputs_cache.past_key_values ,    position_ids=_lowerCAmelCase ,    )

            a			   =model(_lowerCAmelCase	)

            a			   =np.max(np.abs((outputs_cache_next[0][:, -1, :5] - outputs[0][:, -1, :5])	)	)
            self.parent.assertTrue(diff < 1E-3 ,    msg=F'''Max diff is {diff}'''	)






     def        lowerCAmelCase__  (				self ,    _lowerCAmelCase ,    _lowerCAmelCase ,    _lowerCAmelCase ,    _lowerCAmelCase	):
            a			   =20
            a			   =model_class_name(_lowerCAmelCase	)

            a			   =jnp.concatenate(
                [attention_mask, jnp.zeros((attention_mask.shape[0], max_decoder_length - attention_mask.shape[1])	)] ,    axis=-1 ,    )

            a			   =model.init_cache(input_ids.shape[0] ,    _lowerCAmelCase	)
            a			   =jnp.broadcast_to(
                jnp.arange(input_ids.shape[-1] - 1	)[None, :] ,    (input_ids.shape[0], input_ids.shape[-1] - 1)	)

            a			   =model(
                input_ids[:, :-1] ,    attention_mask=_lowerCAmelCase ,    past_key_values=_lowerCAmelCase ,    position_ids=_lowerCAmelCase ,    )
            a			   =jnp.array(input_ids.shape[0] * [[input_ids.shape[-1] - 1]] ,    dtype="""i4"""	)
            a			   =model(
                input_ids[:, -1:] ,    past_key_values=outputs_cache.past_key_values ,    attention_mask=_lowerCAmelCase ,    position_ids=_lowerCAmelCase ,    )

            a			   =model(_lowerCAmelCase ,    attention_mask=_lowerCAmelCase	)

            a			   =np.max(np.abs((outputs_cache_next[0][:, -1, :5] - outputs[0][:, -1, :5])	)	)
            self.parent.assertTrue(diff < 1E-3 ,    msg=F'''Max diff is {diff}'''	)



@require_flax
class        UpperCAmelCase__				(							UpperCAmelCase__      ,		UpperCAmelCase__      ,		unittest.TestCase				):



     '''simple docstring'''
     _SCREAMING_SNAKE_CASE      :					List[str]     					=    (FlaxGPTaModel, FlaxGPTaLMHeadModel) if is_flax_available() else ()
     _SCREAMING_SNAKE_CASE      :					Union[str, Any]     					=    (FlaxGPTaLMHeadModel,) if is_flax_available() else ()
     def        lowerCAmelCase__  (				self	):
            a			   =FlaxGPTaModelTester(self	)
     def        lowerCAmelCase__  (				self	):
            for model_class_name in self.all_model_classes:
                   a   ,      a   ,      a			   =self.model_tester.prepare_config_and_inputs()
                   self.model_tester.check_use_cache_forward(_lowerCAmelCase ,    _lowerCAmelCase ,    _lowerCAmelCase ,    _lowerCAmelCase	)
     def        lowerCAmelCase__  (				self	):
            for model_class_name in self.all_model_classes:
                   a   ,      a   ,      a			   =self.model_tester.prepare_config_and_inputs()
                   self.model_tester.check_use_cache_forward_with_attn_mask(
                       _lowerCAmelCase ,    _lowerCAmelCase ,    _lowerCAmelCase ,    _lowerCAmelCase	)
     @slow
     def        lowerCAmelCase__  (				self	):
            a			   =GPTaTokenizer.from_pretrained("""gpt2""" ,    pad_token="""</s>""" ,    padding_side="""left"""	)
            a			   =tokenizer(["""Hello this is a long string""", """Hey"""] ,    return_tensors="""np""" ,    padding=_lowerCAmelCase ,    truncation=_lowerCAmelCase	)

            a			   =FlaxGPTaLMHeadModel.from_pretrained("""gpt2"""	)
            a			   =False
            a			   =model.config.eos_token_id

            a			   =jax.jit(model.generate	)

            a			   =jit_generate(inputs["""input_ids"""] ,    attention_mask=inputs["""attention_mask"""]	).sequences

            a			   =tokenizer.batch_decode(_lowerCAmelCase ,    skip_special_tokens=_lowerCAmelCase	)

            a			   =[
                """Hello this is a long string of words. I'm going to start with the first one.\n""",
                """Hey, I'm not sure if I'm going to be able to do""",
            ]

            self.assertListEqual(_lowerCAmelCase ,    _lowerCAmelCase	)
     @is_pt_flax_cross_test
     def        lowerCAmelCase__  (				self	):
            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__	):
                          # prepare inputs
                          a			   =self._prepare_for_class(_lowerCAmelCase ,    _lowerCAmelCase	)
                          a			   ={k: torch.tensor(v.tolist()	) for k, v in prepared_inputs_dict.items()}

                          # load corresponding PyTorch class
                          a			   =model_class.__name__[4:]  # Skip the "Flax" at the beginning
                          a			   =getattr(_lowerCAmelCase ,    _lowerCAmelCase	)

                          a   ,      a			   =pt_inputs["""input_ids"""].shape
                          a			   =np.random.randint(0 ,    seq_length - 1 ,    size=(batch_size,)	)
                          for batch_idx, start_index in enumerate(_lowerCAmelCase	):
                                 a			   =0
                                 a			   =1
                                 a			   =0
                                 a			   =1
                          a			   =pt_model_class(_lowerCAmelCase	).eval()
                          a			   =model_class(_lowerCAmelCase ,    dtype=jnp.floataa	)

                          a			   =convert_pytorch_state_dict_to_flax(pt_model.state_dict() ,    _lowerCAmelCase	)
                          a			   =fx_state

                          with torch.no_grad():
                                 a			   =pt_model(**_lowerCAmelCase	).to_tuple()

                          a			   =fx_model(**_lowerCAmelCase	).to_tuple()
                          self.assertEqual(len(_lowerCAmelCase	) ,    len(_lowerCAmelCase	) ,    """Output lengths differ between Flax and PyTorch"""	)
                          for fx_output, pt_output in zip(_lowerCAmelCase ,    _lowerCAmelCase	):
                                 self.assert_almost_equals(fx_output[:, -1] ,    pt_output[:, -1].numpy() ,    4E-2	)

                          with tempfile.TemporaryDirectory() as tmpdirname:
                                 pt_model.save_pretrained(_lowerCAmelCase	)
                                 a			   =model_class.from_pretrained(_lowerCAmelCase ,    from_pt=_lowerCAmelCase	)

                          a			   =fx_model_loaded(**_lowerCAmelCase	).to_tuple()
                          self.assertEqual(
                              len(_lowerCAmelCase	) ,    len(_lowerCAmelCase	) ,    """Output lengths differ between Flax and PyTorch"""	)
                          for fx_output_loaded, pt_output in zip(_lowerCAmelCase ,    _lowerCAmelCase	):
                                 self.assert_almost_equals(fx_output_loaded[:, -1] ,    pt_output[:, -1].numpy() ,    4E-2	)
     @is_pt_flax_cross_test
     def        lowerCAmelCase__  (				self	):
            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__	):
                          # prepare inputs
                          a			   =self._prepare_for_class(_lowerCAmelCase ,    _lowerCAmelCase	)
                          a			   ={k: torch.tensor(v.tolist()	) for k, v in prepared_inputs_dict.items()}

                          # load corresponding PyTorch class
                          a			   =model_class.__name__[4:]  # Skip the "Flax" at the beginning
                          a			   =getattr(_lowerCAmelCase ,    _lowerCAmelCase	)

                          a			   =pt_model_class(_lowerCAmelCase	).eval()
                          a			   =model_class(_lowerCAmelCase ,    dtype=jnp.floataa	)

                          a			   =load_flax_weights_in_pytorch_model(_lowerCAmelCase ,    fx_model.params	)
                          a   ,      a			   =pt_inputs["""input_ids"""].shape
                          a			   =np.random.randint(0 ,    seq_length - 1 ,    size=(batch_size,)	)
                          for batch_idx, start_index in enumerate(_lowerCAmelCase	):
                                 a			   =0
                                 a			   =1
                                 a			   =0
                                 a			   =1

                          # make sure weights are tied in PyTorch
                          pt_model.tie_weights()

                          with torch.no_grad():
                                 a			   =pt_model(**_lowerCAmelCase	).to_tuple()

                          a			   =fx_model(**_lowerCAmelCase	).to_tuple()
                          self.assertEqual(len(_lowerCAmelCase	) ,    len(_lowerCAmelCase	) ,    """Output lengths differ between Flax and PyTorch"""	)
                          for fx_output, pt_output in zip(_lowerCAmelCase ,    _lowerCAmelCase	):
                                 self.assert_almost_equals(fx_output[:, -1] ,    pt_output[:, -1].numpy() ,    4E-2	)

                          with tempfile.TemporaryDirectory() as tmpdirname:
                                 fx_model.save_pretrained(_lowerCAmelCase	)
                                 a			   =pt_model_class.from_pretrained(_lowerCAmelCase ,    from_flax=_lowerCAmelCase	)

                          with torch.no_grad():
                                 a			   =pt_model_loaded(**_lowerCAmelCase	).to_tuple()

                          self.assertEqual(
                              len(_lowerCAmelCase	) ,    len(_lowerCAmelCase	) ,    """Output lengths differ between Flax and PyTorch"""	)
                          for fx_output, pt_output in zip(_lowerCAmelCase ,    _lowerCAmelCase	):
                                 self.assert_almost_equals(fx_output[:, -1] ,    pt_output[:, -1].numpy() ,    4E-2	)






     @slow
     def        lowerCAmelCase__  (				self	):
            for model_class_name in self.all_model_classes:
                   a			   =model_class_name.from_pretrained("""gpt2""" ,    from_pt=_lowerCAmelCase	)
                   a			   =model(np.ones((1, 1)	)	)
                   self.assertIsNotNone(_lowerCAmelCase	)






import inspect
import unittest

from transformers import is_torch_available, is_vision_available
from transformers.models.auto import get_values
from transformers.testing_utils import require_torch, slow, torch_device

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


if is_torch_available():
    import torch

    from transformers import MODEL_MAPPING, PoolFormerConfig, PoolFormerForImageClassification, PoolFormerModel
    from transformers.models.poolformer.modeling_poolformer import POOLFORMER_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
    from PIL import Image

    from transformers import PoolFormerImageProcessor


class        UpperCAmelCase__				(							UpperCAmelCase__				):



     '''simple docstring'''
     def        lowerCAmelCase__  (				self	):
            a			   =self.config_class(**self.inputs_dict	)
            self.parent.assertTrue(hasattr(_lowerCAmelCase ,    """hidden_sizes"""	)	)
            self.parent.assertTrue(hasattr(_lowerCAmelCase ,    """num_encoder_blocks"""	)	)



class        UpperCAmelCase__				:



     '''simple docstring'''
     def __init__(				self ,    _lowerCAmelCase ,    _lowerCAmelCase=13 ,    _lowerCAmelCase=64 ,    _lowerCAmelCase=3 ,    _lowerCAmelCase=4 ,    _lowerCAmelCase=[2, 2, 2, 2] ,    _lowerCAmelCase=[8, 4, 2, 1] ,    _lowerCAmelCase=[16, 32, 64, 128] ,    _lowerCAmelCase=[1, 4, 8, 16] ,    _lowerCAmelCase=False ,    _lowerCAmelCase=True ,    _lowerCAmelCase="gelu" ,    _lowerCAmelCase=0.1 ,    _lowerCAmelCase=0.02 ,    _lowerCAmelCase=3 ,    _lowerCAmelCase=None ,    ):
            a			   =parent
            a			   =batch_size
            a			   =image_size
            a			   =num_channels
            a			   =num_encoder_blocks
            a			   =sr_ratios
            a			   =depths
            a			   =hidden_sizes
            a			   =downsampling_rates
            a			   =is_training
            a			   =use_labels
            a			   =hidden_act
            a			   =hidden_dropout_prob
            a			   =initializer_range
            a			   =num_labels
            a			   =scope
     def        lowerCAmelCase__  (				self	):
            a			   =floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]	)

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

            a			   =PoolFormerConfig(
                image_size=self.image_size ,    num_channels=self.num_channels ,    num_encoder_blocks=self.num_encoder_blocks ,    depths=self.depths ,    hidden_sizes=self.hidden_sizes ,    hidden_act=self.hidden_act ,    hidden_dropout_prob=self.hidden_dropout_prob ,    initializer_range=self.initializer_range ,    )

            return config, pixel_values, labels
     def        lowerCAmelCase__  (				self ,    _lowerCAmelCase ,    _lowerCAmelCase ,    _lowerCAmelCase	):
            a			   =PoolFormerModel(config=_lowerCAmelCase	)
            model.to(_lowerCAmelCase	)
            model.eval()
            a			   =model(_lowerCAmelCase	)
            a			   =a			   =self.image_size // 32.0
            self.parent.assertEqual(
                result.last_hidden_state.shape ,    (self.batch_size, self.hidden_sizes[-1], expected_height, expected_width)	)






     def        lowerCAmelCase__  (				self	):
            a			   =self.prepare_config_and_inputs()
            a   ,      a   ,      a			   =config_and_inputs
            a			   ={"""pixel_values""": pixel_values}
            return config, inputs_dict



@require_torch
class        UpperCAmelCase__				(							UpperCAmelCase__      ,		UpperCAmelCase__      ,		unittest.TestCase				):



     '''simple docstring'''
     _SCREAMING_SNAKE_CASE      :					Tuple     					=    (PoolFormerModel, PoolFormerForImageClassification) if is_torch_available() else ()
     _SCREAMING_SNAKE_CASE      :					Any     					=    (
         {"feature-extraction": PoolFormerModel, "image-classification": PoolFormerForImageClassification}
         if is_torch_available()
         else {}
     )

     _SCREAMING_SNAKE_CASE      :					Union[str, Any]     					=    False
     _SCREAMING_SNAKE_CASE      :					Tuple     					=    False
     _SCREAMING_SNAKE_CASE      :					str     					=    False
     _SCREAMING_SNAKE_CASE      :					Optional[int]     					=    False
     _SCREAMING_SNAKE_CASE      :					Optional[int]     					=    False
     def        lowerCAmelCase__  (				self	):
            a			   =PoolFormerModelTester(self	)
            a			   =PoolFormerConfigTester(self ,    config_class=_lowerCAmelCase	)
     def        lowerCAmelCase__  (				self	):
            self.config_tester.run_common_tests()
     def        lowerCAmelCase__  (				self	):
            a			   =self.model_tester.prepare_config_and_inputs()
            self.model_tester.create_and_check_model(*_lowerCAmelCase	)
     @unittest.skip("""PoolFormer does not use inputs_embeds"""	)
     def        lowerCAmelCase__  (				self	):
            pass
     @unittest.skip("""PoolFormer does not have get_input_embeddings method and get_output_embeddings methods"""	)
     def        lowerCAmelCase__  (				self	):
            pass
     def        lowerCAmelCase__  (				self	):
            def check_hidden_states_output(_lowerCAmelCase ,    _lowerCAmelCase ,    _lowerCAmelCase	):
                   a			   =model_class(_lowerCAmelCase	)
                   model.to(_lowerCAmelCase	)
                   model.eval()

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

                   a			   =outputs.hidden_states

                   a			   =self.model_tester.num_encoder_blocks
                   self.assertEqual(len(_lowerCAmelCase	) ,    _lowerCAmelCase	)

                   # verify the first hidden states (first block)
                   self.assertListEqual(
                       list(hidden_states[0].shape[-3:]	) ,    [
                           self.model_tester.hidden_sizes[0],
                           self.model_tester.image_size // 4,
                           self.model_tester.image_size // 4,
                       ] ,    )

            a   ,      a			   =self.model_tester.prepare_config_and_inputs_for_common()

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

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

                   check_hidden_states_output(_lowerCAmelCase ,    _lowerCAmelCase ,    _lowerCAmelCase	)
     def        lowerCAmelCase__  (				self	):
            if not self.model_tester.is_training:
                   return

            a   ,      a			   =self.model_tester.prepare_config_and_inputs_for_common()
            a			   =True

            for model_class in self.all_model_classes:
                   if model_class in get_values(_lowerCAmelCase	):
                          continue
                   a			   =model_class(_lowerCAmelCase	)
                   model.to(_lowerCAmelCase	)
                   model.train()
                   a			   =self._prepare_for_class(_lowerCAmelCase ,    _lowerCAmelCase ,    return_labels=_lowerCAmelCase	)
                   a			   =model(**_lowerCAmelCase	).loss
                   loss.backward()
     def        lowerCAmelCase__  (				self	):
            a   ,      a			   =self.model_tester.prepare_config_and_inputs_for_common()

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

                   a			   =["""pixel_values"""]
                   self.assertListEqual(arg_names[:1] ,    _lowerCAmelCase	)






     @slow
     def        lowerCAmelCase__  (				self	):
            for model_name in POOLFORMER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
                   a			   =PoolFormerModel.from_pretrained(_lowerCAmelCase	)
                   self.assertIsNotNone(_lowerCAmelCase	)





def 							lowerCamelCase    (							)->			Tuple:

       """simple docstring"""




       a			   =Image.open("""./tests/fixtures/tests_samples/COCO/000000039769.png"""      )
       return image






@require_torch
class        UpperCAmelCase__				(							unittest.TestCase				):



     '''simple docstring'''
     @slow
     def        lowerCAmelCase__  (				self	):
            a			   =PoolFormerImageProcessor()
            a			   =PoolFormerForImageClassification.from_pretrained("""sail/poolformer_s12"""	).to(_lowerCAmelCase	)

            a			   =image_processor(images=prepare_img() ,    return_tensors="""pt"""	).to(_lowerCAmelCase	)

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

            # verify the logits
            a			   =torch.Size((1, 1_000)	)
            self.assertEqual(outputs.logits.shape ,    _lowerCAmelCase	)

            a			   =torch.tensor([-0.61_13, 0.16_85, -0.04_92]	).to(_lowerCAmelCase	)
            self.assertTrue(torch.allclose(outputs.logits[0, :3] ,    _lowerCAmelCase ,    atol=1E-4	)	)
