





'''simple docstring'''







import tempfile
import unittest

import numpy as np

import transformers
from transformers import GPTaTokenizer, GPTJConfig, is_flax_available, is_torch_available
from transformers.testing_utils import is_pt_flax_cross_test, require_flax, tooslow

from ...generation.test_flax_utils import FlaxGenerationTesterMixin
from ...test_modeling_flax_common import FlaxModelTesterMixin, 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.gptj.modeling_flax_gptj import FlaxGPTJForCausalLM, FlaxGPTJModel

if is_torch_available():
    import torch



class 				__SCREAMING_SNAKE_CASE      :
      '''simple docstring'''



      def __init__(				self     :			List[Any] ,UpperCamelCase     :			List[Any] ,UpperCamelCase     :			Tuple=14 ,UpperCamelCase     :			Tuple=7 ,UpperCamelCase     :			str=True ,UpperCamelCase     :			str=True ,UpperCamelCase     :			Union[str, Any]=False ,UpperCamelCase     :			str=True ,UpperCamelCase     :			Tuple=99 ,UpperCamelCase     :			Union[str, Any]=32 ,UpperCamelCase     :			Tuple=4 ,UpperCamelCase     :			Union[str, Any]=4 ,UpperCamelCase     :			Optional[int]=4 ,UpperCamelCase     :			Union[str, Any]=37 ,UpperCamelCase     :			Optional[int]="gelu" ,UpperCamelCase     :			List[Any]=0.1 ,UpperCamelCase     :			str=0.1 ,UpperCamelCase     :			Union[str, Any]=512 ,UpperCamelCase     :			int=0.0_2 ,)		->	int:
           _lowercase				:   str	    =				parent
           _lowercase				:   Optional[Any]	    =				batch_size
           _lowercase				:   int	    =				seq_length
           _lowercase				:   Union[str, Any]	    =				is_training
           _lowercase				:   Tuple	    =				use_input_mask
           _lowercase				:   List[Any]	    =				use_token_type_ids
           _lowercase				:   int	    =				use_labels
           _lowercase				:   Any	    =				vocab_size
           _lowercase				:   Any	    =				hidden_size
           _lowercase				:   Optional[int]	    =				rotary_dim
           _lowercase				:   Dict	    =				num_hidden_layers
           _lowercase				:   List[Any]	    =				num_attention_heads
           _lowercase				:   Tuple	    =				intermediate_size
           _lowercase				:   Any	    =				hidden_act
           _lowercase				:   Dict	    =				hidden_dropout_prob
           _lowercase				:   List[Any]	    =				attention_probs_dropout_prob
           _lowercase				:   Dict	    =				max_position_embeddings
           _lowercase				:   int	    =				initializer_range
           _lowercase				:   List[str]	    =				None
           _lowercase				:   List[Any]	    =				vocab_size - 1
           _lowercase				:   Any	    =				vocab_size - 1
           _lowercase				:   Optional[int]	    =				vocab_size - 1

      def 	_lowerCamelCase		(				self     :			int	)		->	Any:
           _lowercase				:   str	    =				ids_tensor([self.batch_size, self.seq_length] ,self.vocab_size	)

           _lowercase				:   Dict	    =				None
           if self.use_input_mask:
                _lowercase				:   str	    =				random_attention_mask([self.batch_size, self.seq_length]	)

           _lowercase				:   Union[str, Any]	    =				GPTJConfig(
               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=__a ,bos_token_id=self.bos_token_id ,eos_token_id=self.eos_token_id ,pad_token_id=self.pad_token_id ,rotary_dim=self.rotary_dim ,)

           return (config, input_ids, input_mask)

      def 	_lowerCamelCase		(				self     :			Union[str, Any]	)		->	str:
           _lowercase				:   str	    =				self.prepare_config_and_inputs()
           _lowercase    ,			_lowercase    ,			_lowercase				:   List[Any]	    =				config_and_inputs
           _lowercase				:   Dict	    =				{'input_ids': input_ids, 'attention_mask': attention_mask}
           return config, inputs_dict

      def 	_lowerCamelCase		(				self     :			Union[str, Any] ,UpperCamelCase     :			Optional[Any] ,UpperCamelCase     :			Union[str, Any] ,UpperCamelCase     :			List[Any] ,UpperCamelCase     :			Union[str, Any]	)		->	List[str]:
           _lowercase				:   Union[str, Any]	    =				20
           _lowercase				:   Tuple	    =				model_class_name(__a	)

           _lowercase				:   str	    =				model.init_cache(input_ids.shape[0] ,__a	)
           _lowercase				:   Dict	    =				jnp.ones((input_ids.shape[0], max_decoder_length) ,dtype='i4'	)

           _lowercase				:   str	    =				jnp.broadcast_to(
               jnp.arange(input_ids.shape[-1] - 1	)[None, :] ,(input_ids.shape[0], input_ids.shape[-1] - 1)	)
           _lowercase				:   List[str]	    =				model(
               input_ids[:, :-1] ,attention_mask=__a ,past_key_values=__a ,position_ids=__a ,)

           _lowercase				:   Union[str, Any]	    =				jnp.array(input_ids.shape[0] * [[input_ids.shape[-1] - 1]] ,dtype='i4'	)
           _lowercase				:   str	    =				model(
               input_ids[:, -1:] ,attention_mask=__a ,past_key_values=outputs_cache.past_key_values ,position_ids=__a ,)

           _lowercase				:   List[str]	    =				model(__a	)

           _lowercase				:   Any	    =				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     :			int ,UpperCamelCase     :			Dict ,UpperCamelCase     :			int ,UpperCamelCase     :			Optional[Any] ,UpperCamelCase     :			List[Any]	)		->	Dict:
           _lowercase				:   int	    =				20
           _lowercase				:   Optional[int]	    =				model_class_name(__a	)

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

           _lowercase				:   int	    =				model.init_cache(input_ids.shape[0] ,__a	)
           _lowercase				:   Optional[int]	    =				jnp.broadcast_to(
               jnp.arange(input_ids.shape[-1] - 1	)[None, :] ,(input_ids.shape[0], input_ids.shape[-1] - 1)	)

           _lowercase				:   str	    =				model(
               input_ids[:, :-1] ,attention_mask=__a ,past_key_values=__a ,position_ids=__a ,)
           _lowercase				:   Any	    =				jnp.array(input_ids.shape[0] * [[input_ids.shape[-1] - 1]] ,dtype='i4'	)
           _lowercase				:   Optional[int]	    =				model(
               input_ids[:, -1:] ,past_key_values=outputs_cache.past_key_values ,attention_mask=__a ,position_ids=__a ,)

           _lowercase				:   Union[str, Any]	    =				model(__a ,attention_mask=__a	)

           _lowercase				:   List[Any]	    =				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 				__SCREAMING_SNAKE_CASE      (							lowerCAmelCase__				, lowerCAmelCase__				, unittest.TestCase   ):
      '''simple docstring'''



      lowerCAmelCase__							:   int					=		(FlaxGPTJModel, FlaxGPTJForCausalLM) if is_flax_available() else ()
      lowerCAmelCase__							:   Dict					=		(FlaxGPTJForCausalLM,) if is_flax_available() else ()

      def 	_lowerCamelCase		(				self     :			int	)		->	Dict:
           _lowercase				:   Optional[int]	    =				FlaxGPTJModelTester(self	)

      def 	_lowerCamelCase		(				self     :			List[Any]	)		->	Optional[Any]:
           for model_class_name in self.all_model_classes:
                _lowercase    ,			_lowercase    ,			_lowercase				:   Any	    =				self.model_tester.prepare_config_and_inputs()
                self.model_tester.check_use_cache_forward(__a ,__a ,__a ,__a	)

      def 	_lowerCamelCase		(				self     :			Any	)		->	Union[str, Any]:
           for model_class_name in self.all_model_classes:
                _lowercase    ,			_lowercase    ,			_lowercase				:   str	    =				self.model_tester.prepare_config_and_inputs()
                self.model_tester.check_use_cache_forward_with_attn_mask(
                    __a ,__a ,__a ,__a	)

      @tooslow
      def 	_lowerCamelCase		(				self     :			Tuple	)		->	Tuple:
           _lowercase				:   List[str]	    =				GPTaTokenizer.from_pretrained('gpt2' ,pad_token='<|endoftext|>' ,padding_side='left'	)
           _lowercase				:   List[Any]	    =				tokenizer(['Hello this is a long string', 'Hey'] ,return_tensors='np' ,padding=__a ,truncation=__a	)

           _lowercase				:   Union[str, Any]	    =				FlaxGPTJForCausalLM.from_pretrained('EleutherAI/gpt-j-6B'	)
           _lowercase				:   Dict	    =				False
           _lowercase				:   Optional[int]	    =				model.config.eos_token_id

           _lowercase				:   List[Any]	    =				jax.jit(model.generate	)

           _lowercase				:   Tuple	    =				jit_generate(
               inputs['input_ids'] ,attention_mask=inputs['attention_mask'] ,pad_token_id=tokenizer.pad_token_id	).sequences

           _lowercase				:   Any	    =				tokenizer.batch_decode(__a ,skip_special_tokens=__a	)

           _lowercase				:   Union[str, Any]	    =				[
               'Hello this is a long string of text.\n\nI\'m trying to get the text of the',
               'Hey, I\'m a little late to the party. I\'m going to',
           ]

           self.assertListEqual(__a ,__a	)

      @is_pt_flax_cross_test
      def 	_lowerCamelCase		(				self     :			str	)		->	str:
           _lowercase    ,			_lowercase				:   Optional[int]	    =				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
                     _lowercase				:   str	    =				self._prepare_for_class(__a ,__a	)
                     _lowercase				:   Dict	    =				{k: torch.tensor(v.tolist()	) for k, v in prepared_inputs_dict.items()}

                     # load corresponding PyTorch class
                     _lowercase				:   Union[str, Any]	    =				model_class.__name__[4:]  # Skip the "Flax" at the beginning
                     _lowercase				:   List[str]	    =				getattr(__a ,__a	)

                     _lowercase    ,			_lowercase				:   Dict	    =				pt_inputs['input_ids'].shape
                     _lowercase				:   Optional[int]	    =				np.random.randint(0 ,seq_length - 1 ,size=(batch_size,)	)
                     for batch_idx, start_index in enumerate(__a	):
                          _lowercase				:   List[str]	    =				0
                          _lowercase				:   Optional[Any]	    =				1
                          _lowercase				:   Tuple	    =				0
                          _lowercase				:   Union[str, Any]	    =				1
                     _lowercase				:   str	    =				pt_model_class(__a	).eval()
                     _lowercase				:   int	    =				model_class(__a ,dtype=jnp.floataa	)

                     _lowercase				:   int	    =				convert_pytorch_state_dict_to_flax(pt_model.state_dict() ,__a	)
                     _lowercase				:   List[str]	    =				fx_state

                     with torch.no_grad():
                          _lowercase				:   Any	    =				pt_model(**__a	).to_tuple()

                     _lowercase				:   Dict	    =				fx_model(**__a	).to_tuple()
                     self.assertEqual(len(__a	) ,len(__a	) ,'Output lengths differ between Flax and PyTorch'	)
                     for fx_output, pt_output in zip(__a ,__a	):
                          self.assert_almost_equals(fx_output[:, -1] ,pt_output[:, -1].numpy() ,4e-2	)

                     with tempfile.TemporaryDirectory() as tmpdirname:
                          pt_model.save_pretrained(__a	)
                          _lowercase				:   Optional[Any]	    =				model_class.from_pretrained(__a ,from_pt=__a	)

                     _lowercase				:   Optional[Any]	    =				fx_model_loaded(**__a	).to_tuple()
                     self.assertEqual(
                         len(__a	) ,len(__a	) ,'Output lengths differ between Flax and PyTorch'	)
                     for fx_output_loaded, pt_output in zip(__a ,__a	):
                          self.assert_almost_equals(fx_output_loaded[:, -1] ,pt_output[:, -1].numpy() ,4e-2	)

      @is_pt_flax_cross_test
      def 	_lowerCamelCase		(				self     :			Optional[Any]	)		->	List[Any]:
           _lowercase    ,			_lowercase				:   str	    =				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
                     _lowercase				:   Any	    =				self._prepare_for_class(__a ,__a	)
                     _lowercase				:   Dict	    =				{k: torch.tensor(v.tolist()	) for k, v in prepared_inputs_dict.items()}

                     # load corresponding PyTorch class
                     _lowercase				:   List[str]	    =				model_class.__name__[4:]  # Skip the "Flax" at the beginning
                     _lowercase				:   Dict	    =				getattr(__a ,__a	)

                     _lowercase				:   Any	    =				pt_model_class(__a	).eval()
                     _lowercase				:   int	    =				model_class(__a ,dtype=jnp.floataa	)

                     _lowercase				:   Optional[Any]	    =				load_flax_weights_in_pytorch_model(__a ,fx_model.params	)
                     _lowercase    ,			_lowercase				:   int	    =				pt_inputs['input_ids'].shape
                     _lowercase				:   Optional[int]	    =				np.random.randint(0 ,seq_length - 1 ,size=(batch_size,)	)
                     for batch_idx, start_index in enumerate(__a	):
                          _lowercase				:   Optional[int]	    =				0
                          _lowercase				:   Optional[int]	    =				1
                          _lowercase				:   Any	    =				0
                          _lowercase				:   Union[str, Any]	    =				1

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

                     with torch.no_grad():
                          _lowercase				:   List[str]	    =				pt_model(**__a	).to_tuple()

                     _lowercase				:   List[str]	    =				fx_model(**__a	).to_tuple()
                     self.assertEqual(len(__a	) ,len(__a	) ,'Output lengths differ between Flax and PyTorch'	)
                     for fx_output, pt_output in zip(__a ,__a	):
                          self.assert_almost_equals(fx_output[:, -1] ,pt_output[:, -1].numpy() ,4e-2	)

                     with tempfile.TemporaryDirectory() as tmpdirname:
                          fx_model.save_pretrained(__a	)
                          _lowercase				:   Any	    =				pt_model_class.from_pretrained(__a ,from_flax=__a	)

                     with torch.no_grad():
                          _lowercase				:   Union[str, Any]	    =				pt_model_loaded(**__a	).to_tuple()

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



      @tooslow
      def 	_lowerCamelCase		(				self     :			List[Any]	)		->	Dict:
           for model_class_name in self.all_model_classes:
                _lowercase				:   Dict	    =				model_class_name.from_pretrained('EleutherAI/gpt-j-6B'	)
                _lowercase				:   Dict	    =				model(np.ones((1, 1)	)	)
                self.assertIsNotNone(__a	)



