




'''simple docstring'''




import tempfile
import unittest

import numpy as np

import transformers
from transformers import GPTaTokenizer, GPTNeoConfig, 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, 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.gpt_neo.modeling_flax_gpt_neo import FlaxGPTNeoForCausalLM, FlaxGPTNeoModel

if is_torch_available():
							import torch






class      UpperCamelCase_  :







						"""simple docstring"""


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





											"""simple docstring"""







											A_			:  Dict			=				parent
											A_			:  Tuple			=				batch_size
											A_			:  Optional[int]			=				seq_length
											A_			:  Optional[Any]			=				is_training
											A_			:  Tuple			=				use_input_mask
											A_			:  int			=				use_token_type_ids
											A_			:  List[Any]			=				use_labels
											A_			:  Tuple			=				vocab_size
											A_			:  List[Any]			=				hidden_size
											A_			:  Optional[Any]			=				num_hidden_layers
											A_			:  Optional[Any]			=				num_attention_heads
											A_			:  Optional[int]			=				attention_types
											A_			:  List[str]			=				intermediate_size
											A_			:  Optional[int]			=				hidden_act
											A_			:  Optional[Any]			=				hidden_dropout_prob
											A_			:  Optional[int]			=				attention_probs_dropout_prob
											A_			:  Dict			=				max_position_embeddings
											A_			:  int			=				window_size
											A_			:  str			=				initializer_range
											A_			:  Any			=				None
											A_			:  int			=				vocab_size - 1
											A_			:  Tuple			=				vocab_size - 1
											A_			:  Dict			=				vocab_size - 1







						def    _a     (  self   :							Union[str, Any]     ):





											"""simple docstring"""







											A_			:  str			=				ids_tensor([self.batch_size, self.seq_length]				,			self.vocab_size     )

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

											A_			:  Optional[Any]			=				GPTNeoConfig(
											    vocab_size=self.vocab_size				,			hidden_size=self.hidden_size				,			num_layers=self.num_hidden_layers				,			num_heads=self.num_attention_heads				,			max_position_embeddings=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				,			window_size=self.window_size				,			attention_types=self.attention_types				,			)

											return (config, input_ids, input_mask)







						def    _a     (  self   :							int     ):





											"""simple docstring"""







											A_			:  Union[str, Any]			=				self.prepare_config_and_inputs()
											A_			,A_			,A_			:  Optional[Any]			=				config_and_inputs
											A_			:  Union[str, Any]			=				{'''input_ids''': input_ids, '''attention_mask''': attention_mask}
											return config, inputs_dict







						def    _a     (  self   :							Optional[int]				,			_lowerCamelCase   :							Optional[Any]				,			_lowerCamelCase   :							Any				,			_lowerCamelCase   :							Optional[Any]				,			_lowerCamelCase   :							Any     ):





											"""simple docstring"""







											A_			:  Dict			=				20
											A_			:  Union[str, Any]			=				model_class_name(_lowerCamelCase     )

											A_			:  Optional[int]			=				model.init_cache(input_ids.shape[0]				,			_lowerCamelCase     )
											A_			:  Dict			=				jnp.ones((input_ids.shape[0], max_decoder_length)				,			dtype='''i4'''     )

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

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

											A_			:  Optional[int]			=				model(_lowerCamelCase     )

											A_			:  Optional[int]			=				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    _a     (  self   :							Optional[int]				,			_lowerCamelCase   :							Optional[int]				,			_lowerCamelCase   :							Dict				,			_lowerCamelCase   :							str				,			_lowerCamelCase   :							str     ):





											"""simple docstring"""







											A_			:  List[Any]			=				20
											A_			:  List[Any]			=				model_class_name(_lowerCamelCase     )

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

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

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

											A_			:  Optional[int]			=				model(_lowerCamelCase				,			attention_mask=_lowerCamelCase     )

											A_			:  Optional[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      UpperCamelCase_  (a__,   a__,   unittest.TestCase  ):







						"""simple docstring"""


						_lowerCAmelCase         =       (FlaxGPTNeoModel, FlaxGPTNeoForCausalLM) if is_flax_available() else ()
						_lowerCAmelCase         =       (FlaxGPTNeoForCausalLM,) if is_flax_available() else ()







						def    _a     (  self   :							Optional[Any]     ):





											"""simple docstring"""







											A_			:  Tuple			=				FlaxGPTNeoModelTester(self     )







						def    _a     (  self   :							str     ):





											"""simple docstring"""







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







						def    _a     (  self   :							Optional[int]     ):





											"""simple docstring"""







											for model_class_name in self.all_model_classes:
																A_			,A_			,A_			:  List[str]			=				self.model_tester.prepare_config_and_inputs()
																self.model_tester.check_use_cache_forward_with_attn_mask(
																    _lowerCamelCase				,			_lowerCamelCase				,			_lowerCamelCase				,			_lowerCamelCase     )







						@slow
						def    _a     (  self   :							Optional[int]     ):





											"""simple docstring"""







											A_			:  List[Any]			=				GPTaTokenizer.from_pretrained('''gpt2'''				,			pad_token='''<|endoftext|>'''				,			padding_side='''left'''     )
											A_			:  List[str]			=				tokenizer(['''Hello this is a long string''', '''Hey''']				,			return_tensors='''np'''				,			padding=_lowerCamelCase				,			truncation=_lowerCamelCase     )

											A_			:  List[Any]			=				FlaxGPTNeoForCausalLM.from_pretrained('''EleutherAI/gpt-neo-125M'''     )
											A_			:  Union[str, Any]			=				False
											A_			:  Tuple			=				model.config.eos_token_id

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

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

											A_			:  Optional[int]			=				tokenizer.batch_decode(_lowerCamelCase				,			skip_special_tokens=_lowerCamelCase     )

											A_			:  List[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(_lowerCamelCase				,			_lowerCamelCase     )







						@is_pt_flax_cross_test
						def    _a     (  self   :							Any     ):





											"""simple docstring"""







											A_			,A_			:  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
																					A_			:  Union[str, Any]			=				self._prepare_for_class(_lowerCamelCase				,			_lowerCamelCase     )
																					A_			:  Dict			=				{k: torch.tensor(v.tolist()     ) for k, v in prepared_inputs_dict.items()}

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

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

																					A_			:  List[Any]			=				convert_pytorch_state_dict_to_flax(pt_model.state_dict()				,			_lowerCamelCase     )
																					A_			:  Dict			=				fx_state

																					with torch.no_grad():
																										A_			:  Optional[Any]			=				pt_model(**_lowerCamelCase     ).to_tuple()

																					A_			:  Union[str, Any]			=				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     )
																										A_			:  Any			=				model_class.from_pretrained(_lowerCamelCase				,			from_pt=_lowerCamelCase     )

																					A_			:  Dict			=				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    _a     (  self   :							Tuple     ):





											"""simple docstring"""







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

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

																					A_			:  Any			=				pt_model_class(_lowerCamelCase     ).eval()
																					A_			:  Union[str, Any]			=				model_class(_lowerCamelCase				,			dtype=jnp.floataa     )

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

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

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

																					A_			:  Union[str, Any]			=				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     )
																										A_			:  Dict			=				pt_model_class.from_pretrained(_lowerCamelCase				,			from_flax=_lowerCamelCase     )

																					with torch.no_grad():
																										A_			:  Dict			=				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    _a     (  self   :							str     ):





											"""simple docstring"""







											for model_class_name in self.all_model_classes:
																A_			:  int			=				model_class_name.from_pretrained('''EleutherAI/gpt-neo-125M'''     )
																A_			:  Any			=				model(np.ones((1, 1)     )     )
																self.assertIsNotNone(_lowerCamelCase     )









