




'''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      :       Tuple							, _lowerCamelCase      :       int							, _lowerCamelCase      :       Union[str, Any]=14							, _lowerCamelCase      :       int=7							, _lowerCamelCase      :       Optional[Any]=True							, _lowerCamelCase      :       str=True							, _lowerCamelCase      :       Any=False							, _lowerCamelCase      :       str=True							, _lowerCamelCase      :       Optional[Any]=99							, _lowerCamelCase      :       Dict=32							, _lowerCamelCase      :       List[str]=5							, _lowerCamelCase      :       Tuple=4							, _lowerCamelCase      :       str=37							, _lowerCamelCase      :       Any="gelu"							, _lowerCamelCase      :       Any=0.1							, _lowerCamelCase      :       int=0.1							, _lowerCamelCase      :       List[str]=512							, _lowerCamelCase      :       Dict=0.02							, ):





											"""simple docstring"""







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







						def       _a    (     self      :       Optional[Any]    ):





											"""simple docstring"""







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

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

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





											"""simple docstring"""







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







						def       _a    (     self      :       List[str]    ):





											"""simple docstring"""







											A_ ,A_ ,A_ :     Union[str, Any] =							self.prepare_config_and_inputs()

											A_ :     List[str] =							floats_tensor([self.batch_size, self.seq_length, self.hidden_size]    )
											A_ :     int =							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      :       Tuple							, _lowerCamelCase      :       int							, _lowerCamelCase      :       Union[str, Any]							, _lowerCamelCase      :       Optional[int]							, _lowerCamelCase      :       List[Any]    ):





											"""simple docstring"""







											A_ :     Tuple =							20
											A_ :     Tuple =							model_class_name(_lowerCamelCase    )

											A_ :     Any =							model.init_cache(input_ids.shape[0]							, _lowerCamelCase    )
											A_ :     List[str] =							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_ :     Union[str, Any] =							model(
											    input_ids[:, :-1]							, attention_mask=_lowerCamelCase							, past_key_values=_lowerCamelCase							, position_ids=_lowerCamelCase							, )

											A_ :     Tuple =							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_ :     Tuple =							model(_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}'    )







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





											"""simple docstring"""







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

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

											A_ :     Optional[int] =							model.init_cache(input_ids.shape[0]							, _lowerCamelCase    )
											A_ :     Optional[int] =							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_ :     int =							jnp.array(input_ids.shape[0] * [[input_ids.shape[-1] - 1]]							, dtype='''i4'''    )
											A_ :     Optional[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     UpperCamelCase_     (a__,      a__,      unittest.TestCase     ):







						"""simple docstring"""


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







						def       _a    (     self      :       Optional[int]    ):





											"""simple docstring"""







											A_ :     Optional[int] =							FlaxGPTaModelTester(self    )







						def       _a    (     self      :       int    ):





											"""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_ :     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      :       Optional[Any]    ):





											"""simple docstring"""







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

											A_ :     str =							FlaxGPTaLMHeadModel.from_pretrained('''gpt2'''    )
											A_ :     Optional[int] =							False
											A_ :     List[str] =							model.config.eos_token_id

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

											A_ :     List[Any] =							jit_generate(inputs['''input_ids''']							, attention_mask=inputs['''attention_mask''']    ).sequences

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

											A_ :     Tuple =							[
											    '''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      :       List[str]    ):





											"""simple docstring"""







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

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

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

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

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

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

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

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

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

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

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





											"""simple docstring"""







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









'''simple docstring'''




from __future__ import annotations

snake_case__          =						[2, 1, 4, 5, 6, 0, 8, 9, 1, 2, 0, 6, 4, 2, 0, 6, 5, 3, 2, 7]






class     UpperCamelCase_     :







						"""simple docstring"""


						def __init__(     self      :       Dict							, _lowerCamelCase      :       int    ):





											"""simple docstring"""







											A_ :     int =							-1
											A_ :     int =							-1
											A_ :     list[int] =							[-1] * length
											A_ :     Node | None =							None
											A_ :     Node | None =							None







						def __repr__(     self      :       List[str]    ):





											"""simple docstring"""







											return f'Node(min_value={self.minn} max_value={self.maxx})'






