
"""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_     :



							"""simple docstring"""






							def __init__(  self  :List[str]					,							lowerCamelCase_  :Tuple					,							lowerCamelCase_  :Union[str, Any]=14					,							lowerCamelCase_  :int=7					,							lowerCamelCase_  :List[Any]=True					,							lowerCamelCase_  :Tuple=True					,							lowerCamelCase_  :Optional[int]=True					,							lowerCamelCase_  :Tuple=99					,							lowerCamelCase_  :List[str]=32					,							lowerCamelCase_  :Dict=5					,							lowerCamelCase_  :int=4					,							lowerCamelCase_  :Optional[Any]=37					,							lowerCamelCase_  :Dict="gelu"					,							lowerCamelCase_  :Tuple=0.1					,							lowerCamelCase_  :Optional[int]=0.1					,							lowerCamelCase_  :str=512					,							lowerCamelCase_  :Union[str, Any]=0.02					,							lowerCamelCase_  :int=None					,							):





										"""simple docstring"""

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





							def 				UpperCAmelCase__							(  self  :Dict	):





										"""simple docstring"""

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

										lowerCamelCase__     :					Tuple		=None
										if self.use_input_mask:
													lowerCamelCase__     :					List[str]		=random_attention_mask([self.batch_size, self.seq_length]	)

										lowerCamelCase__     :					Union[str, 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=lowerCamelCase_					,							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  :Tuple	):





										"""simple docstring"""

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





							def 				UpperCAmelCase__							(  self  :List[Any]	):





										"""simple docstring"""

										lowerCamelCase__     :					List[str]		=self.prepare_config_and_inputs()

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





										"""simple docstring"""

										lowerCamelCase__     :					Union[str, Any]		=20
										lowerCamelCase__     :					Union[str, Any]		=model_class_name(lowerCamelCase_	)

										lowerCamelCase__     :					Tuple		=model.init_cache(input_ids.shape[0]					,							lowerCamelCase_	)
										lowerCamelCase__     :					Dict		=jnp.ones((input_ids.shape[0], max_decoder_length)					,							dtype='i4'	)

										lowerCamelCase__     :					List[Any]		=jnp.broadcast_to(
										    jnp.arange(input_ids.shape[-1] - 1	)[None, :]					,							(input_ids.shape[0], input_ids.shape[-1] - 1)	)
										lowerCamelCase__     :					Optional[int]		=model(
										    input_ids[:, :-1]					,							attention_mask=lowerCamelCase_					,							past_key_values=lowerCamelCase_					,							position_ids=lowerCamelCase_					,							)

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

										lowerCamelCase__     :					Any		=model(lowerCamelCase_	)

										lowerCamelCase__     :					List[str]		=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  :List[str]					,							lowerCamelCase_  :Optional[int]					,							lowerCamelCase_  :List[Any]					,							lowerCamelCase_  :Tuple					,							lowerCamelCase_  :List[Any]	):





										"""simple docstring"""

										lowerCamelCase__     :					Union[str, Any]		=20
										lowerCamelCase__     :					Optional[Any]		=model_class_name(lowerCamelCase_	)

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

										lowerCamelCase__     :					Union[str, Any]		=model.init_cache(input_ids.shape[0]					,							lowerCamelCase_	)
										lowerCamelCase__     :					Optional[int]		=jnp.broadcast_to(
										    jnp.arange(input_ids.shape[-1] - 1	)[None, :]					,							(input_ids.shape[0], input_ids.shape[-1] - 1)	)

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

										lowerCamelCase__     :					Tuple		=model(lowerCamelCase_					,							attention_mask=lowerCamelCase_	)
										lowerCamelCase__     :					Union[str, 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}"""	)






