






"""simple docstring"""







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       ,  snake_case__       ,  snake_case__=14       ,  snake_case__=7       ,  snake_case__=True       ,  snake_case__=True       ,  snake_case__=False       ,  snake_case__=True       ,  snake_case__=99       ,  snake_case__=32       ,  snake_case__=5       ,  snake_case__=4       ,  snake_case__=37       ,  snake_case__="gelu"       ,  snake_case__=0.1       ,  snake_case__=0.1       ,  snake_case__=512       ,  snake_case__=0.02       ,  ):
												'''simple docstring'''
												lowercase__				:      Tuple=       parent
												lowercase__				:      Union[str, Any]=       batch_size
												lowercase__				:      Any=       seq_length
												lowercase__				:      Tuple=       is_training
												lowercase__				:      Union[str, Any]=       use_input_mask
												lowercase__				:      int=       use_token_type_ids
												lowercase__				:      int=       use_labels
												lowercase__				:      Union[str, Any]=       vocab_size
												lowercase__				:      Tuple=       hidden_size
												lowercase__				:      List[str]=       num_hidden_layers
												lowercase__				:      Tuple=       num_attention_heads
												lowercase__				:      Any=       intermediate_size
												lowercase__				:      str=       hidden_act
												lowercase__				:      int=       hidden_dropout_prob
												lowercase__				:      Optional[Any]=       attention_probs_dropout_prob
												lowercase__				:      List[str]=       max_position_embeddings
												lowercase__				:      List[str]=       initializer_range
												lowercase__				:      int=       None
												lowercase__				:      int=       vocab_size - 1
												lowercase__				:      Any=       vocab_size - 1
												lowercase__				:      List[str]=       vocab_size - 1




						def 		UpperCAmelCase_       (   self			):
												'''simple docstring'''
												lowercase__				:      List[Any]=       ids_tensor([self.batch_size, self.seq_length]       ,  self.vocab_size			)

												lowercase__				:      List[str]=       None
												if self.use_input_mask:
																		lowercase__				:      Optional[Any]=       random_attention_mask([self.batch_size, self.seq_length]			)

												lowercase__				:      Optional[Any]=       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=snake_case__       ,  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 		UpperCAmelCase_       (   self			):
												'''simple docstring'''
												lowercase__				:      List[str]=       self.prepare_config_and_inputs()
												lowercase__,	lowercase__,	lowercase__				:      Dict=       config_and_inputs
												lowercase__				:      Optional[int]=       {"input_ids": input_ids, "attention_mask": attention_mask}
												return config, inputs_dict




						def 		UpperCAmelCase_       (   self			):
												'''simple docstring'''
												lowercase__,	lowercase__,	lowercase__				:      Dict=       self.prepare_config_and_inputs()

												lowercase__				:      Optional[Any]=       floats_tensor([self.batch_size, self.seq_length, self.hidden_size]			)
												lowercase__				:      Optional[Any]=       ids_tensor([self.batch_size, self.seq_length]       ,  vocab_size=2			)

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




						def 		UpperCAmelCase_       (   self       ,  snake_case__       ,  snake_case__       ,  snake_case__       ,  snake_case__			):
												'''simple docstring'''
												lowercase__				:      Any=       20
												lowercase__				:      Dict=       model_class_name(snake_case__			)

												lowercase__				:      Dict=       model.init_cache(input_ids.shape[0]       ,  snake_case__			)
												lowercase__				:      List[Any]=       jnp.ones((input_ids.shape[0], max_decoder_length)       ,  dtype="i4"			)

												lowercase__				:      Optional[int]=       jnp.broadcast_to(
												    jnp.arange(input_ids.shape[-1] - 1			)[None, :]       ,  (input_ids.shape[0], input_ids.shape[-1] - 1)			)
												lowercase__				:      int=       model(
												    input_ids[:, :-1]       ,  attention_mask=snake_case__       ,  past_key_values=snake_case__       ,  position_ids=snake_case__       ,  )

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

												lowercase__				:      int=       model(snake_case__			)

												lowercase__				:      Dict=       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 		UpperCAmelCase_       (   self       ,  snake_case__       ,  snake_case__       ,  snake_case__       ,  snake_case__			):
												'''simple docstring'''
												lowercase__				:      List[str]=       20
												lowercase__				:      Tuple=       model_class_name(snake_case__			)

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

												lowercase__				:      str=       model.init_cache(input_ids.shape[0]       ,  snake_case__			)
												lowercase__				:      List[str]=       jnp.broadcast_to(
												    jnp.arange(input_ids.shape[-1] - 1			)[None, :]       ,  (input_ids.shape[0], input_ids.shape[-1] - 1)			)

												lowercase__				:      Optional[int]=       model(
												    input_ids[:, :-1]       ,  attention_mask=snake_case__       ,  past_key_values=snake_case__       ,  position_ids=snake_case__       ,  )
												lowercase__				:      Optional[int]=       jnp.array(input_ids.shape[0] * [[input_ids.shape[-1] - 1]]       ,  dtype="i4"			)
												lowercase__				:      str=       model(
												    input_ids[:, -1:]       ,  past_key_values=outputs_cache.past_key_values       ,  attention_mask=snake_case__       ,  position_ids=snake_case__       ,  )

												lowercase__				:      Union[str, Any]=       model(snake_case__       ,  attention_mask=snake_case__			)

												lowercase__				:      str=       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(   SCREAMING_SNAKE_CASE__  ,  SCREAMING_SNAKE_CASE__  ,  unittest.TestCase							):




						"""simple docstring"""




						__lowerCamelCase				       =					(FlaxGPTaModel, FlaxGPTaLMHeadModel) if is_flax_available() else ()
						__lowerCamelCase				       =					(FlaxGPTaLMHeadModel,) if is_flax_available() else ()




						def 		UpperCAmelCase_       (   self			):
												'''simple docstring'''
												lowercase__				:      Union[str, Any]=       FlaxGPTaModelTester(self			)




						def 		UpperCAmelCase_       (   self			):
												'''simple docstring'''
												for model_class_name in self.all_model_classes:
																		lowercase__,	lowercase__,	lowercase__				:      int=       self.model_tester.prepare_config_and_inputs()
																		self.model_tester.check_use_cache_forward(snake_case__       ,  snake_case__       ,  snake_case__       ,  snake_case__			)




						def 		UpperCAmelCase_       (   self			):
												'''simple docstring'''
												for model_class_name in self.all_model_classes:
																		lowercase__,	lowercase__,	lowercase__				:      Optional[int]=       self.model_tester.prepare_config_and_inputs()
																		self.model_tester.check_use_cache_forward_with_attn_mask(
																		    snake_case__       ,  snake_case__       ,  snake_case__       ,  snake_case__			)




						@slow
						def 		UpperCAmelCase_       (   self			):
												'''simple docstring'''
												lowercase__				:      Dict=       GPTaTokenizer.from_pretrained("gpt2"       ,  pad_token="</s>"       ,  padding_side="left"			)
												lowercase__				:      List[str]=       tokenizer(["Hello this is a long string", "Hey"]       ,  return_tensors="np"       ,  padding=snake_case__       ,  truncation=snake_case__			)

												lowercase__				:      Union[str, Any]=       FlaxGPTaLMHeadModel.from_pretrained("gpt2"			)
												lowercase__				:      int=       False
												lowercase__				:      Tuple=       model.config.eos_token_id

												lowercase__				:      Tuple=       jax.jit(model.generate			)

												lowercase__				:      List[Any]=       jit_generate(inputs["input_ids"]       ,  attention_mask=inputs["attention_mask"]			).sequences

												lowercase__				:      List[str]=       tokenizer.batch_decode(snake_case__       ,  skip_special_tokens=snake_case__			)

												lowercase__				:      Optional[int]=       [
												    "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(snake_case__       ,  snake_case__			)




						@is_pt_flax_cross_test
						def 		UpperCAmelCase_       (   self			):
												'''simple docstring'''
												lowercase__,	lowercase__				:      Tuple=       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__				:      Dict=       self._prepare_for_class(snake_case__       ,  snake_case__			)
																								lowercase__				:      Tuple=       {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__				:      Optional[int]=       getattr(snake_case__       ,  snake_case__			)

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

																								lowercase__				:      List[Any]=       convert_pytorch_state_dict_to_flax(pt_model.state_dict()       ,  snake_case__			)
																								lowercase__				:      Tuple=       fx_state

																								with torch.no_grad():
																														lowercase__				:      Optional[Any]=       pt_model(**snake_case__			).to_tuple()

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

																								with tempfile.TemporaryDirectory() as tmpdirname:
																														pt_model.save_pretrained(snake_case__			)
																														lowercase__				:      int=       model_class.from_pretrained(snake_case__       ,  from_pt=snake_case__			)

																								lowercase__				:      Union[str, Any]=       fx_model_loaded(**snake_case__			).to_tuple()
																								self.assertEqual(
																								    len(snake_case__			)       ,  len(snake_case__			)       ,  "Output lengths differ between Flax and PyTorch"			)
																								for fx_output_loaded, pt_output in zip(snake_case__       ,  snake_case__			):
																														self.assert_almost_equals(fx_output_loaded[:, -1]       ,  pt_output[:, -1].numpy()       ,  4e-2			)




						@is_pt_flax_cross_test
						def 		UpperCAmelCase_       (   self			):
												'''simple docstring'''
												lowercase__,	lowercase__				:      Any=       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(snake_case__       ,  snake_case__			)
																								lowercase__				:      Optional[Any]=       {k: torch.tensor(v.tolist()			) for k, v in prepared_inputs_dict.items()}

																								# load corresponding PyTorch class
																								lowercase__				:      List[Any]=       model_class.__name__[4:]  # Skip the "Flax" at the beginning
																								lowercase__				:      Tuple=       getattr(snake_case__       ,  snake_case__			)

																								lowercase__				:      List[Any]=       pt_model_class(snake_case__			).eval()
																								lowercase__				:      int=       model_class(snake_case__       ,  dtype=jnp.floataa			)

																								lowercase__				:      str=       load_flax_weights_in_pytorch_model(snake_case__       ,  fx_model.params			)
																								lowercase__,	lowercase__				:      Dict=       pt_inputs["input_ids"].shape
																								lowercase__				:      int=       np.random.randint(0       ,  seq_length - 1       ,  size=(batch_size,)			)
																								for batch_idx, start_index in enumerate(snake_case__			):
																														lowercase__				:      Optional[int]=       0
																														lowercase__				:      Any=       1
																														lowercase__				:      List[str]=       0
																														lowercase__				:      Optional[int]=       1

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

																								with torch.no_grad():
																														lowercase__				:      Dict=       pt_model(**snake_case__			).to_tuple()

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

																								with tempfile.TemporaryDirectory() as tmpdirname:
																														fx_model.save_pretrained(snake_case__			)
																														lowercase__				:      Tuple=       pt_model_class.from_pretrained(snake_case__       ,  from_flax=snake_case__			)

																								with torch.no_grad():
																														lowercase__				:      Optional[Any]=       pt_model_loaded(**snake_case__			).to_tuple()

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




						@slow
						def 		UpperCAmelCase_       (   self			):
												'''simple docstring'''
												for model_class_name in self.all_model_classes:
																		lowercase__				:      str=       model_class_name.from_pretrained("gpt2"       ,  from_pt=snake_case__			)
																		lowercase__				:      List[str]=       model(np.ones((1, 1)			)			)
																		self.assertIsNotNone(snake_case__			)














"""simple docstring"""







def 						lowercase__(A						)				->int:




						"""simple docstring"""



						lowercase__				:      str=       0
						for plain_chr in plain_text:
												lowercase__				:      List[str]=       (
												    ord(A						) + (hash_value << 6) + (hash_value << 16) - hash_value
												)
						return hash_value






