






import os
import unittest

from transformers import is_torch_available
from transformers.testing_utils import require_sentencepiece, require_tokenizers, require_torch, slow, torch_device


if is_torch_available():
	import torch

	from transformers import GPTaLMHeadModel


@require_torch
@require_sentencepiece
@require_tokenizers
class        lowercase__       (     unittest.TestCase		):

	@slow
	@unittest.skip("Model is not available."    )
	def 						UpperCAmelCase					(							self    )->   Tuple:







							'''simple docstring'''





							lowerCAmelCase__	 =  '''nvidia/megatron-gpt2-345m/'''
							if "MYDIR" in os.environ:
													lowerCAmelCase__	 =  os.path.join(os.environ["MYDIR"]						,      _A    )
							lowerCAmelCase__	 =  GPTaLMHeadModel.from_pretrained(_A    )
							model.to(_A    )
							model.half()

							lowerCAmelCase__	 =  torch.tensor(
							    [[101, 7110, 1005, 1056, 2023, 11333, 17413, 1029, 102]]						,      device=_A						,      dtype=torch.long						,      )

							with torch.no_grad():
													lowerCAmelCase__	 =  model(_A    ).logits

							lowerCAmelCase__	 =  torch.Size((1, 9, 50257)    )
							self.assertEqual(output.shape						,      _A    )

							lowerCAmelCase__	 =  torch.tensor(
							    [
							        4.9_414,
							        -0.2_920,
							        -1.2_148,
							        -4.0_273,
							        -0.5_161,
							        -5.2_109,
							        -1.2_412,
							        -1.8_301,
							        -1.7_734,
							        -4.7_148,
							        -0.2_317,
							        -1.0_811,
							        -2.1_777,
							        0.4_141,
							        -3.7_969,
							        -4.0_586,
							        -2.5_332,
							        -3.3_809,
							        4.3_867,
							    ]						,      device=_A						,      dtype=torch.half						,      )

							for i in range(19    ):
													lowerCAmelCase__	 =  8 * i // 17, 2792 * i  # along the diagonal
													lowerCAmelCase__	 =  output[0, r, c], expected_diag[i]
													lowerCAmelCase__	 =  F"row={r} col={c} computed={computed} expected={expected}"
													self.assertAlmostEqual(_A						,      _A						,      delta=1E-4						,      msg=_A    )













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        lowercase__       :

	def __init__(							self						,      __UpperCAmelCase						,      __UpperCAmelCase=14						,      __UpperCAmelCase=7						,      __UpperCAmelCase=True						,      __UpperCAmelCase=True						,      __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						,      __UpperCAmelCase=None						,      )->   Optional[Any]:







							'''simple docstring'''





							lowerCAmelCase__	 =  parent
							lowerCAmelCase__	 =  batch_size
							lowerCAmelCase__	 =  seq_length
							lowerCAmelCase__	 =  is_training
							lowerCAmelCase__	 =  use_input_mask
							lowerCAmelCase__	 =  use_labels
							lowerCAmelCase__	 =  vocab_size
							lowerCAmelCase__	 =  d_model
							lowerCAmelCase__	 =  num_hidden_layers
							lowerCAmelCase__	 =  num_attention_heads
							lowerCAmelCase__	 =  ffn_dim
							lowerCAmelCase__	 =  activation_function
							lowerCAmelCase__	 =  activation_dropout
							lowerCAmelCase__	 =  attention_dropout
							lowerCAmelCase__	 =  max_position_embeddings
							lowerCAmelCase__	 =  initializer_range
							lowerCAmelCase__	 =  None
							lowerCAmelCase__	 =  0
							lowerCAmelCase__	 =  2
							lowerCAmelCase__	 =  1

	def 						UpperCAmelCase					(							self    )->   int:







							'''simple docstring'''





							lowerCAmelCase__	 =  np.clip(ids_tensor([self.batch_size, self.seq_length]						,      self.vocab_size    )						,      3						,      self.vocab_size    )

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

							lowerCAmelCase__	 =  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=__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 						UpperCAmelCase					(							self    )->   int:







							'''simple docstring'''





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

	def 						UpperCAmelCase					(							self    )->   Union[str, Any]:







							'''simple docstring'''





							lowerCAmelCase__ ,							lowerCAmelCase__ ,							lowerCAmelCase__	 =  self.prepare_config_and_inputs()

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

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

	def 						UpperCAmelCase					(							self						,      __UpperCAmelCase						,      __UpperCAmelCase						,      __UpperCAmelCase						,      __UpperCAmelCase    )->   str:







							'''simple docstring'''





							lowerCAmelCase__	 =  20
							lowerCAmelCase__	 =  model_class_name(__UpperCAmelCase    )

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

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

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

							lowerCAmelCase__	 =  model(__UpperCAmelCase    )

							lowerCAmelCase__	 =  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 						UpperCAmelCase					(							self						,      __UpperCAmelCase						,      __UpperCAmelCase						,      __UpperCAmelCase						,      __UpperCAmelCase    )->   List[str]:







							'''simple docstring'''





							lowerCAmelCase__	 =  20
							lowerCAmelCase__	 =  model_class_name(__UpperCAmelCase    )

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

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

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

							lowerCAmelCase__	 =  model(__UpperCAmelCase						,      attention_mask=__UpperCAmelCase    )
							lowerCAmelCase__	 =  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        lowercase__       (     _UpperCAmelCase,   _UpperCAmelCase,   unittest.TestCase		):
	a_             =(FlaxXGLMModel, FlaxXGLMForCausalLM) if is_flax_available() else ()
	a_             =(FlaxXGLMForCausalLM,) if is_flax_available() else ()

	def 						UpperCAmelCase					(							self    )->   Optional[int]:







							'''simple docstring'''





							lowerCAmelCase__	 =  FlaxXGLMModelTester(self    )

	def 						UpperCAmelCase					(							self    )->   Optional[int]:







							'''simple docstring'''





							for model_class_name in self.all_model_classes:
													lowerCAmelCase__ ,							lowerCAmelCase__ ,							lowerCAmelCase__	 =  self.model_tester.prepare_config_and_inputs()
													self.model_tester.check_use_cache_forward(__UpperCAmelCase						,      __UpperCAmelCase						,      __UpperCAmelCase						,      __UpperCAmelCase    )

	def 						UpperCAmelCase					(							self    )->   int:







							'''simple docstring'''





							for model_class_name in self.all_model_classes:
													lowerCAmelCase__ ,							lowerCAmelCase__ ,							lowerCAmelCase__	 =  self.model_tester.prepare_config_and_inputs()
													self.model_tester.check_use_cache_forward_with_attn_mask(
													    __UpperCAmelCase						,      __UpperCAmelCase						,      __UpperCAmelCase						,      __UpperCAmelCase    )

	@slow
	def 						UpperCAmelCase					(							self    )->   List[str]:







							'''simple docstring'''





							lowerCAmelCase__	 =  XGLMTokenizer.from_pretrained("XGLM"						,      padding_side="left"    )
							lowerCAmelCase__	 =  tokenizer(["Hello this is a long string", "Hey"]						,      return_tensors="np"						,      padding=__UpperCAmelCase						,      truncation=__UpperCAmelCase    )

							lowerCAmelCase__	 =  FlaxXGLMForCausalLM.from_pretrained("facebook/xglm-564M"    )
							lowerCAmelCase__	 =  1
							lowerCAmelCase__	 =  False

							lowerCAmelCase__	 =  jax.jit(model.generate    )

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

							lowerCAmelCase__	 =  tokenizer.batch_decode(__UpperCAmelCase						,      skip_special_tokens=__UpperCAmelCase    )

							lowerCAmelCase__	 =  [
							    "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(__UpperCAmelCase						,      __UpperCAmelCase    )

	@is_pt_flax_cross_test
	def 						UpperCAmelCase					(							self    )->   Union[str, Any]:







							'''simple docstring'''





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

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

																			lowerCAmelCase__ ,							lowerCAmelCase__	 =  pt_inputs["input_ids"].shape
																			lowerCAmelCase__	 =  np.random.randint(0						,      seq_length - 1						,      size=(batch_size,)    )
																			for batch_idx, start_index in enumerate(__UpperCAmelCase    ):
																									lowerCAmelCase__	 =  0
																									lowerCAmelCase__	 =  1
																									lowerCAmelCase__	 =  0
																									lowerCAmelCase__	 =  1
																			lowerCAmelCase__	 =  pt_model_class(__UpperCAmelCase    ).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.
																			lowerCAmelCase__	 =  False
																			lowerCAmelCase__	 =  model_class(__UpperCAmelCase						,      dtype=jnp.floataa    )
																			lowerCAmelCase__	 =  convert_pytorch_state_dict_to_flax(pt_model.state_dict()						,      __UpperCAmelCase    )
																			lowerCAmelCase__	 =  fx_state

																			with torch.no_grad():
																									lowerCAmelCase__	 =  pt_model(**__UpperCAmelCase    ).to_tuple()

																			lowerCAmelCase__	 =  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    )
																									lowerCAmelCase__	 =  model_class.from_pretrained(__UpperCAmelCase						,      from_pt=__UpperCAmelCase    )

																			lowerCAmelCase__	 =  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 						UpperCAmelCase					(							self    )->   Tuple:







							'''simple docstring'''





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

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

																			lowerCAmelCase__	 =  pt_model_class(__UpperCAmelCase    ).eval()
																			lowerCAmelCase__	 =  False
																			lowerCAmelCase__	 =  model_class(__UpperCAmelCase						,      dtype=jnp.floataa    )

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

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

																			with torch.no_grad():
																									lowerCAmelCase__	 =  pt_model(**__UpperCAmelCase    ).to_tuple()

																			lowerCAmelCase__	 =  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    )
																									lowerCAmelCase__	 =  pt_model_class.from_pretrained(__UpperCAmelCase						,      from_flax=__UpperCAmelCase    )

																			with torch.no_grad():
																									lowerCAmelCase__	 =  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 						UpperCAmelCase					(							self    )->   Dict:







							'''simple docstring'''





							for model_class_name in self.all_model_classes:
													lowerCAmelCase__	 =  model_class_name.from_pretrained("facebook/xglm-564M"    )
													lowerCAmelCase__	 =  model(np.ones((1, 1)    )    )
													self.assertIsNotNone(__UpperCAmelCase    )