'''simple docstring'''




import unittest

import numpy as np

from transformers.testing_utils import require_torch, require_vision
from transformers.utils import is_torch_available, is_vision_available

from ...test_image_processing_common import ImageProcessingSavingTestMixin


if is_torch_available():
							import torch

if is_vision_available():
							from PIL import Image

							from transformers import SwinaSRImageProcessor
							from transformers.image_transforms import get_image_size






class      UpperCamelCase_  (unittest.TestCase  ):







						"""simple docstring"""


						def __init__(  self   :							Tuple				,			_lowerCamelCase   :							str				,			_lowerCamelCase   :							str=7				,			_lowerCamelCase   :							Tuple=3				,			_lowerCamelCase   :							Any=18				,			_lowerCamelCase   :							Optional[int]=30				,			_lowerCamelCase   :							List[Any]=400				,			_lowerCamelCase   :							int=True				,			_lowerCamelCase   :							Dict=1 / 255				,			_lowerCamelCase   :							List[str]=True				,			_lowerCamelCase   :							Optional[int]=8				,			):





											"""simple docstring"""







											A_			:  Tuple			=				parent
											A_			:  str			=				batch_size
											A_			:  Optional[Any]			=				num_channels
											A_			:  Union[str, Any]			=				image_size
											A_			:  Optional[int]			=				min_resolution
											A_			:  List[Any]			=				max_resolution
											A_			:  int			=				do_rescale
											A_			:  Optional[Any]			=				rescale_factor
											A_			:  Any			=				do_pad
											A_			:  Tuple			=				pad_size







						def    _a     (  self   :							List[Any]     ):





											"""simple docstring"""







											return {
											    "do_rescale": self.do_rescale,
											    "rescale_factor": self.rescale_factor,
											    "do_pad": self.do_pad,
											    "pad_size": self.pad_size,
											}







						def    _a     (  self   :							Tuple				,			_lowerCamelCase   :							List[str]=False				,			_lowerCamelCase   :							Tuple=False				,			_lowerCamelCase   :							str=False     ):





											"""simple docstring"""







											assert not (numpify and torchify), "You cannot specify both numpy and PyTorch tensors at the same time"

											if equal_resolution:
																A_			:  List[Any]			=				[]
																for i in range(self.batch_size     ):
																					image_inputs.append(
																					    np.random.randint(
																					        255				,			size=(self.num_channels, self.max_resolution, self.max_resolution)				,			dtype=np.uinta     )     )
											else:
																A_			:  List[Any]			=				[]
																for i in range(self.batch_size     ):
																					A_			,A_			:  List[str]			=				np.random.choice(np.arange(self.min_resolution				,			self.max_resolution     )				,			2     )
																					image_inputs.append(np.random.randint(255				,			size=(self.num_channels, width, height)				,			dtype=np.uinta     )     )

											if not numpify and not torchify:
																# PIL expects the channel dimension as last dimension
																A_			:  Any			=				[Image.fromarray(np.moveaxis(_lowerCamelCase				,			0				,			-1     )     ) for x in image_inputs]

											if torchify:
																A_			:  str			=				[torch.from_numpy(_lowerCamelCase     ) for x in image_inputs]

											return image_inputs