@require_sentencepiece
@require_flax
class    A_     (					A__				, A__				, unittest.TestCase       ):



							"""simple docstring"""






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





							def 				UpperCAmelCase__							(  self  :List[Any]	):





										"""simple docstring"""

										lowerCamelCase__     :					List[Any]		=FlaxXGLMModelTester(self	)





							def 				UpperCAmelCase__							(  self  :int	):





										"""simple docstring"""

										for model_class_name in self.all_model_classes:
													lowerCamelCase__     :					Optional[Any]		=self.model_tester.prepare_config_and_inputs()
													self.model_tester.check_use_cache_forward(lowerCamelCase_					,							lowerCamelCase_					,							lowerCamelCase_					,							lowerCamelCase_	)





							def 				UpperCAmelCase__							(  self  :Any	):





										"""simple docstring"""

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





							@slow
							def 				UpperCAmelCase__							(  self  :Tuple	):





										"""simple docstring"""

										lowerCamelCase__     :					Optional[int]		=XGLMTokenizer.from_pretrained('XGLM'					,							padding_side='left'	)
										lowerCamelCase__     :					Dict		=tokenizer(['Hello this is a long string', 'Hey']					,							return_tensors='np'					,							padding=lowerCamelCase_					,							truncation=lowerCamelCase_	)

										lowerCamelCase__     :					Optional[int]		=FlaxXGLMForCausalLM.from_pretrained('facebook/xglm-564M'	)
										lowerCamelCase__     :					Optional[int]		=1
										lowerCamelCase__     :					Optional[Any]		=False

										lowerCamelCase__     :					Union[str, Any]		=jax.jit(model.generate	)

										lowerCamelCase__     :					Optional[Any]		=jit_generate(inputs['input_ids']					,							attention_mask=inputs['attention_mask']	).sequences

										lowerCamelCase__     :					Tuple		=tokenizer.batch_decode(lowerCamelCase_					,							skip_special_tokens=lowerCamelCase_	)

										lowerCamelCase__     :					Optional[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(lowerCamelCase_					,							lowerCamelCase_	)





							@is_pt_flax_cross_test
							def 				UpperCAmelCase__							(  self  :Tuple	):





										"""simple docstring"""

										lowerCamelCase__     :					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
																lowerCamelCase__     :					Tuple		=self._prepare_for_class(lowerCamelCase_					,							lowerCamelCase_	)
																lowerCamelCase__     :					List[Any]		={k: torch.tensor(v.tolist()	) for k, v in prepared_inputs_dict.items()}

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

																lowerCamelCase__     :					List[str]		=pt_inputs['input_ids'].shape
																lowerCamelCase__     :					List[Any]		=np.random.randint(0					,							seq_length - 1					,							size=(batch_size,)	)
																for batch_idx, start_index in enumerate(lowerCamelCase_	):
																			lowerCamelCase__     :					Union[str, Any]		=0
																			lowerCamelCase__     :					str		=1
																			lowerCamelCase__     :					Union[str, Any]		=0
																			lowerCamelCase__     :					int		=1
																lowerCamelCase__     :					str		=pt_model_class(lowerCamelCase_	).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__     :					Tuple		=False
																lowerCamelCase__     :					Optional[int]		=model_class(lowerCamelCase_					,							dtype=jnp.floataa	)
																lowerCamelCase__     :					Dict		=convert_pytorch_state_dict_to_flax(pt_model.state_dict()					,							lowerCamelCase_	)
																lowerCamelCase__     :					Dict		=fx_state

																with torch.no_grad():
																			lowerCamelCase__     :					Optional[Any]		=pt_model(**lowerCamelCase_	).to_tuple()

																lowerCamelCase__     :					str		=fx_model(**lowerCamelCase_	).to_tuple()
																self.assertEqual(len(lowerCamelCase_	)					,							len(lowerCamelCase_	)					,							'Output lengths differ between Flax and PyTorch'	)
																for fx_output, pt_output in zip(lowerCamelCase_					,							lowerCamelCase_	):
																			self.assert_almost_equals(fx_output[:, -1]					,							pt_output[:, -1].numpy()					,							4e-2	)

																with tempfile.TemporaryDirectory() as tmpdirname:
																			pt_model.save_pretrained(lowerCamelCase_	)
																			lowerCamelCase__     :					int		=model_class.from_pretrained(lowerCamelCase_					,							from_pt=lowerCamelCase_	)

																lowerCamelCase__     :					Optional[int]		=fx_model_loaded(**lowerCamelCase_	).to_tuple()
																self.assertEqual(
																    len(lowerCamelCase_	)					,							len(lowerCamelCase_	)					,							'Output lengths differ between Flax and PyTorch'	)
																for fx_output_loaded, pt_output in zip(lowerCamelCase_					,							lowerCamelCase_	):
																			self.assert_almost_equals(fx_output_loaded[:, -1]					,							pt_output[:, -1].numpy()					,							4e-2	)





							@is_pt_flax_cross_test
							def 				UpperCAmelCase__							(  self  :Optional[int]	):





										"""simple docstring"""

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

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

																lowerCamelCase__     :					Optional[int]		=pt_model_class(lowerCamelCase_	).eval()
																lowerCamelCase__     :					Any		=False
																lowerCamelCase__     :					List[Any]		=model_class(lowerCamelCase_					,							dtype=jnp.floataa	)

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

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

																with torch.no_grad():
																			lowerCamelCase__     :					Any		=pt_model(**lowerCamelCase_	).to_tuple()

																lowerCamelCase__     :					str		=fx_model(**lowerCamelCase_	).to_tuple()
																self.assertEqual(len(lowerCamelCase_	)					,							len(lowerCamelCase_	)					,							'Output lengths differ between Flax and PyTorch'	)
																for fx_output, pt_output in zip(lowerCamelCase_					,							lowerCamelCase_	):
																			self.assert_almost_equals(fx_output[:, -1]					,							pt_output[:, -1].numpy()					,							4e-2	)

																with tempfile.TemporaryDirectory() as tmpdirname:
																			fx_model.save_pretrained(lowerCamelCase_	)
																			lowerCamelCase__     :					List[str]		=pt_model_class.from_pretrained(lowerCamelCase_					,							from_flax=lowerCamelCase_	)

																with torch.no_grad():
																			lowerCamelCase__     :					str		=pt_model_loaded(**lowerCamelCase_	).to_tuple()

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





							@slow
							def 				UpperCAmelCase__							(  self  :List[Any]	):





										"""simple docstring"""

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

"""simple docstring"""


from typing import List, Optional, Tuple

from tokenizers import pre_tokenizers

from ...tokenization_utils_fast import PreTrainedTokenizerFast
from ...utils import logging
from .tokenization_clip import CLIPTokenizer


lowerCAmelCase								=						logging.get_logger(__name__)

lowerCAmelCase								=						{"""vocab_file""": """vocab.json""", """merges_file""": """merges.txt""", """tokenizer_file""": """tokenizer.json"""}

lowerCAmelCase								=						{
    """vocab_file""": {
        """openai/clip-vit-base-patch32""": """https://huggingface.co/openai/clip-vit-base-patch32/resolve/main/vocab.json""",
    },
    """merges_file""": {
        """openai/clip-vit-base-patch32""": """https://huggingface.co/openai/clip-vit-base-patch32/resolve/main/merges.txt""",
    },
    """tokenizer_file""": {
        """openai/clip-vit-base-patch32""": (
            """https://huggingface.co/openai/clip-vit-base-patch32/resolve/main/tokenizer.json"""
        ),
    },
}

lowerCAmelCase								=						{
    """openai/clip-vit-base-patch32""": 77,
}

class    A_     (					A__       ):



							"""simple docstring"""






							SCREAMING_SNAKE_CASE_     				=       VOCAB_FILES_NAMES
							SCREAMING_SNAKE_CASE_     				=       PRETRAINED_VOCAB_FILES_MAP
							SCREAMING_SNAKE_CASE_     				=       PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
							SCREAMING_SNAKE_CASE_     				=       ["""input_ids""", """attention_mask"""]
							SCREAMING_SNAKE_CASE_     				=       CLIPTokenizer





							def __init__(  self  :List[str]					,							lowerCamelCase_  :Tuple=None					,							lowerCamelCase_  :Any=None					,							lowerCamelCase_  :Optional[Any]=None					,							lowerCamelCase_  :List[Any]="<|endoftext|>"					,							lowerCamelCase_  :Tuple="<|startoftext|>"					,							lowerCamelCase_  :Dict="<|endoftext|>"					,							lowerCamelCase_  :str="<|endoftext|>"					,							**lowerCamelCase_  :int					,							):





										"""simple docstring"""

										super().__init__(
										    lowerCamelCase_					,							lowerCamelCase_					,							tokenizer_file=lowerCamelCase_					,							unk_token=lowerCamelCase_					,							bos_token=lowerCamelCase_					,							eos_token=lowerCamelCase_					,							pad_token=lowerCamelCase_					,							**lowerCamelCase_					,							)

										if not isinstance(self.backend_tokenizer.pre_tokenizer					,							pre_tokenizers.Sequence	):
													raise ValueError(
													    'The `backend_tokenizer` provided does not match the expected format. The CLIP tokenizer has been'
													    ' heavily modified from transformers version 4.17.0. You need to convert the tokenizer you are using'
													    ' to be compatible with this version.The easiest way to do so is'
													    ' `CLIPTokenizerFast.from_pretrained("path_to_local_folder_or_hub_repo, from_slow=True)`. If you want'
													    ' to use your existing tokenizer, you will have to revert to a version prior to 4.17.0 of'
													    ' transformers.'	)

										self._wrap_decode_method_backend_tokenizer()





							def 				UpperCAmelCase__							(  self  :Union[str, Any]	):





										"""simple docstring"""

										lowerCamelCase__     :					List[str]		=self.backend_tokenizer.decode

										def new_decode_method(*lowerCamelCase_  :Optional[int]					,							**lowerCamelCase_  :List[Any]	):
													lowerCamelCase__     :					Any		=orig_decode_method(*lowerCamelCase_					,							**lowerCamelCase_	)
													lowerCamelCase__     :					Optional[int]		=text.replace(self.backend_tokenizer.model.end_of_word_suffix					,							' '	).strip()
													return text

										lowerCamelCase__     :					Optional[Any]		=new_decode_method





							def 				UpperCAmelCase__							(  self  :List[str]					,							lowerCamelCase_  :List[int]					,							lowerCamelCase_  :Optional[List[int]] = None	):





										"""simple docstring"""

										lowerCamelCase__     :					Any		=[self.bos_token_id]
										lowerCamelCase__     :					Any		=[self.eos_token_id]

										if token_ids_a is None:
													return bos_token + token_ids_a + eos_token
										return bos_token + token_ids_a + eos_token + eos_token + token_ids_a + eos_token





							def 				UpperCAmelCase__							(  self  :List[Any]					,							lowerCamelCase_  :List[int]					,							lowerCamelCase_  :Optional[List[int]] = None	):





										"""simple docstring"""

										lowerCamelCase__     :					List[str]		=[self.bos_token_id]
										lowerCamelCase__     :					Optional[Any]		=[self.eos_token_id]

										if token_ids_a is None:
													return len(bos_token + token_ids_a + eos_token	) * [0]
										return len(bos_token + token_ids_a + eos_token + eos_token + token_ids_a + eos_token	) * [0]





							def 				UpperCAmelCase__							(  self  :Tuple					,							lowerCamelCase_  :str					,							lowerCamelCase_  :Optional[str] = None	):





										"""simple docstring"""

										lowerCamelCase__     :					int		=self._tokenizer.model.save(lowerCamelCase_					,							name=lowerCamelCase_	)
										return tuple(lowerCamelCase_	)