




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







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



					"""simple docstring"""

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







				def    a_					(							self				:   Dict    ):



					"""simple docstring"""

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

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

					A_					:			Tuple						=						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    a_					(							self				:   Optional[Any]    ):



					"""simple docstring"""

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







				def    a_					(							self				:   List[str]    ):



					"""simple docstring"""

					A_				,	A_				,	A_					:			Optional[int]						=						self.prepare_config_and_inputs()

					A_					:			List[str]						=						floats_tensor([self.batch_size, self.seq_length, self.hidden_size]    )
					A_					:			str						=						ids_tensor([self.batch_size, self.seq_length] ,       vocab_size=2    )

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







				def    a_					(							self				:   Dict ,       _lowerCamelCase				:   Any ,       _lowerCamelCase				:   int ,       _lowerCamelCase				:   Dict ,       _lowerCamelCase				:   int    ):



					"""simple docstring"""

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

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

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

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

					A_					:			List[Any]						=						model(_lowerCamelCase    )

					A_					:			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}"""    )







				def    a_					(							self				:   Any ,       _lowerCamelCase				:   int ,       _lowerCamelCase				:   Any ,       _lowerCamelCase				:   List[Any] ,       _lowerCamelCase				:   Any    ):



					"""simple docstring"""

					A_					:			str						=						20
					A_					:			int						=						model_class_name(_lowerCamelCase    )

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

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

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

					A_					:			int						=						model(_lowerCamelCase ,       attention_mask=_lowerCamelCase    )

					A_					:			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}"""    )

