




'''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 					lowerCamelCase_							:




					'''simple docstring'''






					def __init__(   self,  __UpperCamelCase,  __UpperCamelCase=14,  __UpperCamelCase=7,  __UpperCamelCase=True,  __UpperCamelCase=True,  __UpperCamelCase=False,  __UpperCamelCase=True,  __UpperCamelCase=99,  __UpperCamelCase=32,  __UpperCamelCase=5,  __UpperCamelCase=4,  __UpperCamelCase=37,  __UpperCamelCase="gelu",  __UpperCamelCase=0.1,  __UpperCamelCase=0.1,  __UpperCamelCase=512,  __UpperCamelCase=0.02,  ):
							_snake_case     	=	parent
							_snake_case     	=	batch_size
							_snake_case     	=	seq_length
							_snake_case     	=	is_training
							_snake_case     	=	use_input_mask
							_snake_case     	=	use_token_type_ids
							_snake_case     	=	use_labels
							_snake_case     	=	vocab_size
							_snake_case     	=	hidden_size
							_snake_case     	=	num_hidden_layers
							_snake_case     	=	num_attention_heads
							_snake_case     	=	intermediate_size
							_snake_case     	=	hidden_act
							_snake_case     	=	hidden_dropout_prob
							_snake_case     	=	attention_probs_dropout_prob
							_snake_case     	=	max_position_embeddings
							_snake_case     	=	initializer_range
							_snake_case     	=	None
							_snake_case     	=	vocab_size - 1
							_snake_case     	=	vocab_size - 1
							_snake_case     	=	vocab_size - 1






					def     __SCREAMING_SNAKE_CASE     (   self			):
							_snake_case     	=	ids_tensor([self.batch_size, self.seq_length],  self.vocab_size			)

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

							_snake_case     	=	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=__UpperCamelCase,  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     __SCREAMING_SNAKE_CASE     (   self			):
							_snake_case     	=	self.prepare_config_and_inputs()
							_snake_case	,   _snake_case	,   _snake_case     	=	config_and_inputs
							_snake_case     	=	{"input_ids": input_ids, "attention_mask": attention_mask}
							return config, inputs_dict






					def     __SCREAMING_SNAKE_CASE     (   self			):
							_snake_case	,   _snake_case	,   _snake_case     	=	self.prepare_config_and_inputs()

							_snake_case     	=	floats_tensor([self.batch_size, self.seq_length, self.hidden_size]			)
							_snake_case     	=	ids_tensor([self.batch_size, self.seq_length],  vocab_size=2			)

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






					def     __SCREAMING_SNAKE_CASE     (   self,  __UpperCamelCase,  __UpperCamelCase,  __UpperCamelCase,  __UpperCamelCase			):
							_snake_case     	=	20
							_snake_case     	=	model_class_name(__UpperCamelCase			)

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

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

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

							_snake_case     	=	model(__UpperCamelCase			)

							_snake_case     	=	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     __SCREAMING_SNAKE_CASE     (   self,  __UpperCamelCase,  __UpperCamelCase,  __UpperCamelCase,  __UpperCamelCase			):
							_snake_case     	=	20
							_snake_case     	=	model_class_name(__UpperCamelCase			)

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

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

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

							_snake_case     	=	model(__UpperCamelCase,  attention_mask=__UpperCamelCase			)

							_snake_case     	=	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 					lowerCamelCase_							(	UpperCAmelCase_,	UpperCAmelCase_,	unittest.TestCase							):




					'''simple docstring'''






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






					def     __SCREAMING_SNAKE_CASE     (   self			):
							_snake_case     	=	FlaxGPTaModelTester(self			)






					def     __SCREAMING_SNAKE_CASE     (   self			):
							for model_class_name in self.all_model_classes:
									_snake_case	,   _snake_case	,   _snake_case     	=	self.model_tester.prepare_config_and_inputs()
									self.model_tester.check_use_cache_forward(__UpperCamelCase,  __UpperCamelCase,  __UpperCamelCase,  __UpperCamelCase			)






					def     __SCREAMING_SNAKE_CASE     (   self			):
							for model_class_name in self.all_model_classes:
									_snake_case	,   _snake_case	,   _snake_case     	=	self.model_tester.prepare_config_and_inputs()
									self.model_tester.check_use_cache_forward_with_attn_mask(
									    __UpperCamelCase,  __UpperCamelCase,  __UpperCamelCase,  __UpperCamelCase			)






					@slow
					def     __SCREAMING_SNAKE_CASE     (   self			):
							_snake_case     	=	GPTaTokenizer.from_pretrained("gpt2",  pad_token="</s>",  padding_side="left"			)
							_snake_case     	=	tokenizer(["Hello this is a long string", "Hey"],  return_tensors="np",  padding=__UpperCamelCase,  truncation=__UpperCamelCase			)

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

							_snake_case     	=	jax.jit(model.generate			)

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

							_snake_case     	=	tokenizer.batch_decode(__UpperCamelCase,  skip_special_tokens=__UpperCamelCase			)

							_snake_case     	=	[
							    "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(__UpperCamelCase,  __UpperCamelCase			)






					@is_pt_flax_cross_test
					def     __SCREAMING_SNAKE_CASE     (   self			):
							_snake_case	,   _snake_case     	=	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
											_snake_case     	=	self._prepare_for_class(__UpperCamelCase,  __UpperCamelCase			)
											_snake_case     	=	{k: torch.tensor(v.tolist()			) for k, v in prepared_inputs_dict.items()}

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

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

											_snake_case     	=	convert_pytorch_state_dict_to_flax(pt_model.state_dict(),  __UpperCamelCase			)
											_snake_case     	=	fx_state

											with torch.no_grad():
													_snake_case     	=	pt_model(**__UpperCamelCase			).to_tuple()

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

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

											_snake_case     	=	fx_model_loaded(**__UpperCamelCase			).to_tuple()
											self.assertEqual(
											    len(__UpperCamelCase			),  len(__UpperCamelCase			),  "Output lengths differ between Flax and PyTorch"			)
											for fx_output_loaded, pt_output in zip(__UpperCamelCase,  __UpperCamelCase			):
													self.assert_almost_equals(fx_output_loaded[:, -1],  pt_output[:, -1].numpy(),  4E-2			)






					@is_pt_flax_cross_test
					def     __SCREAMING_SNAKE_CASE     (   self			):
							_snake_case	,   _snake_case     	=	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
											_snake_case     	=	self._prepare_for_class(__UpperCamelCase,  __UpperCamelCase			)
											_snake_case     	=	{k: torch.tensor(v.tolist()			) for k, v in prepared_inputs_dict.items()}

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

											_snake_case     	=	pt_model_class(__UpperCamelCase			).eval()
											_snake_case     	=	model_class(__UpperCamelCase,  dtype=jnp.floataa			)

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

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

											with torch.no_grad():
													_snake_case     	=	pt_model(**__UpperCamelCase			).to_tuple()

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

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

											with torch.no_grad():
													_snake_case     	=	pt_model_loaded(**__UpperCamelCase			).to_tuple()

											self.assertEqual(
											    len(__UpperCamelCase			),  len(__UpperCamelCase			),  "Output lengths differ between Flax and PyTorch"			)
											for fx_output, pt_output in zip(__UpperCamelCase,  __UpperCamelCase			):
													self.assert_almost_equals(fx_output[:, -1],  pt_output[:, -1].numpy(),  4E-2			)






					@slow
					def     __SCREAMING_SNAKE_CASE     (   self			):
							for model_class_name in self.all_model_classes:
									_snake_case     	=	model_class_name.from_pretrained("gpt2",  from_pt=__UpperCamelCase			)
									_snake_case     	=	model(np.ones((1, 1)			)			)
									self.assertIsNotNone(__UpperCamelCase			)











'''simple docstring'''







from collections import OrderedDict
from typing import Mapping

from ...configuration_utils import PretrainedConfig
from ...onnx import OnnxConfig
from ...utils import logging


_lowerCamelCase							:						Optional[int]							     =  logging.get_logger(__name__)

_lowerCamelCase							:						int							     =  {
    'flaubert/flaubert_small_cased': 'https://huggingface.co/flaubert/flaubert_small_cased/resolve/main/config.json',
    'flaubert/flaubert_base_uncased': 'https://huggingface.co/flaubert/flaubert_base_uncased/resolve/main/config.json',
    'flaubert/flaubert_base_cased': 'https://huggingface.co/flaubert/flaubert_base_cased/resolve/main/config.json',
    'flaubert/flaubert_large_cased': 'https://huggingface.co/flaubert/flaubert_large_cased/resolve/main/config.json',
}





class 					lowerCamelCase_							(	UpperCAmelCase_							):




					'''simple docstring'''






					lowercase__   =			"""flaubert"""
					lowercase__   =			{
					    """hidden_size""": """emb_dim""",
					    """num_attention_heads""": """n_heads""",
					    """num_hidden_layers""": """n_layers""",
					    """n_words""": """vocab_size""",  # For backward compatibility
					}






					def __init__(   self,  __UpperCamelCase=False,  __UpperCamelCase=0.0,  __UpperCamelCase=30_145,  __UpperCamelCase=2_048,  __UpperCamelCase=12,  __UpperCamelCase=16,  __UpperCamelCase=0.1,  __UpperCamelCase=0.1,  __UpperCamelCase=True,  __UpperCamelCase=False,  __UpperCamelCase=False,  __UpperCamelCase=False,  __UpperCamelCase=1,  __UpperCamelCase=True,  __UpperCamelCase=512,  __UpperCamelCase=2_048**-0.5,  __UpperCamelCase=1E-12,  __UpperCamelCase=0.02,  __UpperCamelCase=0,  __UpperCamelCase=1,  __UpperCamelCase=2,  __UpperCamelCase=3,  __UpperCamelCase=5,  __UpperCamelCase=True,  __UpperCamelCase="first",  __UpperCamelCase=True,  __UpperCamelCase=None,  __UpperCamelCase=True,  __UpperCamelCase=0.1,  __UpperCamelCase=5,  __UpperCamelCase=5,  __UpperCamelCase=0,  __UpperCamelCase=0,  __UpperCamelCase=2,  __UpperCamelCase=0,  **__UpperCamelCase,  ):
							_snake_case     	=	pre_norm
							_snake_case     	=	layerdrop
							_snake_case     	=	vocab_size
							_snake_case     	=	emb_dim
							_snake_case     	=	n_layers
							_snake_case     	=	n_heads
							_snake_case     	=	dropout
							_snake_case     	=	attention_dropout
							_snake_case     	=	gelu_activation
							_snake_case     	=	sinusoidal_embeddings
							_snake_case     	=	causal
							_snake_case     	=	asm
							_snake_case     	=	n_langs
							_snake_case     	=	use_lang_emb
							_snake_case     	=	layer_norm_eps
							_snake_case     	=	bos_index
							_snake_case     	=	eos_index
							_snake_case     	=	pad_index
							_snake_case     	=	unk_index
							_snake_case     	=	mask_index
							_snake_case     	=	is_encoder
							_snake_case     	=	max_position_embeddings
							_snake_case     	=	embed_init_std
							_snake_case     	=	init_std
							_snake_case     	=	summary_type
							_snake_case     	=	summary_use_proj
							_snake_case     	=	summary_activation
							_snake_case     	=	summary_proj_to_labels
							_snake_case     	=	summary_first_dropout
							_snake_case     	=	start_n_top
							_snake_case     	=	end_n_top
							_snake_case     	=	mask_token_id
							_snake_case     	=	lang_id

							if "n_words" in kwargs:
									_snake_case     	=	kwargs["n_words"]

							super().__init__(pad_token_id=__UpperCamelCase,  bos_token_id=__UpperCamelCase,  **__UpperCamelCase			)





class 					lowerCamelCase_							(	UpperCAmelCase_							):




					'''simple docstring'''






					@property
					def     __SCREAMING_SNAKE_CASE     (   self			):
							if self.task == "multiple-choice":
									_snake_case     	=	{0: "batch", 1: "choice", 2: "sequence"}
							else:
									_snake_case     	=	{0: "batch", 1: "sequence"}
							return OrderedDict(
							    [
							        ("input_ids", dynamic_axis),
							        ("attention_mask", dynamic_axis),
							    ]			)