def 						snake_case__    ( lowerCamelCase__				:   list[int]			)   ->   Node | None:
					A_ :     Optional[Any] =							Node(len(lowerCamelCase__			)			)
					A_ ,A_ :     str =							min(lowerCamelCase__			), max(lowerCamelCase__			)
					# Leaf node case where the node contains only one unique value
					if root.minn == root.maxx:
										return root
					A_ :     Optional[int] =							(root.minn + root.maxx) // 2

					A_ :     list[int] =							[]
					A_ :     list[int] =							[]

					for index, num in enumerate(lowerCamelCase__			):
										if num <= pivot:
															left_arr.append(lowerCamelCase__			)
										else:
															right_arr.append(lowerCamelCase__			)
										A_ :     Union[str, Any] =							len(lowerCamelCase__			)
					A_ :     Optional[int] =							build_tree(lowerCamelCase__			)
					A_ :     Optional[int] =							build_tree(lowerCamelCase__			)
					return root




def 						snake_case__    ( lowerCamelCase__				:   Node | None			,		lowerCamelCase__				:   int			,		lowerCamelCase__				:   int			)   ->   int:
					if index < 0 or node is None:
										return 0
					# Leaf node cases
					if node.minn == node.maxx:
										return index + 1 if node.minn == num else 0
					A_ :     List[Any] =							(node.minn + node.maxx) // 2
					if num <= pivot:
										# go the left subtree and map index to the left subtree
										return rank_till_index(node.left			,		lowerCamelCase__			,		node.map_left[index] - 1			)
					else:
										# go to the right subtree and map index to the right subtree
										return rank_till_index(node.right			,		lowerCamelCase__			,		index - node.map_left[index]			)




def 						snake_case__    ( lowerCamelCase__				:   Node | None			,		lowerCamelCase__				:   int			,		lowerCamelCase__				:   int			,		lowerCamelCase__				:   int			)   ->   int:
					if start > end:
										return 0
					A_ :     Union[str, Any] =							rank_till_index(lowerCamelCase__			,		lowerCamelCase__			,		lowerCamelCase__			)
					A_ :     Union[str, Any] =							rank_till_index(lowerCamelCase__			,		lowerCamelCase__			,		start - 1			)
					return rank_till_end - rank_before_start




def 						snake_case__    ( lowerCamelCase__				:   Node | None			,		lowerCamelCase__				:   int			,		lowerCamelCase__				:   int			,		lowerCamelCase__				:   int			)   ->   int:
					if index > (end - start) or start > end or node is None:
										return -1
					# Leaf node case
					if node.minn == node.maxx:
										return node.minn
					# Number of elements in the left subtree in interval [start, end]
					A_ :     List[str] =							node.map_left[end] - (
					    node.map_left[start - 1] if start else 0
					)
					if num_elements_in_left_tree > index:
										return quantile(
										    node.left			,		lowerCamelCase__			,		(node.map_left[start - 1] if start else 0)			,		node.map_left[end] - 1			,		)
					else:
										return quantile(
										    node.right			,		index - num_elements_in_left_tree			,		start - (node.map_left[start - 1] if start else 0)			,		end - node.map_left[end]			,		)




def 						snake_case__    ( lowerCamelCase__				:   Node | None			,		lowerCamelCase__				:   int			,		lowerCamelCase__				:   int			,		lowerCamelCase__				:   int			,		lowerCamelCase__				:   int			)   ->   int:
					if (
					    start > end
					    or node is None
					    or start_num > end_num
					    or node.minn > end_num
					    or node.maxx < start_num
					):
										return 0
					if start_num <= node.minn and node.maxx <= end_num:
										return end - start + 1
					A_ :     int =							range_counting(
					    node.left			,		(node.map_left[start - 1] if start else 0)			,		node.map_left[end] - 1			,		lowerCamelCase__			,		lowerCamelCase__			,		)
					A_ :     List[str] =							range_counting(
					    node.right			,		start - (node.map_left[start - 1] if start else 0)			,		end - node.map_left[end]			,		lowerCamelCase__			,		lowerCamelCase__			,		)
					return left + right


if __name__ == "__main__":
							import doctest

							doctest.testmod()



