





'''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							:
		def __init__(  self,__SCREAMING_SNAKE_CASE,__SCREAMING_SNAKE_CASE=14,__SCREAMING_SNAKE_CASE=7,__SCREAMING_SNAKE_CASE=True,__SCREAMING_SNAKE_CASE=True,__SCREAMING_SNAKE_CASE=False,__SCREAMING_SNAKE_CASE=True,__SCREAMING_SNAKE_CASE=99,__SCREAMING_SNAKE_CASE=32,__SCREAMING_SNAKE_CASE=5,__SCREAMING_SNAKE_CASE=4,__SCREAMING_SNAKE_CASE=37,__SCREAMING_SNAKE_CASE="gelu",__SCREAMING_SNAKE_CASE=0.1,__SCREAMING_SNAKE_CASE=0.1,__SCREAMING_SNAKE_CASE=5_12,__SCREAMING_SNAKE_CASE=0.02,):



					'''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   lowerCamelCase__			(  self			):



					'''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=__SCREAMING_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   lowerCamelCase__			(  self			):



					'''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   lowerCamelCase__			(  self			):



					'''simple docstring'''

					__lowerCAmelCase						,       __lowerCAmelCase						,       __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   lowerCamelCase__			(  self,__SCREAMING_SNAKE_CASE,__SCREAMING_SNAKE_CASE,__SCREAMING_SNAKE_CASE,__SCREAMING_SNAKE_CASE			):



					'''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   lowerCamelCase__			(  self,__SCREAMING_SNAKE_CASE,__SCREAMING_SNAKE_CASE,__SCREAMING_SNAKE_CASE,__SCREAMING_SNAKE_CASE			):



					'''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     _UpperCAmelCase							(						lowerCAmelCase_	,					lowerCAmelCase_	,					unittest.TestCase   ):
		a   :   int     =(FlaxGPTaModel, FlaxGPTaLMHeadModel) if is_flax_available() else ()
		a   :   Optional[int]     =(FlaxGPTaLMHeadModel,) if is_flax_available() else ()
		def   lowerCamelCase__			(  self			):



					'''simple docstring'''

					__lowerCAmelCase											=				FlaxGPTaModelTester(self			)
		def   lowerCamelCase__			(  self			):



					'''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   lowerCamelCase__			(  self			):



					'''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			)
		@slow
		def   lowerCamelCase__			(  self			):



					'''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=__SCREAMING_SNAKE_CASE,truncation=__SCREAMING_SNAKE_CASE			)

					__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(__SCREAMING_SNAKE_CASE,skip_special_tokens=__SCREAMING_SNAKE_CASE			)

					__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(__SCREAMING_SNAKE_CASE,__SCREAMING_SNAKE_CASE			)
		@is_pt_flax_cross_test
		def   lowerCamelCase__			(  self			):



					'''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   lowerCamelCase__			(  self			):



					'''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			)







		@slow
		def   lowerCamelCase__			(  self			):



					'''simple docstring'''

					for model_class_name in self.all_model_classes:
								__lowerCAmelCase											=				model_class_name.from_pretrained("""gpt2""",from_pt=__SCREAMING_SNAKE_CASE			)
								__lowerCAmelCase											=				model(np.ones((1, 1)			)			)
								self.assertIsNotNone(__SCREAMING_SNAKE_CASE			)








'''simple docstring'''







import unittest
from pathlib import Path
from tempfile import TemporaryDirectory
from unittest.mock import Mock, patch

import diffusers.utils.hub_utils





class     _UpperCAmelCase							(						unittest.TestCase   ):
		@patch("""diffusers.utils.hub_utils.get_full_repo_name"""			)
		def   lowerCamelCase__			(  self,__SCREAMING_SNAKE_CASE			):



					'''simple docstring'''

					__lowerCAmelCase											=				"""full_repo_name"""
					with TemporaryDirectory() as tmpdir:
								# Dummy args values
								__lowerCAmelCase											=				Mock()
								__lowerCAmelCase											=				tmpdir
								__lowerCAmelCase											=				0
								__lowerCAmelCase											=				"""hub_token"""
								__lowerCAmelCase											=				"""dataset_name"""
								__lowerCAmelCase											=				0.01
								__lowerCAmelCase											=				10_00_00
								__lowerCAmelCase											=				1_00_00
								__lowerCAmelCase											=				0.01
								__lowerCAmelCase											=				0.02
								__lowerCAmelCase											=				0.03
								__lowerCAmelCase											=				0.0005
								__lowerCAmelCase											=				0.00_0001
								__lowerCAmelCase											=				1
								__lowerCAmelCase											=				10
								__lowerCAmelCase											=				0.001
								__lowerCAmelCase											=				0.1
								__lowerCAmelCase											=				0.2
								__lowerCAmelCase											=				True

								# Model card mush be rendered and saved
								diffusers.utils.hub_utils.create_model_card(__SCREAMING_SNAKE_CASE,model_name="""model_name"""			)
								self.assertTrue((Path(__SCREAMING_SNAKE_CASE			) / """README.md""").is_file()			)