"""simple docstring"""


import unittest

from transformers import MODEL_FOR_VISUAL_QUESTION_ANSWERING_MAPPING, is_vision_available
from transformers.pipelines import pipeline
from transformers.testing_utils import (
    is_pipeline_test,
    nested_simplify,
    require_tf,
    require_torch,
    require_vision,
    slow,
)

from .test_pipelines_common import ANY


if is_vision_available():
	from PIL import Image
else:
	class     _UpperCamelCase    :
								'''simple docstring'''







								@staticmethod
								def      snake_case							(	*__a       ,  **__a      ):
											pass
@is_pipeline_test
@require_torch
@require_vision
class     _UpperCamelCase    (  unittest.TestCase						):
							'''simple docstring'''







							__UpperCAmelCase		:       str							 =MODEL_FOR_VISUAL_QUESTION_ANSWERING_MAPPING

							def      snake_case							(	self       ,  __a       ,  __a       ,  __a      ):
										__lowerCAmelCase							       =     pipeline("visual-question-answering"       ,  model="hf-internal-testing/tiny-vilt-random-vqa"      )
										__lowerCAmelCase							       =     [
										    {
										        "image": Image.open("./tests/fixtures/tests_samples/COCO/000000039769.png"      ),
										        "question": "How many cats are there?",
										    },
										    {
										        "image": "./tests/fixtures/tests_samples/COCO/000000039769.png",
										        "question": "How many cats are there?",
										    },
										]
										return vqa_pipeline, examples

							def      snake_case							(	self       ,  __a       ,  __a      ):
										__lowerCAmelCase							       =     vqa_pipeline(__a       ,  top_k=1      )
										self.assertEqual(
										    __a       ,  [
										        [{"score": ANY(__a      ), "answer": ANY(__a      )}],
										        [{"score": ANY(__a      ), "answer": ANY(__a      )}],
										    ]       ,  )

							@require_torch
							def      snake_case							(	self      ):
										__lowerCAmelCase							       =     pipeline("visual-question-answering"       ,  model="hf-internal-testing/tiny-vilt-random-vqa"      )
										__lowerCAmelCase							       =     "./tests/fixtures/tests_samples/COCO/000000039769.png"
										__lowerCAmelCase							       =     "How many cats are there?"

										__lowerCAmelCase							       =     vqa_pipeline(image=__a       ,  question="How many cats are there?"       ,  top_k=2      )
										self.assertEqual(
										    __a       ,  [{"score": ANY(__a      ), "answer": ANY(__a      )}, {"score": ANY(__a      ), "answer": ANY(__a      )}]      )

										__lowerCAmelCase							       =     vqa_pipeline({"image": image, "question": question}       ,  top_k=2      )
										self.assertEqual(
										    __a       ,  [{"score": ANY(__a      ), "answer": ANY(__a      )}, {"score": ANY(__a      ), "answer": ANY(__a      )}]      )

							@slow
							@require_torch
							def      snake_case							(	self      ):
										__lowerCAmelCase							       =     pipeline("visual-question-answering"       ,  model="dandelin/vilt-b32-finetuned-vqa"      )
										__lowerCAmelCase							       =     "./tests/fixtures/tests_samples/COCO/000000039769.png"
										__lowerCAmelCase							       =     "How many cats are there?"

										__lowerCAmelCase							       =     vqa_pipeline(image=__a       ,  question=__a       ,  top_k=2      )
										self.assertEqual(
										    nested_simplify(__a       ,  decimals=4      )       ,  [{"score": 0.8_7_9_9, "answer": "2"}, {"score": 0.2_9_6, "answer": "1"}]      )

										__lowerCAmelCase							       =     vqa_pipeline({"image": image, "question": question}       ,  top_k=2      )
										self.assertEqual(
										    nested_simplify(__a       ,  decimals=4      )       ,  [{"score": 0.8_7_9_9, "answer": "2"}, {"score": 0.2_9_6, "answer": "1"}]      )

										__lowerCAmelCase							       =     vqa_pipeline(
										    [{"image": image, "question": question}, {"image": image, "question": question}]       ,  top_k=2      )
										self.assertEqual(
										    nested_simplify(__a       ,  decimals=4      )       ,  [[{"score": 0.8_7_9_9, "answer": "2"}, {"score": 0.2_9_6, "answer": "1"}]] * 2       ,  )



							@require_tf
							@unittest.skip("Visual question answering not implemented in TF"      )
							def      snake_case							(	self      ):
										pass




