






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 							a_				:


							'''simple docstring'''






							def __init__(	self    ,    lowercase_    ,    lowercase_=1_4    ,    lowercase_=7    ,    lowercase_=True    ,    lowercase_=True    ,    lowercase_=False    ,    lowercase_=True    ,    lowercase_=9_9    ,    lowercase_=3_2    ,    lowercase_=5    ,    lowercase_=4    ,    lowercase_=3_7    ,    lowercase_="gelu"    ,    lowercase_=0.1    ,    lowercase_=0.1    ,    lowercase_=5_1_2    ,    lowercase_=0.02    ,    )      ->      List[str]:
														'''simple docstring'''

														lowerCAmelCase_       		=  parent
														lowerCAmelCase_       		=  batch_size
														lowerCAmelCase_       		=  seq_length
														lowerCAmelCase_       		=  is_training
														lowerCAmelCase_       		=  use_input_mask
														lowerCAmelCase_       		=  use_token_type_ids
														lowerCAmelCase_       		=  use_labels
														lowerCAmelCase_       		=  vocab_size
														lowerCAmelCase_       		=  hidden_size
														lowerCAmelCase_       		=  num_hidden_layers
														lowerCAmelCase_       		=  num_attention_heads
														lowerCAmelCase_       		=  intermediate_size
														lowerCAmelCase_       		=  hidden_act
														lowerCAmelCase_       		=  hidden_dropout_prob
														lowerCAmelCase_       		=  attention_probs_dropout_prob
														lowerCAmelCase_       		=  max_position_embeddings
														lowerCAmelCase_       		=  initializer_range
														lowerCAmelCase_       		=  None
														lowerCAmelCase_       		=  vocab_size - 1
														lowerCAmelCase_       		=  vocab_size - 1
														lowerCAmelCase_       		=  vocab_size - 1

							def   _lowercase (	self     )      ->      Union[str, Any]:
														'''simple docstring'''

														lowerCAmelCase_       		=  ids_tensor([self.batch_size, self.seq_length]    ,    self.vocab_size     )

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

														lowerCAmelCase_       		=  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   _lowercase (	self     )      ->      Dict:
														'''simple docstring'''

														lowerCAmelCase_       		=  self.prepare_config_and_inputs()
														lowerCAmelCase_       		=  config_and_inputs
														lowerCAmelCase_       		=  {"""input_ids""": input_ids, """attention_mask""": attention_mask}
														return config, inputs_dict

							def   _lowercase (	self     )      ->      Optional[Any]:
														'''simple docstring'''

														lowerCAmelCase_       		=  self.prepare_config_and_inputs()

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

														return (
														    config,
														    input_ids,
														    attention_mask,
														    encoder_hidden_states,
														    encoder_attention_mask,
														)

							def   _lowercase (	self    ,    lowercase_    ,    lowercase_    ,    lowercase_    ,    lowercase_     )      ->      Optional[Any]:
														'''simple docstring'''

														lowerCAmelCase_       		=  2_0
														lowerCAmelCase_       		=  model_class_name(_lowerCamelCase     )

														lowerCAmelCase_       		=  model.init_cache(input_ids.shape[0]    ,    _lowerCamelCase     )
														lowerCAmelCase_       		=  jnp.ones((input_ids.shape[0], max_decoder_length)    ,    dtype='i4'     )

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

														lowerCAmelCase_       		=  jnp.array(input_ids.shape[0] * [[input_ids.shape[-1] - 1]]    ,    dtype='i4'     )
														lowerCAmelCase_       		=  model(
														    input_ids[:, -1:]    ,    attention_mask=_lowerCamelCase    ,    past_key_values=outputs_cache.past_key_values    ,    position_ids=_lowerCamelCase    ,    )

														lowerCAmelCase_       		=  model(_lowerCamelCase     )

														lowerCAmelCase_       		=  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   _lowercase (	self    ,    lowercase_    ,    lowercase_    ,    lowercase_    ,    lowercase_     )      ->      Optional[Any]:
														'''simple docstring'''

														lowerCAmelCase_       		=  2_0
														lowerCAmelCase_       		=  model_class_name(_lowerCamelCase     )

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

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

														lowerCAmelCase_       		=  model(
														    input_ids[:, :-1]    ,    attention_mask=_lowerCamelCase    ,    past_key_values=_lowerCamelCase    ,    position_ids=_lowerCamelCase    ,    )
														lowerCAmelCase_       		=  jnp.array(input_ids.shape[0] * [[input_ids.shape[-1] - 1]]    ,    dtype='i4'     )
														lowerCAmelCase_       		=  model(
														    input_ids[:, -1:]    ,    past_key_values=outputs_cache.past_key_values    ,    attention_mask=_lowerCamelCase    ,    position_ids=_lowerCamelCase    ,    )

														lowerCAmelCase_       		=  model(_lowerCamelCase    ,    attention_mask=_lowerCamelCase     )

														lowerCAmelCase_       		=  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 							a_				(		lowerCAmelCase__			,			lowerCAmelCase__			,			unittest.TestCase	):


							'''simple docstring'''






							__a:			Optional[Any]									=					(FlaxGPTaModel, FlaxGPTaLMHeadModel) if is_flax_available() else ()
							__a:			int									=					(FlaxGPTaLMHeadModel,) if is_flax_available() else ()

							def   _lowercase (	self     )      ->      str:
														'''simple docstring'''

														lowerCAmelCase_       		=  FlaxGPTaModelTester(self     )

							def   _lowercase (	self     )      ->      Tuple:
														'''simple docstring'''

														for model_class_name in self.all_model_classes:
																					lowerCAmelCase_       		=  self.model_tester.prepare_config_and_inputs()
																					self.model_tester.check_use_cache_forward(_lowerCamelCase    ,    _lowerCamelCase    ,    _lowerCamelCase    ,    _lowerCamelCase     )

							def   _lowercase (	self     )      ->      List[Any]:
														'''simple docstring'''

														for model_class_name in self.all_model_classes:
																					lowerCAmelCase_       		=  self.model_tester.prepare_config_and_inputs()
																					self.model_tester.check_use_cache_forward_with_attn_mask(
																					    _lowerCamelCase    ,    _lowerCamelCase    ,    _lowerCamelCase    ,    _lowerCamelCase     )

							@slow
							def   _lowercase (	self     )      ->      Any:
														'''simple docstring'''

														lowerCAmelCase_       		=  GPTaTokenizer.from_pretrained('gpt2'    ,    pad_token='</s>'    ,    padding_side='left'     )
														lowerCAmelCase_       		=  tokenizer(['Hello this is a long string', 'Hey']    ,    return_tensors='np'    ,    padding=_lowerCamelCase    ,    truncation=_lowerCamelCase     )

														lowerCAmelCase_       		=  FlaxGPTaLMHeadModel.from_pretrained('gpt2'     )
														lowerCAmelCase_       		=  False
														lowerCAmelCase_       		=  model.config.eos_token_id

														lowerCAmelCase_       		=  jax.jit(model.generate     )

														lowerCAmelCase_       		=  jit_generate(inputs['input_ids']    ,    attention_mask=inputs['attention_mask']     ).sequences

														lowerCAmelCase_       		=  tokenizer.batch_decode(_lowerCamelCase    ,    skip_special_tokens=_lowerCamelCase     )

														lowerCAmelCase_       		=  [
														    """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   _lowercase (	self     )      ->      str:
														'''simple docstring'''

														lowerCAmelCase_       		=  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
																												lowerCAmelCase_       		=  self._prepare_for_class(_lowerCamelCase    ,    _lowerCamelCase     )
																												lowerCAmelCase_       		=  {k: torch.tensor(v.tolist()     ) for k, v in prepared_inputs_dict.items()}

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

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

																												lowerCAmelCase_       		=  convert_pytorch_state_dict_to_flax(pt_model.state_dict()    ,    _lowerCamelCase     )
																												lowerCAmelCase_       		=  fx_state

																												with torch.no_grad():
																																			lowerCAmelCase_       		=  pt_model(**_lowerCamelCase     ).to_tuple()

																												lowerCAmelCase_       		=  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     )
																																			lowerCAmelCase_       		=  model_class.from_pretrained(_lowerCamelCase    ,    from_pt=_lowerCamelCase     )

																												lowerCAmelCase_       		=  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   _lowercase (	self     )      ->      Dict:
														'''simple docstring'''

														lowerCAmelCase_       		=  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
																												lowerCAmelCase_       		=  self._prepare_for_class(_lowerCamelCase    ,    _lowerCamelCase     )
																												lowerCAmelCase_       		=  {k: torch.tensor(v.tolist()     ) for k, v in prepared_inputs_dict.items()}

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

																												lowerCAmelCase_       		=  pt_model_class(_lowerCamelCase     ).eval()
																												lowerCAmelCase_       		=  model_class(_lowerCamelCase    ,    dtype=jnp.floataa     )

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

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

																												with torch.no_grad():
																																			lowerCAmelCase_       		=  pt_model(**_lowerCamelCase     ).to_tuple()

																												lowerCAmelCase_       		=  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     )
																																			lowerCAmelCase_       		=  pt_model_class.from_pretrained(_lowerCamelCase    ,    from_flax=_lowerCamelCase     )

																												with torch.no_grad():
																																			lowerCAmelCase_       		=  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   _lowercase (	self     )      ->      Optional[Any]:
														'''simple docstring'''

														for model_class_name in self.all_model_classes:
																					lowerCAmelCase_       		=  model_class_name.from_pretrained('gpt2'    ,    from_pt=_lowerCamelCase     )
																					lowerCAmelCase_       		=  model(np.ones((1, 1)     )     )
																					self.assertIsNotNone(_lowerCamelCase     )














import warnings

from ...processing_utils import ProcessorMixin
from ...tokenization_utils_base import BatchEncoding




class 							a_				(		a_	):


							'''simple docstring'''






							__a:			int									=					['''image_processor''', '''tokenizer''']
							__a:			int									=					'''VideoMAEImageProcessor'''
							__a:			Dict									=					('''CLIPTokenizer''', '''CLIPTokenizerFast''')

							def __init__(	self    ,    lowercase_=None    ,    lowercase_=None    ,    **lowercase_     )      ->      Optional[Any]:
														'''simple docstring'''

														lowerCAmelCase_       		=  None
														if "feature_extractor" in kwargs:
																					warnings.warn(
																					    'The `feature_extractor` argument is deprecated and will be removed in v5, use `image_processor`'
																					    ' instead.'    ,    lowercase_    ,    )
																					lowerCAmelCase_       		=  kwargs.pop('feature_extractor'     )

														lowerCAmelCase_       		=  image_processor if image_processor is not None else feature_extractor
														if image_processor is None:
																					raise ValueError('You need to specify an `image_processor`.'     )
														if tokenizer is None:
																					raise ValueError('You need to specify a `tokenizer`.'     )

														super().__init__(lowercase_    ,    lowercase_     )
														lowerCAmelCase_       		=  self.image_processor

							def __call__(	self    ,    lowercase_=None    ,    lowercase_=None    ,    lowercase_=None    ,    **lowercase_     )      ->      Optional[int]:
														'''simple docstring'''

														if text is None and videos is None:
																					raise ValueError('You have to specify either text or videos. Both cannot be none.'     )

														if text is not None:
																					lowerCAmelCase_       		=  self.tokenizer(lowercase_    ,    return_tensors=lowercase_    ,    **lowercase_     )

														if videos is not None:
																					lowerCAmelCase_       		=  self.image_processor(lowercase_    ,    return_tensors=lowercase_    ,    **lowercase_     )

														if text is not None and videos is not None:
																					lowerCAmelCase_       		=  image_features.pixel_values
																					return encoding
														elif text is not None:
																					return encoding
														else:
																					return BatchEncoding(data=dict(**lowercase_     )    ,    tensor_type=lowercase_     )

							def   _lowercase (	self    ,    *lowercase_    ,    **lowercase_     )      ->      Union[str, Any]:
														'''simple docstring'''

														return self.tokenizer.batch_decode(*lowercase_    ,    **lowercase_     )

							def   _lowercase (	self    ,    *lowercase_    ,    **lowercase_     )      ->      str:
														'''simple docstring'''

														return self.tokenizer.decode(*lowercase_    ,    **lowercase_     )

							@property
							def   _lowercase (	self     )      ->      Any:
														'''simple docstring'''

														return ["input_ids", "attention_mask", "position_ids", "pixel_values"]

							@property
							def   _lowercase (	self     )      ->      List[str]:
														'''simple docstring'''

														warnings.warn(
														    '`feature_extractor_class` is deprecated and will be removed in v5. Use `image_processor_class` instead.'    ,    lowercase_    ,    )
														return self.image_processor_class






							@property
							def   _lowercase (	self     )      ->      int:
														'''simple docstring'''

														warnings.warn(
														    '`feature_extractor` is deprecated and will be removed in v5. Use `image_processor` instead.'    ,    lowercase_    ,    )
														return self.image_processor



