



from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_tokenizers_available, is_torch_available


UpperCAmelCase				:							Optional[int]								={
    """configuration_realm""": ["""REALM_PRETRAINED_CONFIG_ARCHIVE_MAP""", """RealmConfig"""],
    """tokenization_realm""": ["""RealmTokenizer"""],
}

try:
	if not is_tokenizers_available():
		raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
	pass
else:
	UpperCAmelCase				:							Optional[int]								=["""RealmTokenizerFast"""]

try:
	if not is_torch_available():
		raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
	pass
else:
	UpperCAmelCase				:							Union[str, Any]								=[
	    """REALM_PRETRAINED_MODEL_ARCHIVE_LIST""",
	    """RealmEmbedder""",
	    """RealmForOpenQA""",
	    """RealmKnowledgeAugEncoder""",
	    """RealmPreTrainedModel""",
	    """RealmReader""",
	    """RealmScorer""",
	    """load_tf_weights_in_realm""",
	]
	UpperCAmelCase				:							List[str]								=["""RealmRetriever"""]


if TYPE_CHECKING:
	from .configuration_realm import REALM_PRETRAINED_CONFIG_ARCHIVE_MAP, RealmConfig
	from .tokenization_realm import RealmTokenizer

	try:
		if not is_tokenizers_available():
			raise OptionalDependencyNotAvailable()
	except OptionalDependencyNotAvailable:
		pass
	else:
		from .tokenization_realm import RealmTokenizerFast

	try:
		if not is_torch_available():
			raise OptionalDependencyNotAvailable()
	except OptionalDependencyNotAvailable:
		pass
	else:
		from .modeling_realm import (
		    REALM_PRETRAINED_MODEL_ARCHIVE_LIST,
		    RealmEmbedder,
		    RealmForOpenQA,
		    RealmKnowledgeAugEncoder,
		    RealmPreTrainedModel,
		    RealmReader,
		    RealmScorer,
		    load_tf_weights_in_realm,
		)
		from .retrieval_realm import RealmRetriever


