




'''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       _a  :
       def __init__(     self						:     Optional[int]       ,    lowercase						:     Optional[int]       ,    lowercase						:     Any=14       ,    lowercase						:     Union[str, Any]=7       ,    lowercase						:     str=True       ,    lowercase						:     Tuple=True       ,    lowercase						:     str=True       ,    lowercase						:     str=99       ,    lowercase						:     Dict=32       ,    lowercase						:     Union[str, Any]=5       ,    lowercase						:     Optional[Any]=4       ,    lowercase						:     List[str]=37       ,    lowercase						:     List[Any]="gelu"       ,    lowercase						:     Union[str, Any]=0.1       ,    lowercase						:     int=0.1       ,    lowercase						:     Optional[int]=512       ,    lowercase						:     Optional[int]=0.02       ,    lowercase						:     List[str]=None       ,    ):



           '''simple docstring'''

           UpperCAmelCase      	=		parent
           UpperCAmelCase      	=		batch_size
           UpperCAmelCase      	=		seq_length
           UpperCAmelCase      	=		is_training
           UpperCAmelCase      	=		use_input_mask
           UpperCAmelCase      	=		use_labels
           UpperCAmelCase      	=		vocab_size
           UpperCAmelCase      	=		d_model
           UpperCAmelCase      	=		num_hidden_layers
           UpperCAmelCase      	=		num_attention_heads
           UpperCAmelCase      	=		ffn_dim
           UpperCAmelCase      	=		activation_function
           UpperCAmelCase      	=		activation_dropout
           UpperCAmelCase      	=		attention_dropout
           UpperCAmelCase      	=		max_position_embeddings
           UpperCAmelCase      	=		initializer_range
           UpperCAmelCase      	=		None
           UpperCAmelCase      	=		0
           UpperCAmelCase      	=		2
           UpperCAmelCase      	=		1
       def    A				(     self						:     str    ):



           '''simple docstring'''

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

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

           UpperCAmelCase      	=		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    A				(     self						:     Any    ):



           '''simple docstring'''

           UpperCAmelCase      	=		self.prepare_config_and_inputs()
           UpperCAmelCase      	=		config_and_inputs
           UpperCAmelCase      	=		{"""input_ids""": input_ids, """attention_mask""": attention_mask}
           return config, inputs_dict
       def    A				(     self						:     Optional[Any]    ):



           '''simple docstring'''

           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    A				(     self						:     List[Any]       ,    lowercase						:     Optional[Any]       ,    lowercase						:     Tuple       ,    lowercase						:     List[str]       ,    lowercase						:     Union[str, Any]    ):



           '''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    A				(     self						:     List[str]       ,    lowercase						:     Dict       ,    lowercase						:     Tuple       ,    lowercase						:     Optional[Any]       ,    lowercase						:     int    ):



           '''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_sentencepiece
@require_flax
class       _a  (     _a							,							_a							,							unittest.TestCase			):
       __a    :					Optional[Any]           = (FlaxXGLMModel, FlaxXGLMForCausalLM) if is_flax_available() else ()
       __a    :					Tuple           = (FlaxXGLMForCausalLM,) if is_flax_available() else ()
       def    A				(     self						:     List[Any]    ):



           '''simple docstring'''

           UpperCAmelCase      	=		FlaxXGLMModelTester(self    )
       def    A				(     self						:     List[str]    ):



           '''simple docstring'''

           for model_class_name in self.all_model_classes:
               UpperCAmelCase      	=		self.model_tester.prepare_config_and_inputs()
               self.model_tester.check_use_cache_forward(snake_case_       ,    snake_case_       ,    snake_case_       ,    snake_case_    )
       def    A				(     self						:     List[str]    ):



           '''simple docstring'''

           for model_class_name in self.all_model_classes:
               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    A				(     self						:     Dict    ):



           '''simple docstring'''

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

           UpperCAmelCase      	=		FlaxXGLMForCausalLM.from_pretrained('''facebook/xglm-564M'''    )
           UpperCAmelCase      	=		1
           UpperCAmelCase      	=		False

           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 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    A				(     self						:     List[Any]    ):



           '''simple docstring'''

           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_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()
                   # 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      	=		False
                   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    A				(     self						:     Dict    ):



           '''simple docstring'''

           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      	=		False
                   UpperCAmelCase      	=		model_class(snake_case_       ,    dtype=jnp.floataa    )

                   UpperCAmelCase      	=		load_flax_weights_in_pytorch_model(snake_case_       ,    fx_model.params    )
                   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    A				(     self						:     int    ):



           '''simple docstring'''

           for model_class_name in self.all_model_classes:
               UpperCAmelCase      	=		model_class_name.from_pretrained('''facebook/xglm-564M'''    )
               UpperCAmelCase      	=		model(np.ones((1, 1)    )    )
               self.assertIsNotNone(snake_case_    )










'''simple docstring'''







import unittest

from transformers import is_torch_available, load_tool

from .test_tools_common import ToolTesterMixin


if is_torch_available():
   import torch


class       _a  (     unittest.TestCase							,							__a			):
       def    A				(     self						:     str    ):



           '''simple docstring'''

           UpperCAmelCase      	=		load_tool('''speech-to-text'''    )
           self.tool.setup()
       def    A				(     self						:     Dict    ):



           '''simple docstring'''

           UpperCAmelCase      	=		self.tool(torch.ones(3_000    )    )
           self.assertEqual(lowercase       ,    ''' you'''    )







       def    A				(     self						:     Dict    ):



           '''simple docstring'''

           UpperCAmelCase      	=		self.tool(audio=torch.ones(3_000    )    )
           self.assertEqual(lowercase       ,    ''' you'''    )




