





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      ,       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.0_2      ,       UpperCAmelCase_=None      ,       ):
								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 _snake_case				(						self							):
								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 _snake_case				(						self							):
								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 _snake_case				(						self							):
								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 _snake_case				(						self      ,       UpperCAmelCase_      ,       UpperCAmelCase_      ,       UpperCAmelCase_      ,       UpperCAmelCase_							):
								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 _snake_case				(						self      ,       UpperCAmelCase_      ,       UpperCAmelCase_      ,       UpperCAmelCase_      ,       UpperCAmelCase_							):
								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 					__A   (  a			,    a			,    unittest.TestCase			):
				__A            = (FlaxXGLMModel, FlaxXGLMForCausalLM) if is_flax_available() else ()
				__A            = (FlaxXGLMForCausalLM,) if is_flax_available() else ()


				def _snake_case				(						self							):
								lowerCamelCase		       =FlaxXGLMModelTester(self							)


				def _snake_case				(						self							):
								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 _snake_case				(						self							):
								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 _snake_case				(						self							):
								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 _snake_case				(						self							):
								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 _snake_case				(						self							):
								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 _snake_case				(						self							):
								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_							)












import os
from shutil import copyfile
from typing import List, Optional, Tuple

from ...tokenization_utils import AddedToken
from ...tokenization_utils_fast import PreTrainedTokenizerFast
from ...utils import is_sentencepiece_available, logging


if is_sentencepiece_available():
							from .tokenization_xlm_roberta import XLMRobertaTokenizer
else:
							UpperCAmelCase__  :	Union[str, Any]     =None


UpperCAmelCase__  :	str     =logging.get_logger(__name__)

UpperCAmelCase__  :	Optional[int]     ={'''vocab_file''': '''sentencepiece.bpe.model''', '''tokenizer_file''': '''tokenizer.json'''}

UpperCAmelCase__  :	str     ={
    '''vocab_file''': {
        '''xlm-roberta-base''': '''https://huggingface.co/xlm-roberta-base/resolve/main/sentencepiece.bpe.model''',
        '''xlm-roberta-large''': '''https://huggingface.co/xlm-roberta-large/resolve/main/sentencepiece.bpe.model''',
        '''xlm-roberta-large-finetuned-conll02-dutch''': (
            '''https://huggingface.co/xlm-roberta-large-finetuned-conll02-dutch/resolve/main/sentencepiece.bpe.model'''
        ),
        '''xlm-roberta-large-finetuned-conll02-spanish''': (
            '''https://huggingface.co/xlm-roberta-large-finetuned-conll02-spanish/resolve/main/sentencepiece.bpe.model'''
        ),
        '''xlm-roberta-large-finetuned-conll03-english''': (
            '''https://huggingface.co/xlm-roberta-large-finetuned-conll03-english/resolve/main/sentencepiece.bpe.model'''
        ),
        '''xlm-roberta-large-finetuned-conll03-german''': (
            '''https://huggingface.co/xlm-roberta-large-finetuned-conll03-german/resolve/main/sentencepiece.bpe.model'''
        ),
    },
    '''tokenizer_file''': {
        '''xlm-roberta-base''': '''https://huggingface.co/xlm-roberta-base/resolve/main/tokenizer.json''',
        '''xlm-roberta-large''': '''https://huggingface.co/xlm-roberta-large/resolve/main/tokenizer.json''',
        '''xlm-roberta-large-finetuned-conll02-dutch''': (
            '''https://huggingface.co/xlm-roberta-large-finetuned-conll02-dutch/resolve/main/tokenizer.json'''
        ),
        '''xlm-roberta-large-finetuned-conll02-spanish''': (
            '''https://huggingface.co/xlm-roberta-large-finetuned-conll02-spanish/resolve/main/tokenizer.json'''
        ),
        '''xlm-roberta-large-finetuned-conll03-english''': (
            '''https://huggingface.co/xlm-roberta-large-finetuned-conll03-english/resolve/main/tokenizer.json'''
        ),
        '''xlm-roberta-large-finetuned-conll03-german''': (
            '''https://huggingface.co/xlm-roberta-large-finetuned-conll03-german/resolve/main/tokenizer.json'''
        ),
    },
}

