
"""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 			lowercase					:






			def __init__(self  :	str      ,SCREAMING_SNAKE_CASE_  :	Union[str, Any]      ,SCREAMING_SNAKE_CASE_  :	Dict=14      ,SCREAMING_SNAKE_CASE_  :	Any=7      ,SCREAMING_SNAKE_CASE_  :	Any=True      ,SCREAMING_SNAKE_CASE_  :	Any=True      ,SCREAMING_SNAKE_CASE_  :	Optional[int]=False      ,SCREAMING_SNAKE_CASE_  :	Dict=True      ,SCREAMING_SNAKE_CASE_  :	Optional[int]=99      ,SCREAMING_SNAKE_CASE_  :	Tuple=32      ,SCREAMING_SNAKE_CASE_  :	str=4      ,SCREAMING_SNAKE_CASE_  :	List[str]=4      ,SCREAMING_SNAKE_CASE_  :	Optional[int]=4      ,SCREAMING_SNAKE_CASE_  :	List[Any]=37      ,SCREAMING_SNAKE_CASE_  :	Tuple="gelu"      ,SCREAMING_SNAKE_CASE_  :	Dict=0.1      ,SCREAMING_SNAKE_CASE_  :	Union[str, Any]=0.1      ,SCREAMING_SNAKE_CASE_  :	List[Any]=512      ,SCREAMING_SNAKE_CASE_  :	Optional[Any]=0.02      ,)       ->	Optional[int]:





				"""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												=  rotary_dim
				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 					UpperCAmelCase							(self  :	Optional[int]			)       ->	int:





				"""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												=  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=SCREAMING_SNAKE_CASE_      ,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 					UpperCAmelCase							(self  :	str			)       ->	Union[str, Any]:





				"""simple docstring"""


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






			def 					UpperCAmelCase							(self  :	Optional[int]      ,SCREAMING_SNAKE_CASE_  :	Optional[Any]      ,SCREAMING_SNAKE_CASE_  :	Tuple      ,SCREAMING_SNAKE_CASE_  :	str      ,SCREAMING_SNAKE_CASE_  :	Tuple			)       ->	Dict:





				"""simple docstring"""


				lowerCAmelCase												=  20
				lowerCAmelCase												=  model_class_name(SCREAMING_SNAKE_CASE_			)

				lowerCAmelCase												=  model.init_cache(input_ids.shape[0]      ,SCREAMING_SNAKE_CASE_			)
				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=SCREAMING_SNAKE_CASE_      ,past_key_values=SCREAMING_SNAKE_CASE_      ,position_ids=SCREAMING_SNAKE_CASE_      ,)

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

				lowerCAmelCase												=  model(SCREAMING_SNAKE_CASE_			)

				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 					UpperCAmelCase							(self  :	Optional[int]      ,SCREAMING_SNAKE_CASE_  :	List[str]      ,SCREAMING_SNAKE_CASE_  :	Optional[Any]      ,SCREAMING_SNAKE_CASE_  :	Dict      ,SCREAMING_SNAKE_CASE_  :	Optional[int]			)       ->	Dict:





				"""simple docstring"""


				lowerCAmelCase												=  20
				lowerCAmelCase												=  model_class_name(SCREAMING_SNAKE_CASE_			)

				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]      ,SCREAMING_SNAKE_CASE_			)
				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=SCREAMING_SNAKE_CASE_      ,past_key_values=SCREAMING_SNAKE_CASE_      ,position_ids=SCREAMING_SNAKE_CASE_      ,)
				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=SCREAMING_SNAKE_CASE_      ,position_ids=SCREAMING_SNAKE_CASE_      ,)

				lowerCAmelCase												=  model(SCREAMING_SNAKE_CASE_      ,attention_mask=SCREAMING_SNAKE_CASE_			)

				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 			lowercase					( lowercase__   ,lowercase__   ,unittest.TestCase		):
			lowercase        =     (FlaxGPTJModel, FlaxGPTJForCausalLM) if is_flax_available() else ()
			lowercase        =     (FlaxGPTJForCausalLM,) if is_flax_available() else ()






			def 					UpperCAmelCase							(self  :	Any			)       ->	Union[str, Any]:





				"""simple docstring"""


				lowerCAmelCase												=  FlaxGPTJModelTester(self			)






			def 					UpperCAmelCase							(self  :	int			)       ->	str:





				"""simple docstring"""


				for model_class_name in self.all_model_classes:
					lowerCAmelCase							,       lowerCAmelCase							,       lowerCAmelCase												=  self.model_tester.prepare_config_and_inputs()
					self.model_tester.check_use_cache_forward(SCREAMING_SNAKE_CASE_      ,SCREAMING_SNAKE_CASE_      ,SCREAMING_SNAKE_CASE_      ,SCREAMING_SNAKE_CASE_			)






			def 					UpperCAmelCase							(self  :	Tuple			)       ->	Tuple:





				"""simple docstring"""


				for model_class_name in self.all_model_classes:
					lowerCAmelCase							,       lowerCAmelCase							,       lowerCAmelCase												=  self.model_tester.prepare_config_and_inputs()
					self.model_tester.check_use_cache_forward_with_attn_mask(
					    SCREAMING_SNAKE_CASE_      ,SCREAMING_SNAKE_CASE_      ,SCREAMING_SNAKE_CASE_      ,SCREAMING_SNAKE_CASE_			)






			@tooslow
			def 					UpperCAmelCase							(self  :	Any			)       ->	Dict:





				"""simple docstring"""


				lowerCAmelCase												=  GPTaTokenizer.from_pretrained('''gpt2'''      ,pad_token='''<|endoftext|>'''      ,padding_side='''left'''			)
				lowerCAmelCase												=  tokenizer(['''Hello this is a long string''', '''Hey''']      ,return_tensors='''np'''      ,padding=SCREAMING_SNAKE_CASE_      ,truncation=SCREAMING_SNAKE_CASE_			)

				lowerCAmelCase												=  FlaxGPTJForCausalLM.from_pretrained('''EleutherAI/gpt-j-6B'''			)
				lowerCAmelCase												=  False
				lowerCAmelCase												=  model.config.eos_token_id

				lowerCAmelCase												=  jax.jit(model.generate			)

				lowerCAmelCase												=  jit_generate(
				    inputs['''input_ids''']      ,attention_mask=inputs['''attention_mask''']      ,pad_token_id=tokenizer.pad_token_id			).sequences

				lowerCAmelCase												=  tokenizer.batch_decode(SCREAMING_SNAKE_CASE_      ,skip_special_tokens=SCREAMING_SNAKE_CASE_			)

				lowerCAmelCase												=  [
				    '''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(SCREAMING_SNAKE_CASE_      ,SCREAMING_SNAKE_CASE_			)






			@is_pt_flax_cross_test
			def 					UpperCAmelCase							(self  :	Union[str, Any]			)       ->	str:





				"""simple docstring"""


				lowerCAmelCase							,       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(SCREAMING_SNAKE_CASE_      ,SCREAMING_SNAKE_CASE_			)
						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(SCREAMING_SNAKE_CASE_      ,SCREAMING_SNAKE_CASE_			)

						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(SCREAMING_SNAKE_CASE_			):
							lowerCAmelCase												=  0
							lowerCAmelCase												=  1
							lowerCAmelCase												=  0
							lowerCAmelCase												=  1
						lowerCAmelCase												=  pt_model_class(SCREAMING_SNAKE_CASE_			).eval()
						lowerCAmelCase												=  model_class(SCREAMING_SNAKE_CASE_      ,dtype=jnp.floataa			)

						lowerCAmelCase												=  convert_pytorch_state_dict_to_flax(pt_model.state_dict()      ,SCREAMING_SNAKE_CASE_			)
						lowerCAmelCase												=  fx_state

						with torch.no_grad():
							lowerCAmelCase												=  pt_model(**SCREAMING_SNAKE_CASE_			).to_tuple()

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

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

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






			@is_pt_flax_cross_test
			def 					UpperCAmelCase							(self  :	Optional[Any]			)       ->	Optional[Any]:





				"""simple docstring"""


				lowerCAmelCase							,       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(SCREAMING_SNAKE_CASE_      ,SCREAMING_SNAKE_CASE_			)
						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(SCREAMING_SNAKE_CASE_      ,SCREAMING_SNAKE_CASE_			)

						lowerCAmelCase												=  pt_model_class(SCREAMING_SNAKE_CASE_			).eval()
						lowerCAmelCase												=  model_class(SCREAMING_SNAKE_CASE_      ,dtype=jnp.floataa			)

						lowerCAmelCase												=  load_flax_weights_in_pytorch_model(SCREAMING_SNAKE_CASE_      ,fx_model.params			)
						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(SCREAMING_SNAKE_CASE_			):
							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(**SCREAMING_SNAKE_CASE_			).to_tuple()

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

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

						with torch.no_grad():
							lowerCAmelCase												=  pt_model_loaded(**SCREAMING_SNAKE_CASE_			).to_tuple()

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






			@tooslow
			def 					UpperCAmelCase							(self  :	List[Any]			)       ->	Union[str, Any]:





				"""simple docstring"""


				for model_class_name in self.all_model_classes:
					lowerCAmelCase												=  model_class_name.from_pretrained('''EleutherAI/gpt-j-6B'''			)
					lowerCAmelCase												=  model(np.ones((1, 1)			)			)
					self.assertIsNotNone(SCREAMING_SNAKE_CASE_			)




"""simple docstring"""






from __future__ import annotations



class 			lowercase					:






			def __init__(self  :	List[Any]      ,SCREAMING_SNAKE_CASE_  :	int = 0			)       ->	List[Any]:





				"""simple docstring"""


				lowerCAmelCase												=  key






			def 					UpperCAmelCase							(self  :	Tuple      ,SCREAMING_SNAKE_CASE_  :	str      ,SCREAMING_SNAKE_CASE_  :	int			)       ->	list[str]:





				"""simple docstring"""


				assert isinstance(SCREAMING_SNAKE_CASE_      ,SCREAMING_SNAKE_CASE_			) and isinstance(SCREAMING_SNAKE_CASE_      ,SCREAMING_SNAKE_CASE_			)

				lowerCAmelCase												=  key or self.__key or 1

				# make sure key is an appropriate size
				key %= 255

				return [chr(ord(SCREAMING_SNAKE_CASE_			) ^ key			) for ch in content]






			def 					UpperCAmelCase							(self  :	int      ,SCREAMING_SNAKE_CASE_  :	str      ,SCREAMING_SNAKE_CASE_  :	int			)       ->	list[str]:





				"""simple docstring"""


				assert isinstance(SCREAMING_SNAKE_CASE_      ,SCREAMING_SNAKE_CASE_			) and isinstance(SCREAMING_SNAKE_CASE_      ,SCREAMING_SNAKE_CASE_			)

				lowerCAmelCase												=  key or self.__key or 1

				# make sure key is an appropriate size
				key %= 255

				return [chr(ord(SCREAMING_SNAKE_CASE_			) ^ key			) for ch in content]






			def 					UpperCAmelCase							(self  :	Union[str, Any]      ,SCREAMING_SNAKE_CASE_  :	str      ,SCREAMING_SNAKE_CASE_  :	int = 0			)       ->	str:





				"""simple docstring"""


				assert isinstance(SCREAMING_SNAKE_CASE_      ,SCREAMING_SNAKE_CASE_			) and isinstance(SCREAMING_SNAKE_CASE_      ,SCREAMING_SNAKE_CASE_			)

				lowerCAmelCase												=  key or self.__key or 1

				# make sure key can be any size
				while key > 255:
					key -= 255

				# This will be returned
				lowerCAmelCase												=  ''''''

				for ch in content:
					ans += chr(ord(SCREAMING_SNAKE_CASE_			) ^ key			)

				return ans






			def 					UpperCAmelCase							(self  :	Union[str, Any]      ,SCREAMING_SNAKE_CASE_  :	str      ,SCREAMING_SNAKE_CASE_  :	int = 0			)       ->	str:





				"""simple docstring"""


				assert isinstance(SCREAMING_SNAKE_CASE_      ,SCREAMING_SNAKE_CASE_			) and isinstance(SCREAMING_SNAKE_CASE_      ,SCREAMING_SNAKE_CASE_			)

				lowerCAmelCase												=  key or self.__key or 1

				# make sure key can be any size
				while key > 255:
					key -= 255

				# This will be returned
				lowerCAmelCase												=  ''''''

				for ch in content:
					ans += chr(ord(SCREAMING_SNAKE_CASE_			) ^ key			)

				return ans






			def 					UpperCAmelCase							(self  :	List[str]      ,SCREAMING_SNAKE_CASE_  :	str      ,SCREAMING_SNAKE_CASE_  :	int = 0			)       ->	bool:





				"""simple docstring"""


				assert isinstance(SCREAMING_SNAKE_CASE_      ,SCREAMING_SNAKE_CASE_			) and isinstance(SCREAMING_SNAKE_CASE_      ,SCREAMING_SNAKE_CASE_			)

				try:
					with open(SCREAMING_SNAKE_CASE_			) as fin, open('''encrypt.out'''      ,'''w+'''			) as fout:
						# actual encrypt-process
						for line in fin:
							fout.write(self.encrypt_string(SCREAMING_SNAKE_CASE_      ,SCREAMING_SNAKE_CASE_			)			)

				except OSError:
					return False

				return True






			def 					UpperCAmelCase							(self  :	List[Any]      ,SCREAMING_SNAKE_CASE_  :	str      ,SCREAMING_SNAKE_CASE_  :	int			)       ->	bool:





				"""simple docstring"""


				assert isinstance(SCREAMING_SNAKE_CASE_      ,SCREAMING_SNAKE_CASE_			) and isinstance(SCREAMING_SNAKE_CASE_      ,SCREAMING_SNAKE_CASE_			)

				try:
					with open(SCREAMING_SNAKE_CASE_			) as fin, open('''decrypt.out'''      ,'''w+'''			) as fout:
						# actual encrypt-process
						for line in fin:
							fout.write(self.decrypt_string(SCREAMING_SNAKE_CASE_      ,SCREAMING_SNAKE_CASE_			)			)

				except OSError:
					return False

				return True


# Tests
# crypt = XORCipher()
# key = 67

# # test encrypt
# print(crypt.encrypt("hallo welt",key))
# # test decrypt
# print(crypt.decrypt(crypt.encrypt("hallo welt",key), key))

# # test encrypt_string
# print(crypt.encrypt_string("hallo welt",key))

# # test decrypt_string
# print(crypt.decrypt_string(crypt.encrypt_string("hallo welt",key),key))

# if (crypt.encrypt_file("test.txt",key)):
#       print("encrypt successful")
# else:
#       print("encrypt unsuccessful")

# if (crypt.decrypt_file("encrypt.out",key)):
#       print("decrypt successful")
# else:
#       print("decrypt unsuccessful")