@require_flax
class     lowercase (      __UpperCAmelCase       ,		__UpperCAmelCase       ,		unittest.TestCase):
				__lowerCAmelCase      :      Tuple					    =				(FlaxGPTaModel, FlaxGPTaLMHeadModel) if is_flax_available() else ()
				__lowerCAmelCase      :      List[Any]					    =				(FlaxGPTaLMHeadModel,) if is_flax_available() else ()







				def    a_					(							self				:   Dict    ):



					"""simple docstring"""

					A_					:			List[str]						=						FlaxGPTaModelTester(self    )







				def    a_					(							self				:   List[Any]    ):



					"""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				:   Any    ):



					"""simple docstring"""

					for model_class_name in self.all_model_classes:
						A_				,	A_				,	A_					:			List[Any]						=						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				:   str    ):



					"""simple docstring"""

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

					A_					:			Any						=						FlaxGPTaLMHeadModel.from_pretrained('''gpt2'''    )
					A_					:			Dict						=						False
					A_					:			Optional[Any]						=						model.config.eos_token_id

					A_					:			Dict						=						jax.jit(model.generate    )

					A_					:			Optional[int]						=						jit_generate(inputs['''input_ids'''] ,       attention_mask=inputs['''attention_mask''']    ).sequences

					A_					:			Union[str, Any]						=						tokenizer.batch_decode(_lowerCamelCase ,       skip_special_tokens=_lowerCamelCase    )

					A_					:			Any						=						[
					    '''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    a_					(							self				:   Union[str, Any]    ):



					"""simple docstring"""

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

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

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

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

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

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

							A_					:			List[Any]						=						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				:   Optional[Any]    ):



					"""simple docstring"""

					A_				,	A_					:			str						=						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_					:			Optional[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_					:			Tuple						=						model_class.__name__[4:]  # Skip the "Flax" at the beginning
							A_					:			Optional[Any]						=						getattr(_lowerCamelCase ,       _lowerCamelCase    )

							A_					:			Tuple						=						pt_model_class(_lowerCamelCase    ).eval()
							A_					:			Dict						=						model_class(_lowerCamelCase ,       dtype=jnp.floataa    )

							A_					:			Union[str, Any]						=						load_flax_weights_in_pytorch_model(_lowerCamelCase ,       fx_model.params    )
							A_				,	A_					:			List[Any]						=						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_					:			Tuple						=						0
								A_					:			List[Any]						=						1
								A_					:			Any						=						0
								A_					:			Optional[Any]						=						1

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

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

							A_					:			List[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_					:			str						=						pt_model_class.from_pretrained(_lowerCamelCase ,       from_flax=_lowerCamelCase    )

							with torch.no_grad():
								A_					:			int						=						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				:   Optional[int]    ):



					"""simple docstring"""

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











"""simple docstring"""

import bisect
import copy
from collections import defaultdict

import numpy as np
from torch.utils.data import BatchSampler, Sampler

from utils import logger

def 						lowercase_ (						_UpperCAmelCase					, _UpperCAmelCase		):




	"""simple docstring"""






	A_					:			Optional[int]						=						copy.deepcopy(_UpperCAmelCase		)
	A_					:			int						=						sorted(_UpperCAmelCase		)
	A_					:			Tuple						=						[bisect.bisect_right(_UpperCAmelCase					, _UpperCAmelCase		) for y in x]
	return quantized

def 						lowercase_ (						_UpperCAmelCase					, _UpperCAmelCase=0		):




	"""simple docstring"""






	A_					:			List[Any]						=						np.arange(start=3					, stop=_UpperCAmelCase					, step=4		).tolist() if k > 0 else [10]
	A_					:			List[Any]						=						_quantize(_UpperCAmelCase					, _UpperCAmelCase		)
	# count number of elements per group
	A_					:			Tuple						=						np.unique(_UpperCAmelCase					, return_counts=_UpperCAmelCase		)[1]
	A_					:			Optional[int]						=						[0] + bins + [np.inf]
	logger.info('''Using {} as bins for aspect lengths quantization'''.format(_UpperCAmelCase		)		)
	logger.info('''Count of instances per bin: {}'''.format(_UpperCAmelCase		)		)
	return groups

class     lowercase (      __UpperCAmelCase):







				def __init__(							self				:   Union[str, Any] ,       _lowerCamelCase				:   List[Any] ,       _lowerCamelCase				:   Dict ,       _lowerCamelCase				:   Dict    ):



					"""simple docstring"""

					if not isinstance(_lowerCamelCase ,       _lowerCamelCase    ):
						raise ValueError(
						    '''sampler should be an instance of torch.utils.data.Sampler, but got sampler={}'''.format(_lowerCamelCase    )    )
					A_					:			List[str]						=						sampler
					A_					:			str						=						group_ids
					A_					:			Any						=						batch_size







				def __iter__(							self				:   List[str]    ):



					"""simple docstring"""

					A_					:			List[str]						=						defaultdict(_lowerCamelCase    )
					A_					:			Dict						=						defaultdict(_lowerCamelCase    )

					A_					:			List[str]						=						0
					for idx in self.sampler:
						A_					:			Union[str, Any]						=						self.group_ids[idx]
						buffer_per_group[group_id].append(_lowerCamelCase    )
						samples_per_group[group_id].append(_lowerCamelCase    )
						if len(buffer_per_group[group_id]    ) == self.batch_size:
							yield buffer_per_group[group_id]  # TODO
							num_batches += 1
							del buffer_per_group[group_id]
						assert len(buffer_per_group[group_id]    ) < self.batch_size

					# now we have run out of elements that satisfy
					# the group criteria, let's return the remaining
					# elements so that the size of the sampler is
					# deterministic
					A_					:			int						=						len(self    )
					A_					:			Dict						=						expected_num_batches - num_batches
					if num_remaining > 0:
						# for the remaining batches, group the batches by similar lengths
						A_					:			Optional[Any]						=						[]
						for group_id, idxs in sorted(buffer_per_group.items() ,       key=lambda _lowerCamelCase    : x[0]    ):
							batch_idx.extend(_lowerCamelCase    )
							if len(_lowerCamelCase    ) >= self.batch_size:
								yield batch_idx[: self.batch_size]
								A_					:			str						=						batch_idx[self.batch_size :]
								num_remaining -= 1
						if len(_lowerCamelCase    ) > 0:
							yield batch_idx
							num_remaining -= 1
					assert num_remaining == 0







				def __len__(							self				:   List[Any]    ):



					"""simple docstring"""

					return (len(self.sampler    ) + self.batch_size - 1) // self.batch_size





