




"""simple docstring"""

import tempfile
import unittest

import transformers
from transformers import XGLMConfig, XGLMTokenizer, is_flax_available, is_torch_available
from transformers.testing_utils import is_pt_flax_cross_test, require_flax, require_sentencepiece, 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
						import numpy as np

						from transformers.modeling_flax_pytorch_utils import (
						    convert_pytorch_state_dict_to_flax,
						    load_flax_weights_in_pytorch_model,
						)
						from transformers.models.xglm.modeling_flax_xglm import FlaxXGLMForCausalLM, FlaxXGLMModel


if is_torch_available():
						import torch



@require_flax
class  _snake_case				:

				'''simple docstring'''




				def __init__(							self      :  Dict					,							snake_case      :  Tuple					,							snake_case      :  Union[str, Any]=14					,							snake_case      :  str=7					,							snake_case      :  int=True					,							snake_case      :  Dict=True					,							snake_case      :  Optional[int]=True					,							snake_case      :  Dict=99					,							snake_case      :  Dict=32					,							snake_case      :  Dict=5					,							snake_case      :  Dict=4					,							snake_case      :  Union[str, Any]=37					,							snake_case      :  Optional[int]="gelu"					,							snake_case      :  Optional[Any]=0.1					,							snake_case      :  Any=0.1					,							snake_case      :  Any=512					,							snake_case      :  int=0.02					,							snake_case      :  Dict=None					,							):
								UpperCAmelCase_					:List[Any]      			=      parent
								UpperCAmelCase_					:Optional[Any]      			=      batch_size
								UpperCAmelCase_					:Tuple      			=      seq_length
								UpperCAmelCase_					:List[Any]      			=      is_training
								UpperCAmelCase_					:Dict      			=      use_input_mask
								UpperCAmelCase_					:List[str]      			=      use_labels
								UpperCAmelCase_					:str      			=      vocab_size
								UpperCAmelCase_					:str      			=      d_model
								UpperCAmelCase_					:Union[str, Any]      			=      num_hidden_layers
								UpperCAmelCase_					:Tuple      			=      num_attention_heads
								UpperCAmelCase_					:Tuple      			=      ffn_dim
								UpperCAmelCase_					:List[str]      			=      activation_function
								UpperCAmelCase_					:Any      			=      activation_dropout
								UpperCAmelCase_					:Dict      			=      attention_dropout
								UpperCAmelCase_					:Any      			=      max_position_embeddings
								UpperCAmelCase_					:Any      			=      initializer_range
								UpperCAmelCase_					:Optional[Any]      			=      None
								UpperCAmelCase_					:Any      			=      0
								UpperCAmelCase_					:int      			=      2
								UpperCAmelCase_					:int      			=      1





				def 				snake_case_ (							self      :  str     ):
								UpperCAmelCase_					:Union[str, Any]      			=      np.clip(ids_tensor([self.batch_size, self.seq_length]					,							self.vocab_size     )					,							3					,							self.vocab_size     )

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

								UpperCAmelCase_					:List[Any]      			=      XGLMConfig(
								    vocab_size=self.vocab_size					,							d_model=self.hidden_size					,							num_layers=self.num_hidden_layers					,							attention_heads=self.num_attention_heads					,							ffn_dim=self.ffn_dim					,							activation_function=self.activation_function					,							activation_dropout=self.activation_dropout					,							attention_dropout=self.attention_dropout					,							max_position_embeddings=self.max_position_embeddings					,							initializer_range=self.initializer_range					,							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 				snake_case_ (							self      :  Dict     ):
								UpperCAmelCase_					:Union[str, Any]      			=      self.prepare_config_and_inputs()
								UpperCAmelCase_					:Optional[Any]      			=      config_and_inputs
								UpperCAmelCase_					:str      			=      {'''input_ids''': input_ids, '''attention_mask''': attention_mask}
								return config, inputs_dict





				def 				snake_case_ (							self      :  List[str]     ):
								UpperCAmelCase_					:Union[str, Any]      			=      self.prepare_config_and_inputs()

								UpperCAmelCase_					:Tuple      			=      floats_tensor([self.batch_size, self.seq_length, self.hidden_size]     )
								UpperCAmelCase_					:Tuple      			=      ids_tensor([self.batch_size, self.seq_length]					,							vocab_size=2     )

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





				def 				snake_case_ (							self      :  Dict					,							snake_case      :  Union[str, Any]					,							snake_case      :  List[Any]					,							snake_case      :  Union[str, Any]					,							snake_case      :  Optional[Any]     ):
								UpperCAmelCase_					:List[Any]      			=      20
								UpperCAmelCase_					:List[str]      			=      model_class_name(snake_case     )

								UpperCAmelCase_					:List[str]      			=      model.init_cache(input_ids.shape[0]					,							snake_case     )
								UpperCAmelCase_					:Tuple      			=      jnp.ones((input_ids.shape[0], max_decoder_length)					,							dtype='''i4'''     )

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

								UpperCAmelCase_					:Dict      			=      jnp.array(input_ids.shape[0] * [[input_ids.shape[-1] - 1]]					,							dtype='''i4'''     )
								UpperCAmelCase_					:List[str]      			=      model(
								    input_ids[:, -1:]					,							attention_mask=snake_case					,							past_key_values=outputs_cache.past_key_values					,							position_ids=snake_case					,							)

								UpperCAmelCase_					:int      			=      model(snake_case     )

								UpperCAmelCase_					: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}'     )





				def 				snake_case_ (							self      :  str					,							snake_case      :  Tuple					,							snake_case      :  List[str]					,							snake_case      :  List[str]					,							snake_case      :  Optional[Any]     ):
								UpperCAmelCase_					:Tuple      			=      20
								UpperCAmelCase_					:Dict      			=      model_class_name(snake_case     )

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

								UpperCAmelCase_					:Optional[int]      			=      model.init_cache(input_ids.shape[0]					,							snake_case     )
								UpperCAmelCase_					:str      			=      jnp.broadcast_to(
								    jnp.arange(input_ids.shape[-1] - 1     )[None, :]					,							(input_ids.shape[0], input_ids.shape[-1] - 1)     )

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

								UpperCAmelCase_					:Union[str, Any]      			=      model(snake_case					,							attention_mask=snake_case     )
								UpperCAmelCase_					:Dict      			=      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_sentencepiece