@require_torch
@require_vision
class      UpperCamelCase_  (a__,   unittest.TestCase  ):







						"""simple docstring"""


						_lowerCAmelCase         =       SwinaSRImageProcessor if is_vision_available() else None







						def    _a     (  self   :							List[str]     ):





											"""simple docstring"""







											A_			:  int			=				SwinaSRImageProcessingTester(self     )







						@property
						def    _a     (  self   :							Union[str, Any]     ):





											"""simple docstring"""







											return self.image_processor_tester.prepare_image_processor_dict()







						def    _a     (  self   :							Any     ):





											"""simple docstring"""







											A_			:  Optional[int]			=				self.image_processing_class(**self.image_processor_dict     )
											self.assertTrue(hasattr(_lowerCamelCase				,			'''do_rescale'''     )     )
											self.assertTrue(hasattr(_lowerCamelCase				,			'''rescale_factor'''     )     )
											self.assertTrue(hasattr(_lowerCamelCase				,			'''do_pad'''     )     )
											self.assertTrue(hasattr(_lowerCamelCase				,			'''pad_size'''     )     )







						def    _a     (  self   :							Tuple     ):





											"""simple docstring"""







											pass







						def    _a     (  self   :							Optional[int]				,			_lowerCamelCase   :							Optional[int]     ):





											"""simple docstring"""







											A_			,A_			:  Any			=				get_image_size(_lowerCamelCase     )
											A_			:  List[Any]			=				self.image_processor_tester.pad_size

											A_			:  Dict			=				(old_height // size + 1) * size - old_height
											A_			:  Optional[Any]			=				(old_width // size + 1) * size - old_width
											return old_height + pad_height, old_width + pad_width







						def    _a     (  self   :							int     ):





											"""simple docstring"""







											A_			:  Dict			=				self.image_processing_class(**self.image_processor_dict     )
											# create random PIL images
											A_			:  Optional[int]			=				self.image_processor_tester.prepare_inputs(equal_resolution=_lowerCamelCase     )
											for image in image_inputs:
																self.assertIsInstance(_lowerCamelCase				,			Image.Image     )

											# Test not batched input
											A_			:  Any			=				image_processor(image_inputs[0]				,			return_tensors='''pt'''     ).pixel_values
											A_			,A_			:  List[Any]			=				self.calculate_expected_size(np.array(image_inputs[0]     )     )
											self.assertEqual(
											    encoded_images.shape				,			(
											        1,
											        self.image_processor_tester.num_channels,
											        expected_height,
											        expected_width,
											    )				,			)







						def    _a     (  self   :							Optional[Any]     ):





											"""simple docstring"""







											A_			:  Optional[Any]			=				self.image_processing_class(**self.image_processor_dict     )
											# create random numpy tensors
											A_			:  Dict			=				self.image_processor_tester.prepare_inputs(equal_resolution=_lowerCamelCase				,			numpify=_lowerCamelCase     )
											for image in image_inputs:
																self.assertIsInstance(_lowerCamelCase				,			np.ndarray     )

											# Test not batched input
											A_			:  Union[str, Any]			=				image_processor(image_inputs[0]				,			return_tensors='''pt'''     ).pixel_values
											A_			,A_			:  Union[str, Any]			=				self.calculate_expected_size(image_inputs[0]     )
											self.assertEqual(
											    encoded_images.shape				,			(
											        1,
											        self.image_processor_tester.num_channels,
											        expected_height,
											        expected_width,
											    )				,			)







						def    _a     (  self   :							int     ):





											"""simple docstring"""







											A_			:  Union[str, Any]			=				self.image_processing_class(**self.image_processor_dict     )
											# create random PyTorch tensors
											A_			:  List[Any]			=				self.image_processor_tester.prepare_inputs(equal_resolution=_lowerCamelCase				,			torchify=_lowerCamelCase     )
											for image in image_inputs:
																self.assertIsInstance(_lowerCamelCase				,			torch.Tensor     )

											# Test not batched input
											A_			:  int			=				image_processor(image_inputs[0]				,			return_tensors='''pt'''     ).pixel_values
											A_			,A_			:  str			=				self.calculate_expected_size(image_inputs[0]     )
											self.assertEqual(
											    encoded_images.shape				,			(
											        1,
											        self.image_processor_tester.num_channels,
											        expected_height,
											        expected_width,
											    )				,			)



