






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






 '''simple docstring'''




 def __init__(		self							,					_A							,					_A=14							,					_A=7							,					_A=True							,					_A=True							,					_A=True							,					_A=99							,					_A=32							,					_A=5							,					_A=4							,					_A=37							,					_A="gelu"							,					_A=0.1							,					_A=0.1							,					_A=512							,					_A=0.02							,					_A=None							,					)	->						str:

     """simple docstring"""



     _UpperCAmelCase			: Union[str, Any]     					=			parent
     _UpperCAmelCase			: Union[str, Any]     					=			batch_size
     _UpperCAmelCase			: Optional[int]     					=			seq_length
     _UpperCAmelCase			: List[str]     					=			is_training
     _UpperCAmelCase			: Any     					=			use_input_mask
     _UpperCAmelCase			: str     					=			use_labels
     _UpperCAmelCase			: int     					=			vocab_size
     _UpperCAmelCase			: Any     					=			d_model
     _UpperCAmelCase			: Optional[Any]     					=			num_hidden_layers
     _UpperCAmelCase			: Any     					=			num_attention_heads
     _UpperCAmelCase			: Any     					=			ffn_dim
     _UpperCAmelCase			: List[str]     					=			activation_function
     _UpperCAmelCase			: Optional[Any]     					=			activation_dropout
     _UpperCAmelCase			: int     					=			attention_dropout
     _UpperCAmelCase			: Optional[Any]     					=			max_position_embeddings
     _UpperCAmelCase			: Union[str, Any]     					=			initializer_range
     _UpperCAmelCase			: List[Any]     					=			None
     _UpperCAmelCase			: Dict     					=			0
     _UpperCAmelCase			: int     					=			2
     _UpperCAmelCase			: int     					=			1


 def 	snake_case__			(		self)	->						Dict:

     """simple docstring"""



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

     _UpperCAmelCase			: int     					=			None
     if self.use_input_mask:
         _UpperCAmelCase			: List[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=__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 	snake_case__			(		self)	->						Dict:

     """simple docstring"""



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


 def 	snake_case__			(		self)	->						str:

     """simple docstring"""



     _UpperCAmelCase					,	_UpperCAmelCase					,	_UpperCAmelCase			: int     					=			self.prepare_config_and_inputs()

     _UpperCAmelCase			: str     					=			floats_tensor([self.batch_size, self.seq_length, self.hidden_size])
     _UpperCAmelCase			: Dict     					=			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							,					_A							,					_A							,					_A							,					_A)	->						Optional[Any]:

     """simple docstring"""



     _UpperCAmelCase			: Union[str, Any]     					=			20
     _UpperCAmelCase			: Optional[Any]     					=			model_class_name(__UpperCamelCase)

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

     _UpperCAmelCase			: Tuple     					=			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=__UpperCamelCase							,					past_key_values=__UpperCamelCase							,					position_ids=__UpperCamelCase							,					)

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

     _UpperCAmelCase			: Dict     					=			model(__UpperCamelCase)

     _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}''')


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

     """simple docstring"""



     _UpperCAmelCase			: Union[str, Any]     					=			20
     _UpperCAmelCase			: str     					=			model_class_name(__UpperCamelCase)

     _UpperCAmelCase			: List[str]     					=			jnp.concatenate(
         [attention_mask, jnp.zeros((attention_mask.shape[0], max_decoder_length - attention_mask.shape[1]))]							,					axis=-1							,					)

     _UpperCAmelCase			: Union[str, Any]     					=			model.init_cache(input_ids.shape[0]							,					__UpperCamelCase)
     _UpperCAmelCase			: List[Any]     					=			jnp.broadcast_to(
         jnp.arange(input_ids.shape[-1] - 1)[None, :]							,					(input_ids.shape[0], input_ids.shape[-1] - 1))

     _UpperCAmelCase			: Dict     					=			model(
         input_ids[:, :-1]							,					attention_mask=__UpperCamelCase							,					past_key_values=__UpperCamelCase							,					position_ids=__UpperCamelCase							,					)
     _UpperCAmelCase			: Any     					=			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=__UpperCamelCase							,					position_ids=__UpperCamelCase							,					)

     _UpperCAmelCase			: Tuple     					=			model(__UpperCamelCase							,					attention_mask=__UpperCamelCase)
     _UpperCAmelCase			: 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_sentencepiece
@require_flax
class 							A_      (  lowercase_	,				lowercase_	,				unittest.TestCase     ):






 '''simple docstring'''




 _SCREAMING_SNAKE_CASE							:     int	=   (FlaxXGLMModel, FlaxXGLMForCausalLM) if is_flax_available() else ()
 _SCREAMING_SNAKE_CASE							:     Any	=   (FlaxXGLMForCausalLM,) if is_flax_available() else ()


 def 	snake_case__			(		self)	->						List[Any]:

     """simple docstring"""



     _UpperCAmelCase			: List[str]     					=			FlaxXGLMModelTester(self)


 def 	snake_case__			(		self)	->						Optional[Any]:

     """simple docstring"""



     for model_class_name in self.all_model_classes:
         _UpperCAmelCase					,	_UpperCAmelCase					,	_UpperCAmelCase			: Any     					=			self.model_tester.prepare_config_and_inputs()
         self.model_tester.check_use_cache_forward(__UpperCamelCase							,					__UpperCamelCase							,					__UpperCamelCase							,					__UpperCamelCase)


 def 	snake_case__			(		self)	->						Any:

     """simple docstring"""



     for model_class_name in self.all_model_classes:
         _UpperCAmelCase					,	_UpperCAmelCase					,	_UpperCAmelCase			: Union[str, Any]     					=			self.model_tester.prepare_config_and_inputs()
         self.model_tester.check_use_cache_forward_with_attn_mask(
             __UpperCamelCase							,					__UpperCamelCase							,					__UpperCamelCase							,					__UpperCamelCase)


 @slow
 def 	snake_case__			(		self)	->						Any:

     """simple docstring"""



     _UpperCAmelCase			: Union[str, Any]     					=			XGLMTokenizer.from_pretrained('''XGLM'''							,					padding_side='''left''')
     _UpperCAmelCase			: List[Any]     					=			tokenizer(['''Hello this is a long string''', '''Hey''']							,					return_tensors='''np'''							,					padding=__UpperCamelCase							,					truncation=__UpperCamelCase)

     _UpperCAmelCase			: List[str]     					=			FlaxXGLMForCausalLM.from_pretrained('''facebook/xglm-564M''')
     _UpperCAmelCase			: str     					=			1
     _UpperCAmelCase			: List[str]     					=			False

     _UpperCAmelCase			: List[Any]     					=			jax.jit(model.generate)

     _UpperCAmelCase			: Dict     					=			jit_generate(inputs['''input_ids''']							,					attention_mask=inputs['''attention_mask''']).sequences

     _UpperCAmelCase			: Optional[int]     					=			tokenizer.batch_decode(__UpperCamelCase							,					skip_special_tokens=__UpperCamelCase)

     _UpperCAmelCase			: Any     					=			[
         '''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 	snake_case__			(		self)	->						Dict:

     """simple docstring"""



     _UpperCAmelCase					,	_UpperCAmelCase			: List[str]     					=			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(__UpperCamelCase							,					__UpperCamelCase)
             _UpperCAmelCase			: List[str]     					=			{k: torch.tensor(v.tolist()) for k, v in prepared_inputs_dict.items()}

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

             _UpperCAmelCase					,	_UpperCAmelCase			: int     					=			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(__UpperCamelCase):
                 _UpperCAmelCase			: int     					=			0
                 _UpperCAmelCase			: Optional[int]     					=			1
                 _UpperCAmelCase			: Optional[int]     					=			0
                 _UpperCAmelCase			: Optional[int]     					=			1
             _UpperCAmelCase			: Optional[int]     					=			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.
             _UpperCAmelCase			: Optional[int]     					=			False
             _UpperCAmelCase			: Dict     					=			model_class(__UpperCamelCase							,					dtype=jnp.floataa)
             _UpperCAmelCase			: List[str]     					=			convert_pytorch_state_dict_to_flax(pt_model.state_dict()							,					__UpperCamelCase)
             _UpperCAmelCase			: int     					=			fx_state

             with torch.no_grad():
                 _UpperCAmelCase			: List[str]     					=			pt_model(**__UpperCamelCase).to_tuple()

             _UpperCAmelCase			: List[Any]     					=			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)
                 _UpperCAmelCase			: Optional[int]     					=			model_class.from_pretrained(__UpperCamelCase							,					from_pt=__UpperCamelCase)

             _UpperCAmelCase			: Optional[Any]     					=			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 	snake_case__			(		self)	->						int:

     """simple docstring"""



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

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

             _UpperCAmelCase			: Any     					=			pt_model_class(__UpperCamelCase).eval()
             _UpperCAmelCase			: int     					=			False
             _UpperCAmelCase			: List[Any]     					=			model_class(__UpperCamelCase							,					dtype=jnp.floataa)

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

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

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

             _UpperCAmelCase			: str     					=			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)
                 _UpperCAmelCase			: Optional[int]     					=			pt_model_class.from_pretrained(__UpperCamelCase							,					from_flax=__UpperCamelCase)

             with torch.no_grad():
                 _UpperCAmelCase			: Any     					=			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 	snake_case__			(		self)	->						Tuple:

     """simple docstring"""



     for model_class_name in self.all_model_classes:
         _UpperCAmelCase			: Tuple     					=			model_class_name.from_pretrained('''facebook/xglm-564M''')
         _UpperCAmelCase			: Union[str, Any]     					=			model(np.ones((1, 1)))
         self.assertIsNotNone(__UpperCamelCase)












from unittest.mock import patch

import datasets
from datasets import Dataset

def        _lowerCamelCase		(  ) ->   Tuple:
    _UpperCAmelCase			: Union[str, Any]     					=			Dataset.from_dict({'''col_1''': [3, 2, 0, 1]}	)

    with patch('''tqdm.auto.tqdm'''	) as mock_tqdm:
        datasets.disable_progress_bar()
        dset.map(lambda __A	: {"col_2": x["col_1"] + 1}	)
        mock_tqdm.assert_not_called()

        mock_tqdm.reset_mock()

        datasets.enable_progress_bar()
        dset.map(lambda __A	: {"col_2": x["col_1"] + 1}	)
        mock_tqdm.assert_called()