@require_flax
class  _snake_case				(							A__							,						A__							,						unittest.TestCase       ):

				'''simple docstring'''




				UpperCamelCase__		     =(FlaxXGLMModel, FlaxXGLMForCausalLM) if is_flax_available() else ()
				UpperCamelCase__		     =(FlaxXGLMForCausalLM,) if is_flax_available() else ()





				def 				snake_case_ (							self      :  Optional[int]     ):
								UpperCAmelCase_					:Optional[Any]      			=      FlaxXGLMModelTester(self     )





				def 				snake_case_ (							self      :  Any     ):
								for model_class_name in self.all_model_classes:
												UpperCAmelCase_					:Union[str, Any]      			=      self.model_tester.prepare_config_and_inputs()
												self.model_tester.check_use_cache_forward(snake_case					,							snake_case					,							snake_case					,							snake_case     )





				def 				snake_case_ (							self      :  str     ):
								for model_class_name in self.all_model_classes:
												UpperCAmelCase_					:Optional[int]      			=      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 				snake_case_ (							self      :  Union[str, Any]     ):
								UpperCAmelCase_					:Optional[int]      			=      XGLMTokenizer.from_pretrained('''XGLM'''					,							padding_side='''left'''     )
								UpperCAmelCase_					:Any      			=      tokenizer(['''Hello this is a long string''', '''Hey''']					,							return_tensors='''np'''					,							padding=snake_case					,							truncation=snake_case     )

								UpperCAmelCase_					:Any      			=      FlaxXGLMForCausalLM.from_pretrained('''facebook/xglm-564M'''     )
								UpperCAmelCase_					:int      			=      1
								UpperCAmelCase_					:Dict      			=      False

								UpperCAmelCase_					:Optional[int]      			=      jax.jit(model.generate     )

								UpperCAmelCase_					:List[str]      			=      jit_generate(inputs['''input_ids''']					,							attention_mask=inputs['''attention_mask''']     ).sequences

								UpperCAmelCase_					:List[str]      			=      tokenizer.batch_decode(snake_case					,							skip_special_tokens=snake_case     )

								UpperCAmelCase_					:str      			=      [
								    '''Hello this is a long string of questions, but I\'m not sure if I\'m''',
								    '''Hey, I\'m a newbie to the forum and I\'''',
								]

								self.assertListEqual(snake_case					,							snake_case     )





				@is_pt_flax_cross_test
				def 				snake_case_ (							self      :  int     ):
								UpperCAmelCase_					: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
																UpperCAmelCase_					:Any      			=      self._prepare_for_class(snake_case					,							snake_case     )
																UpperCAmelCase_					:Optional[int]      			=      {k: torch.tensor(v.tolist()     ) for k, v in prepared_inputs_dict.items()}

																# load corresponding PyTorch class
																UpperCAmelCase_					:Union[str, Any]      			=      model_class.__name__[4:]  # Skip the "Flax" at the beginning
																UpperCAmelCase_					:Tuple      			=      getattr(snake_case					,							snake_case     )

																UpperCAmelCase_					:Dict      			=      pt_inputs['''input_ids'''].shape
																UpperCAmelCase_					:List[Any]      			=      np.random.randint(0					,							seq_length - 1					,							size=(batch_size,)     )
																for batch_idx, start_index in enumerate(snake_case     ):
																				UpperCAmelCase_					:List[str]      			=      0
																				UpperCAmelCase_					:int      			=      1
																				UpperCAmelCase_					:str      			=      0
																				UpperCAmelCase_					:List[Any]      			=      1
																UpperCAmelCase_					:List[Any]      			=      pt_model_class(snake_case     ).eval()
																# Flax models don't use the `use_cache` option and cache is not returned as a default.
																# So we disable `use_cache` here for PyTorch model.
																UpperCAmelCase_					:Tuple      			=      False
																UpperCAmelCase_					:Union[str, Any]      			=      model_class(snake_case					,							dtype=jnp.floataa     )
																UpperCAmelCase_					:int      			=      convert_pytorch_state_dict_to_flax(pt_model.state_dict()					,							snake_case     )
																UpperCAmelCase_					:Dict      			=      fx_state

																with torch.no_grad():
																				UpperCAmelCase_					:Union[str, Any]      			=      pt_model(**snake_case     ).to_tuple()

																UpperCAmelCase_					:List[str]      			=      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_					:Dict      			=      model_class.from_pretrained(snake_case					,							from_pt=snake_case     )

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

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

																UpperCAmelCase_					:int      			=      pt_model_class(snake_case     ).eval()
																UpperCAmelCase_					:Union[str, Any]      			=      False
																UpperCAmelCase_					:Optional[Any]      			=      model_class(snake_case					,							dtype=jnp.floataa     )

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

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

																with torch.no_grad():
																				UpperCAmelCase_					:Optional[int]      			=      pt_model(**snake_case     ).to_tuple()

																UpperCAmelCase_					:Any      			=      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_					:Dict      			=      pt_model_class.from_pretrained(snake_case					,							from_flax=snake_case     )

																with torch.no_grad():
																				UpperCAmelCase_					:str      			=      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 				snake_case_ (							self      :  Any     ):
								for model_class_name in self.all_model_classes:
												UpperCAmelCase_					:Optional[Any]      			=      model_class_name.from_pretrained('''facebook/xglm-564M'''     )
												UpperCAmelCase_					:int      			=      model(np.ones((1, 1)     )     )
												self.assertIsNotNone(snake_case     )