UpperCAmelCase__  :	int     ={
    '''xlm-roberta-base''': 5_12,
    '''xlm-roberta-large''': 5_12,
    '''xlm-roberta-large-finetuned-conll02-dutch''': 5_12,
    '''xlm-roberta-large-finetuned-conll02-spanish''': 5_12,
    '''xlm-roberta-large-finetuned-conll03-english''': 5_12,
    '''xlm-roberta-large-finetuned-conll03-german''': 5_12,
}

class 					__A   (  a			):
				__A            = VOCAB_FILES_NAMES
				__A            = PRETRAINED_VOCAB_FILES_MAP
				__A            = PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
				__A            = ["""input_ids""", """attention_mask"""]
				__A            = XLMRobertaTokenizer


				def __init__(						self      ,       UpperCAmelCase_=None      ,       UpperCAmelCase_=None      ,       UpperCAmelCase_="<s>"      ,       UpperCAmelCase_="</s>"      ,       UpperCAmelCase_="</s>"      ,       UpperCAmelCase_="<s>"      ,       UpperCAmelCase_="<unk>"      ,       UpperCAmelCase_="<pad>"      ,       UpperCAmelCase_="<mask>"      ,       **UpperCAmelCase_      ,       ):
								# Mask token behave like a normal word, i.e. include the space before it
								lowerCamelCase		       =AddedToken(UpperCAmelCase_      ,       lstrip=UpperCAmelCase_      ,       rstrip=UpperCAmelCase_							) if isinstance(UpperCAmelCase_      ,       UpperCAmelCase_							) else mask_token

								super().__init__(
								    UpperCAmelCase_      ,       tokenizer_file=UpperCAmelCase_      ,       bos_token=UpperCAmelCase_      ,       eos_token=UpperCAmelCase_      ,       sep_token=UpperCAmelCase_      ,       cls_token=UpperCAmelCase_      ,       unk_token=UpperCAmelCase_      ,       pad_token=UpperCAmelCase_      ,       mask_token=UpperCAmelCase_      ,       **UpperCAmelCase_      ,       )

								lowerCamelCase		       =vocab_file
								lowerCamelCase		       =False if not self.vocab_file else True


				def _snake_case				(						self      ,       UpperCAmelCase_      ,       UpperCAmelCase_ = None							):

								if token_ids_a is None:
												return [self.cls_token_id] + token_ids_a + [self.sep_token_id]
								lowerCamelCase		       =[self.cls_token_id]
								lowerCamelCase		       =[self.sep_token_id]
								return cls + token_ids_a + sep + sep + token_ids_a + sep


				def _snake_case				(						self      ,       UpperCAmelCase_      ,       UpperCAmelCase_ = None							):
								lowerCamelCase		       =[self.sep_token_id]
								lowerCamelCase		       =[self.cls_token_id]

								if token_ids_a is None:
												return len(cls + token_ids_a + sep							) * [0]
								return len(cls + token_ids_a + sep + sep + token_ids_a + sep							) * [0]


				def _snake_case				(						self      ,       UpperCAmelCase_      ,       UpperCAmelCase_ = None							):
								if not self.can_save_slow_tokenizer:
												raise ValueError(
												    """Your fast tokenizer does not have the necessary information to save the vocabulary for a slow """
												    """tokenizer."""							)

								if not os.path.isdir(UpperCAmelCase_							):
												logger.error(f"""Vocabulary path ({save_directory}) should be a directory."""							)
												return
								lowerCamelCase		       =os.path.join(
								    UpperCAmelCase_      ,       (filename_prefix + """-""" if filename_prefix else """""") + VOCAB_FILES_NAMES["""vocab_file"""]							)

								if os.path.abspath(self.vocab_file							) != os.path.abspath(UpperCAmelCase_							):
												copyfile(self.vocab_file      ,       UpperCAmelCase_							)

								return (out_vocab_file,)