else:
	import sys

	UpperCAmelCase				:							str								=_LazyModule(__name__, globals()["""__file__"""], _import_structure, module_spec=__spec__)





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				:
				'''simple docstring'''







				def __init__(	self       ,	snake_case__       ,	snake_case__=14       ,	snake_case__=7       ,	snake_case__=True       ,	snake_case__=True       ,	snake_case__=False       ,	snake_case__=True       ,	snake_case__=99       ,	snake_case__=32       ,	snake_case__=5       ,	snake_case__=4       ,	snake_case__=37       ,	snake_case__="gelu"       ,	snake_case__=0.1       ,	snake_case__=0.1       ,	snake_case__=512       ,	snake_case__=0.02       ,	):

									'''simple docstring'''


									UpperCamelCase_											=				parent
									UpperCamelCase_											=				batch_size
									UpperCamelCase_											=				seq_length
									UpperCamelCase_											=				is_training
									UpperCamelCase_											=				use_input_mask
									UpperCamelCase_											=				use_token_type_ids
									UpperCamelCase_											=				use_labels
									UpperCamelCase_											=				vocab_size
									UpperCamelCase_											=				hidden_size
									UpperCamelCase_											=				num_hidden_layers
									UpperCamelCase_											=				num_attention_heads
									UpperCamelCase_											=				intermediate_size
									UpperCamelCase_											=				hidden_act
									UpperCamelCase_											=				hidden_dropout_prob
									UpperCamelCase_											=				attention_probs_dropout_prob
									UpperCamelCase_											=				max_position_embeddings
									UpperCamelCase_											=				initializer_range
									UpperCamelCase_											=				None
									UpperCamelCase_											=				vocab_size - 1
									UpperCamelCase_											=				vocab_size - 1
									UpperCamelCase_											=				vocab_size - 1

				def  _lowerCamelCase   (	self      ):

									'''simple docstring'''


									UpperCamelCase_											=				ids_tensor([self.batch_size, self.seq_length]       ,	self.vocab_size      )

									UpperCamelCase_											=				None
									if self.use_input_mask:
														UpperCamelCase_											=				random_attention_mask([self.batch_size, self.seq_length]      )

									UpperCamelCase_											=				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=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'''


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

				def  _lowerCamelCase   (	self      ):

									'''simple docstring'''


									UpperCamelCase_					,      UpperCamelCase_					,      UpperCamelCase_											=				self.prepare_config_and_inputs()

									UpperCamelCase_											=				floats_tensor([self.batch_size, self.seq_length, self.hidden_size]      )
									UpperCamelCase_											=				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       ,	snake_case__       ,	snake_case__       ,	snake_case__       ,	snake_case__      ):

									'''simple docstring'''


									UpperCamelCase_											=				20
									UpperCamelCase_											=				model_class_name(snake_case__      )

									UpperCamelCase_											=				model.init_cache(input_ids.shape[0]       ,	snake_case__      )
									UpperCamelCase_											=				jnp.ones((input_ids.shape[0], max_decoder_length)       ,	dtype="i4"      )

									UpperCamelCase_											=				jnp.broadcast_to(
									    jnp.arange(input_ids.shape[-1] - 1      )[None, :]       ,	(input_ids.shape[0], input_ids.shape[-1] - 1)      )
									UpperCamelCase_											=				model(
									    input_ids[:, :-1]       ,	attention_mask=snake_case__       ,	past_key_values=snake_case__       ,	position_ids=snake_case__       ,	)

									UpperCamelCase_											=				jnp.array(input_ids.shape[0] * [[input_ids.shape[-1] - 1]]       ,	dtype="i4"      )
									UpperCamelCase_											=				model(
									    input_ids[:, -1:]       ,	attention_mask=snake_case__       ,	past_key_values=outputs_cache.past_key_values       ,	position_ids=snake_case__       ,	)

									UpperCamelCase_											=				model(snake_case__      )

									UpperCamelCase_											=				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       ,	snake_case__       ,	snake_case__       ,	snake_case__       ,	snake_case__      ):

									'''simple docstring'''


									UpperCamelCase_											=				20
									UpperCamelCase_											=				model_class_name(snake_case__      )

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

									UpperCamelCase_											=				model.init_cache(input_ids.shape[0]       ,	snake_case__      )
									UpperCamelCase_											=				jnp.broadcast_to(
									    jnp.arange(input_ids.shape[-1] - 1      )[None, :]       ,	(input_ids.shape[0], input_ids.shape[-1] - 1)      )

									UpperCamelCase_											=				model(
									    input_ids[:, :-1]       ,	attention_mask=snake_case__       ,	past_key_values=snake_case__       ,	position_ids=snake_case__       ,	)
									UpperCamelCase_											=				jnp.array(input_ids.shape[0] * [[input_ids.shape[-1] - 1]]       ,	dtype="i4"      )
									UpperCamelCase_											=				model(
									    input_ids[:, -1:]       ,	past_key_values=outputs_cache.past_key_values       ,	attention_mask=snake_case__       ,	position_ids=snake_case__       ,	)

									UpperCamelCase_											=				model(snake_case__       ,	attention_mask=snake_case__      )

									UpperCamelCase_											=				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				(a_ ,				a_ ,				unittest.TestCase ):
				'''simple docstring'''







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

				def  _lowerCamelCase   (	self      ):

									'''simple docstring'''


									UpperCamelCase_											=				FlaxGPTaModelTester(self      )

				def  _lowerCamelCase   (	self      ):

									'''simple docstring'''


									for model_class_name in self.all_model_classes:
														UpperCamelCase_					,      UpperCamelCase_					,      UpperCamelCase_											=				self.model_tester.prepare_config_and_inputs()
														self.model_tester.check_use_cache_forward(snake_case__       ,	snake_case__       ,	snake_case__       ,	snake_case__      )

				def  _lowerCamelCase   (	self      ):

									'''simple docstring'''


									for model_class_name in self.all_model_classes:
														UpperCamelCase_					,      UpperCamelCase_					,      UpperCamelCase_											=				self.model_tester.prepare_config_and_inputs()
														self.model_tester.check_use_cache_forward_with_attn_mask(
														    snake_case__       ,	snake_case__       ,	snake_case__       ,	snake_case__      )

				@slow
				def  _lowerCamelCase   (	self      ):

									'''simple docstring'''


									UpperCamelCase_											=				GPTaTokenizer.from_pretrained("gpt2"       ,	pad_token="</s>"       ,	padding_side="left"      )
									UpperCamelCase_											=				tokenizer(["Hello this is a long string", "Hey"]       ,	return_tensors="np"       ,	padding=snake_case__       ,	truncation=snake_case__      )

									UpperCamelCase_											=				FlaxGPTaLMHeadModel.from_pretrained("gpt2"      )
									UpperCamelCase_											=				False
									UpperCamelCase_											=				model.config.eos_token_id

									UpperCamelCase_											=				jax.jit(model.generate      )

									UpperCamelCase_											=				jit_generate(inputs["input_ids"]       ,	attention_mask=inputs["attention_mask"]      ).sequences

									UpperCamelCase_											=				tokenizer.batch_decode(snake_case__       ,	skip_special_tokens=snake_case__      )

									UpperCamelCase_											=				[
									    "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(snake_case__       ,	snake_case__      )

				@is_pt_flax_cross_test
				def  _lowerCamelCase   (	self      ):

									'''simple docstring'''


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

																			# load corresponding PyTorch class
																			UpperCamelCase_											=				model_class.__name__[4:]  # Skip the "Flax" at the beginning
																			UpperCamelCase_											=				getattr(snake_case__       ,	snake_case__      )

																			UpperCamelCase_					,      UpperCamelCase_											=				pt_inputs["input_ids"].shape
																			UpperCamelCase_											=				np.random.randint(0       ,	seq_length - 1       ,	size=(batch_size,)      )
																			for batch_idx, start_index in enumerate(snake_case__      ):
																								UpperCamelCase_											=				0
																								UpperCamelCase_											=				1
																								UpperCamelCase_											=				0
																								UpperCamelCase_											=				1
																			UpperCamelCase_											=				pt_model_class(snake_case__      ).eval()
																			UpperCamelCase_											=				model_class(snake_case__       ,	dtype=jnp.floataa      )

																			UpperCamelCase_											=				convert_pytorch_state_dict_to_flax(pt_model.state_dict()       ,	snake_case__      )
																			UpperCamelCase_											=				fx_state

																			with torch.no_grad():
																								UpperCamelCase_											=				pt_model(**snake_case__      ).to_tuple()

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

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

																			UpperCamelCase_											=				fx_model_loaded(**snake_case__      ).to_tuple()
																			self.assertEqual(
																			    len(snake_case__      )       ,	len(snake_case__      )       ,	"Output lengths differ between Flax and PyTorch"      )
																			for fx_output_loaded, pt_output in zip(snake_case__       ,	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'''


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

																			# load corresponding PyTorch class
																			UpperCamelCase_											=				model_class.__name__[4:]  # Skip the "Flax" at the beginning
																			UpperCamelCase_											=				getattr(snake_case__       ,	snake_case__      )

																			UpperCamelCase_											=				pt_model_class(snake_case__      ).eval()
																			UpperCamelCase_											=				model_class(snake_case__       ,	dtype=jnp.floataa      )

																			UpperCamelCase_											=				load_flax_weights_in_pytorch_model(snake_case__       ,	fx_model.params      )
																			UpperCamelCase_					,      UpperCamelCase_											=				pt_inputs["input_ids"].shape
																			UpperCamelCase_											=				np.random.randint(0       ,	seq_length - 1       ,	size=(batch_size,)      )
																			for batch_idx, start_index in enumerate(snake_case__      ):
																								UpperCamelCase_											=				0
																								UpperCamelCase_											=				1
																								UpperCamelCase_											=				0
																								UpperCamelCase_											=				1

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

																			with torch.no_grad():
																								UpperCamelCase_											=				pt_model(**snake_case__      ).to_tuple()

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

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

																			with torch.no_grad():
																								UpperCamelCase_											=				pt_model_loaded(**snake_case__      ).to_tuple()

																			self.assertEqual(
																			    len(snake_case__      )       ,	len(snake_case__      )       ,	"Output lengths differ between Flax and PyTorch"      )
																			for fx_output, pt_output in zip(snake_case__       ,	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:
														UpperCamelCase_											=				model_class_name.from_pretrained("gpt2"       ,	from_pt=snake_case__      )
														UpperCamelCase_											=				model(np.ones((1, 1)      )      )
														self.assertIsNotNone(snake_case__      )