"""simple docstring"""

from collections import OrderedDict
from typing import Mapping

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


__lowerCamelCase								=						logging.get_logger(__name__)

__lowerCamelCase								=						{
    "mnaylor/mega-base-wikitext": "https://huggingface.co/mnaylor/mega-base-wikitext/resolve/main/config.json",
}



class  _snake_case				(							A__       ):

				'''simple docstring'''




				UpperCamelCase__		     ="""mega"""





				def __init__(							self      :  Optional[int]					,							snake_case      :  Union[str, Any]=30_522					,							snake_case      :  Optional[Any]=128					,							snake_case      :  str=4					,							snake_case      :  Dict=256					,							snake_case      :  Optional[Any]=16					,							snake_case      :  Any=True					,							snake_case      :  str=64					,							snake_case      :  Optional[Any]=False					,							snake_case      :  Union[str, Any]=-1					,							snake_case      :  List[str]=None					,							snake_case      :  Any=True					,							snake_case      :  Optional[int]="scalenorm"					,							snake_case      :  Tuple=True					,							snake_case      :  str="silu"					,							snake_case      :  Dict="softmax"					,							snake_case      :  Optional[int]=0.1					,							snake_case      :  Tuple=0.1					,							snake_case      :  Dict=0.1					,							snake_case      :  List[str]=False					,							snake_case      :  Any=True					,							snake_case      :  List[Any]=256					,							snake_case      :  Optional[Any]=True					,							snake_case      :  Optional[int]=0.1					,							snake_case      :  str=2_048					,							snake_case      :  Union[str, Any]=False					,							snake_case      :  List[Any]=2					,							snake_case      :  str=0.02					,							snake_case      :  str=0.2					,							snake_case      :  str=0.02					,							snake_case      :  Union[str, Any]=1.0					,							snake_case      :  Optional[Any]=1					,							snake_case      :  List[Any]=0					,							snake_case      :  List[Any]=2					,							snake_case      :  List[Any]="rotary"					,							snake_case      :  Tuple=None					,							snake_case      :  int=True					,							snake_case      :  Dict=True					,							**snake_case      :  List[str]					,							):
								super().__init__(pad_token_id=snake_case					,							bos_token_id=snake_case					,							eos_token_id=snake_case					,							**snake_case     )

								UpperCAmelCase_					:List[Any]      			=      vocab_size
								UpperCAmelCase_					:Union[str, Any]      			=      hidden_size
								UpperCAmelCase_					:str      			=      num_hidden_layers
								UpperCAmelCase_					:List[Any]      			=      activation
								UpperCAmelCase_					:int      			=      attention_activation
								UpperCAmelCase_					:Tuple      			=      intermediate_size
								UpperCAmelCase_					:Optional[Any]      			=      ema_projection_size
								UpperCAmelCase_					:Optional[int]      			=      bidirectional
								UpperCAmelCase_					:Optional[int]      			=      shared_representation_size
								UpperCAmelCase_					:Optional[Any]      			=      use_chunking
								UpperCAmelCase_					:List[str]      			=      chunk_size
								UpperCAmelCase_					:Optional[Any]      			=      truncation
								UpperCAmelCase_					:List[str]      			=      normalize_before_mega
								UpperCAmelCase_					:List[Any]      			=      normalization_type
								UpperCAmelCase_					:Tuple      			=      norm_affine
								UpperCAmelCase_					:Union[str, Any]      			=      dropout_prob
								UpperCAmelCase_					:Tuple      			=      hidden_dropout_prob
								UpperCAmelCase_					:List[Any]      			=      attention_probs_dropout_prob
								UpperCAmelCase_					:str      			=      use_feature_dropout
								UpperCAmelCase_					:List[Any]      			=      use_normalized_ffn
								UpperCAmelCase_					:Any      			=      nffn_hidden_size
								UpperCAmelCase_					:List[Any]      			=      normalize_before_ffn
								UpperCAmelCase_					:Optional[int]      			=      nffn_activation_dropout_prob
								UpperCAmelCase_					:str      			=      max_positions
								UpperCAmelCase_					:int      			=      add_token_type_embeddings
								UpperCAmelCase_					:List[Any]      			=      type_vocab_size
								UpperCAmelCase_					:List[str]      			=      initializer_range
								UpperCAmelCase_					:int      			=      ema_delta_alpha_range
								UpperCAmelCase_					:Optional[int]      			=      ema_beta_range
								UpperCAmelCase_					:int      			=      ema_gamma_omega_range
								UpperCAmelCase_					:int      			=      relative_positional_bias
								UpperCAmelCase_					:List[Any]      			=      use_cache
								UpperCAmelCase_					:Optional[Any]      			=      classifier_dropout
								UpperCAmelCase_					:Optional[Any]      			=      add_lm_hidden_dense_layer
								UpperCAmelCase_					:str      			=      1  # not used but required by Hugging Face



class  _snake_case				(							A__       ):

				'''simple docstring'''




				@property
				def 				snake_case_ (							self      :  int     ):
								if self.task == "multiple-choice":
												UpperCAmelCase_					:List[str]      			=      {0: '''batch''', 1: '''choice''', 2: '''sequence'''}
								else:
												UpperCAmelCase_					:Tuple      			=      {0: '''batch''', 1: '''sequence'''}
								return OrderedDict(
								    [
								        ('''input_ids''', dynamic_axis),
								        ('''attention_mask''', dynamic_axis),
								    ]     )


