



'''simple docstring'''

import unittest

from transformers import BigBirdConfig, is_flax_available
from transformers.testing_utils import require_flax, slow

from ...test_modeling_flax_common import FlaxModelTesterMixin, ids_tensor, random_attention_mask


if is_flax_available():
			import jax

			from transformers.models.big_bird.modeling_flax_big_bird import (
			    FlaxBigBirdForCausalLM,
			    FlaxBigBirdForMaskedLM,
			    FlaxBigBirdForMultipleChoice,
			    FlaxBigBirdForPreTraining,
			    FlaxBigBirdForQuestionAnswering,
			    FlaxBigBirdForSequenceClassification,
			    FlaxBigBirdForTokenClassification,
			    FlaxBigBirdModel,
			)

class 							a__( unittest.TestCase     ):






			def __init__(     self					,   _UpperCAmelCase					,   _UpperCAmelCase=2					,   _UpperCAmelCase=56					,   _UpperCAmelCase=True					,   _UpperCAmelCase=True					,   _UpperCAmelCase=True					,   _UpperCAmelCase=True					,   _UpperCAmelCase=99					,   _UpperCAmelCase=32					,   _UpperCAmelCase=2					,   _UpperCAmelCase=2					,   _UpperCAmelCase=7					,   _UpperCAmelCase="gelu_new"					,   _UpperCAmelCase=0.1					,   _UpperCAmelCase=0.1					,   _UpperCAmelCase=512					,   _UpperCAmelCase=16					,   _UpperCAmelCase=2					,   _UpperCAmelCase=0.02					,   _UpperCAmelCase=4					,   _UpperCAmelCase="block_sparse"					,   _UpperCAmelCase=True					,   _UpperCAmelCase=False					,   _UpperCAmelCase=2					,   _UpperCAmelCase=3					,   )							->							Optional[int]:
					snake_case__	     =parent
					snake_case__	     =batch_size
					snake_case__	     =seq_length
					snake_case__	     =is_training
					snake_case__	     =use_attention_mask
					snake_case__	     =use_token_type_ids
					snake_case__	     =use_labels
					snake_case__	     =vocab_size
					snake_case__	     =hidden_size
					snake_case__	     =num_hidden_layers
					snake_case__	     =num_attention_heads
					snake_case__	     =intermediate_size
					snake_case__	     =hidden_act
					snake_case__	     =hidden_dropout_prob
					snake_case__	     =attention_probs_dropout_prob
					snake_case__	     =max_position_embeddings
					snake_case__	     =type_vocab_size
					snake_case__	     =type_sequence_label_size
					snake_case__	     =initializer_range
					snake_case__	     =num_choices

					snake_case__	     =rescale_embeddings
					snake_case__	     =attention_type
					snake_case__	     =use_bias
					snake_case__	     =block_size
					snake_case__	     =num_random_blocks






			def 						_lowercase  (     self						)							->							List[Any]:
					snake_case__	     =ids_tensor([self.batch_size, self.seq_length]					,   self.vocab_size						)

					snake_case__	     =None
					if self.use_attention_mask:
							snake_case__	     =random_attention_mask([self.batch_size, self.seq_length]						)

					snake_case__	     =None
					if self.use_token_type_ids:
							snake_case__	     =ids_tensor([self.batch_size, self.seq_length]					,   self.type_vocab_size						)

					snake_case__	     =BigBirdConfig(
					    vocab_size=self.vocab_size					,   hidden_size=self.hidden_size					,   num_hidden_layers=self.num_hidden_layers					,   num_attention_heads=self.num_attention_heads					,   intermediate_size=self.intermediate_size					,   hidden_act=self.hidden_act					,   hidden_dropout_prob=self.hidden_dropout_prob					,   attention_probs_dropout_prob=self.attention_probs_dropout_prob					,   max_position_embeddings=self.max_position_embeddings					,   type_vocab_size=self.type_vocab_size					,   is_decoder=_UpperCAmelCase					,   initializer_range=self.initializer_range					,   attention_type=self.attention_type					,   block_size=self.block_size					,   num_random_blocks=self.num_random_blocks					,   use_bias=self.use_bias					,   rescale_embeddings=self.rescale_embeddings					,   )

					return config, input_ids, token_type_ids, attention_mask






			def 						_lowercase  (     self						)							->							Tuple:
					snake_case__	     =self.prepare_config_and_inputs()
					snake_case__     ,	snake_case__     ,	snake_case__     ,	snake_case__	     =config_and_inputs
					snake_case__	     ={
					    'input_ids': input_ids,
					    'token_type_ids': token_type_ids,
					    'attention_mask': attention_mask,
					}
					return config, inputs_dict





@require_flax
class 							a__( snake_case__   ,      unittest.TestCase     ):
			a_	: str									=      (
			    (
			        FlaxBigBirdForCausalLM,
			        FlaxBigBirdModel,
			        FlaxBigBirdForPreTraining,
			        FlaxBigBirdForMaskedLM,
			        FlaxBigBirdForMultipleChoice,
			        FlaxBigBirdForQuestionAnswering,
			        FlaxBigBirdForSequenceClassification,
			        FlaxBigBirdForTokenClassification,
			    )
			    if is_flax_available()
			    else ()
			)

			a_	: Tuple									=      False
			a_	: Tuple									=      False






			def 						_lowercase  (     self						)							->							Optional[int]:
					snake_case__	     =FlaxBigBirdModelTester(self						)






			@slow
			# copied from `test_modeling_flax_common` because it takes much longer than other models
			def 						_lowercase  (     self						)							->							str:
					super().test_from_pretrained_save_pretrained()






			@slow
			# copied from `test_modeling_flax_common` because it takes much longer than other models
			def 						_lowercase  (     self						)							->							Optional[Any]:
					super().test_from_pretrained_with_no_automatic_init()






			@slow
			# copied from `test_modeling_flax_common` because it takes much longer than other models
			def 						_lowercase  (     self						)							->							Union[str, Any]:
					super().test_no_automatic_init()






			@slow
			# copied from `test_modeling_flax_common` because it takes much longer than other models
			def 						_lowercase  (     self						)							->							int:
					super().test_hidden_states_output()






			@slow
			def 						_lowercase  (     self						)							->							Union[str, Any]:
					for model_class_name in self.all_model_classes:
							snake_case__	     =model_class_name.from_pretrained('google/bigbird-roberta-base'						)
							self.assertIsNotNone(_UpperCAmelCase						)






			def 						_lowercase  (     self						)							->							Tuple:
					if self.test_attn_probs:
							super().test_attention_outputs()






			@slow
			# copied from `test_modeling_flax_common` because it takes much longer than other models
			def 						_lowercase  (     self						)							->							str:
					snake_case__     ,	snake_case__	     =self.model_tester.prepare_config_and_inputs_for_common()

					for model_class in self.all_model_classes:
							with self.subTest(model_class.__name__						):
									snake_case__	     =self._prepare_for_class(_UpperCAmelCase					,   _UpperCAmelCase						)
									snake_case__	     =model_class(_UpperCAmelCase						)

									@jax.jit
									def model_jitted(_UpperCAmelCase					,   _UpperCAmelCase=None					,   **_UpperCAmelCase						):
											return model(input_ids=_UpperCAmelCase					,   attention_mask=_UpperCAmelCase					,   **_UpperCAmelCase						)

									with self.subTest('JIT Enabled'						):
											snake_case__	     =model_jitted(**_UpperCAmelCase						).to_tuple()

									with self.subTest('JIT Disabled'						):
											with jax.disable_jit():
													snake_case__	     =model_jitted(**_UpperCAmelCase						).to_tuple()

									self.assertEqual(len(_UpperCAmelCase						)					,   len(_UpperCAmelCase						)						)
									for jitted_output, output in zip(_UpperCAmelCase					,   _UpperCAmelCase						):
											self.assertEqual(jitted_output.shape					,   output.shape						)






			def 						_lowercase  (     self					,   _UpperCAmelCase					,   _UpperCAmelCase					,   _UpperCAmelCase					,   _UpperCAmelCase=1E-5					,   _UpperCAmelCase="outputs"					,   _UpperCAmelCase=None						)							->							Optional[Any]:
					# `bigbird_block_sparse_attention` in `FlaxBigBird` returns `attention_probs = None`, while in PyTorch version,
					# an effort was done to return `attention_probs` (yet to be verified).
					if name.startswith('outputs.attentions'						):
							return
					else:
							super().check_pt_flax_outputs(_UpperCAmelCase					,   _UpperCAmelCase					,   _UpperCAmelCase					,   _UpperCAmelCase					,   _UpperCAmelCase					,   _UpperCAmelCase						)









'''simple docstring'''

import json
from typing import List, Optional, Tuple

from tokenizers import pre_tokenizers, processors

from ...tokenization_utils_base import AddedToken, BatchEncoding
from ...tokenization_utils_fast import PreTrainedTokenizerFast
from ...utils import logging
from .tokenization_bart import BartTokenizer


SCREAMING_SNAKE_CASE__     :						List[str]	=  logging.get_logger(__name__)


SCREAMING_SNAKE_CASE__     :						Optional[Any]	=  {'''vocab_file''': '''vocab.json''', '''merges_file''': '''merges.txt''', '''tokenizer_file''': '''tokenizer.json'''}

# See all BART models at https://huggingface.co/models?filter=bart
SCREAMING_SNAKE_CASE__     :						Dict	=  {
    '''vocab_file''': {
        '''facebook/bart-base''': '''https://huggingface.co/facebook/bart-base/resolve/main/vocab.json''',
        '''facebook/bart-large''': '''https://huggingface.co/facebook/bart-large/resolve/main/vocab.json''',
        '''facebook/bart-large-mnli''': '''https://huggingface.co/facebook/bart-large-mnli/resolve/main/vocab.json''',
        '''facebook/bart-large-cnn''': '''https://huggingface.co/facebook/bart-large-cnn/resolve/main/vocab.json''',
        '''facebook/bart-large-xsum''': '''https://huggingface.co/facebook/bart-large-xsum/resolve/main/vocab.json''',
        '''yjernite/bart_eli5''': '''https://huggingface.co/yjernite/bart_eli5/resolve/main/vocab.json''',
    },
    '''merges_file''': {
        '''facebook/bart-base''': '''https://huggingface.co/facebook/bart-base/resolve/main/merges.txt''',
        '''facebook/bart-large''': '''https://huggingface.co/facebook/bart-large/resolve/main/merges.txt''',
        '''facebook/bart-large-mnli''': '''https://huggingface.co/facebook/bart-large-mnli/resolve/main/merges.txt''',
        '''facebook/bart-large-cnn''': '''https://huggingface.co/facebook/bart-large-cnn/resolve/main/merges.txt''',
        '''facebook/bart-large-xsum''': '''https://huggingface.co/facebook/bart-large-xsum/resolve/main/merges.txt''',
        '''yjernite/bart_eli5''': '''https://huggingface.co/yjernite/bart_eli5/resolve/main/merges.txt''',
    },
    '''tokenizer_file''': {
        '''facebook/bart-base''': '''https://huggingface.co/facebook/bart-base/resolve/main/tokenizer.json''',
        '''facebook/bart-large''': '''https://huggingface.co/facebook/bart-large/resolve/main/tokenizer.json''',
        '''facebook/bart-large-mnli''': '''https://huggingface.co/facebook/bart-large-mnli/resolve/main/tokenizer.json''',
        '''facebook/bart-large-cnn''': '''https://huggingface.co/facebook/bart-large-cnn/resolve/main/tokenizer.json''',
        '''facebook/bart-large-xsum''': '''https://huggingface.co/facebook/bart-large-xsum/resolve/main/tokenizer.json''',
        '''yjernite/bart_eli5''': '''https://huggingface.co/yjernite/bart_eli5/resolve/main/tokenizer.json''',
    },
}

SCREAMING_SNAKE_CASE__     :						Optional[int]	=  {
    '''facebook/bart-base''': 1_0_2_4,
    '''facebook/bart-large''': 1_0_2_4,
    '''facebook/bart-large-mnli''': 1_0_2_4,
    '''facebook/bart-large-cnn''': 1_0_2_4,
    '''facebook/bart-large-xsum''': 1_0_2_4,
    '''yjernite/bart_eli5''': 1_0_2_4,
}

class 							a__( snake_case__     ):
			a_	: Dict									=      VOCAB_FILES_NAMES
			a_	: Tuple									=      PRETRAINED_VOCAB_FILES_MAP
			a_	: Tuple									=      PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
			a_	: List[Any]									=      ['''input_ids''', '''attention_mask''']
			a_	: Union[str, Any]									=      BartTokenizer






			def __init__(     self					,   _UpperCAmelCase=None					,   _UpperCAmelCase=None					,   _UpperCAmelCase=None					,   _UpperCAmelCase="replace"					,   _UpperCAmelCase="<s>"					,   _UpperCAmelCase="</s>"					,   _UpperCAmelCase="</s>"					,   _UpperCAmelCase="<s>"					,   _UpperCAmelCase="<unk>"					,   _UpperCAmelCase="<pad>"					,   _UpperCAmelCase="<mask>"					,   _UpperCAmelCase=False					,   _UpperCAmelCase=True					,   **_UpperCAmelCase					,   )							->							int:
					super().__init__(
					    _UpperCAmelCase					,   _UpperCAmelCase					,   tokenizer_file=_UpperCAmelCase					,   errors=_UpperCAmelCase					,   bos_token=_UpperCAmelCase					,   eos_token=_UpperCAmelCase					,   sep_token=_UpperCAmelCase					,   cls_token=_UpperCAmelCase					,   unk_token=_UpperCAmelCase					,   pad_token=_UpperCAmelCase					,   mask_token=_UpperCAmelCase					,   add_prefix_space=_UpperCAmelCase					,   trim_offsets=_UpperCAmelCase					,   **_UpperCAmelCase					,   )

					snake_case__	     =json.loads(self.backend_tokenizer.pre_tokenizer.__getstate__()						)
					if pre_tok_state.get('add_prefix_space'					,   _UpperCAmelCase						) != add_prefix_space:
							snake_case__	     =getattr(_UpperCAmelCase					,   pre_tok_state.pop('type'						)						)
							snake_case__	     =add_prefix_space
							snake_case__	     =pre_tok_class(**_UpperCAmelCase						)

					snake_case__	     =add_prefix_space

					# the pre_tokenizer is already updated in the GPT2TokenizerFast `__init__`
					snake_case__	     ='post_processor'
					snake_case__	     =getattr(self.backend_tokenizer					,   _UpperCAmelCase					,   _UpperCAmelCase						)
					if tokenizer_component_instance:
							snake_case__	     =json.loads(tokenizer_component_instance.__getstate__()						)

							# The lists 'sep' and 'cls' must be cased in tuples for the object `post_processor_class`
							if "sep" in state:
									snake_case__	     =tuple(state['sep']						)
							if "cls" in state:
									snake_case__	     =tuple(state['cls']						)

							snake_case__	     =False

							if state.get('add_prefix_space'					,   _UpperCAmelCase						) != add_prefix_space:
									snake_case__	     =add_prefix_space
									snake_case__	     =True

							if state.get('trim_offsets'					,   _UpperCAmelCase						) != trim_offsets:
									snake_case__	     =trim_offsets
									snake_case__	     =True

							if changes_to_apply:
									snake_case__	     =getattr(_UpperCAmelCase					,   state.pop('type'						)						)
									snake_case__	     =component_class(**_UpperCAmelCase						)
									setattr(self.backend_tokenizer					,   _UpperCAmelCase					,   _UpperCAmelCase						)






			@property
			def 						_lowercase  (     self						)							->							str:
					if self._mask_token is None:
							if self.verbose:
									logger.error('Using mask_token, but it is not set yet.'						)
							return None
					return str(self._mask_token						)






			@mask_token.setter
			def 						_lowercase  (     self					,   _UpperCAmelCase						)							->							int:
					snake_case__	     =AddedToken(_UpperCAmelCase					,   lstrip=_UpperCAmelCase					,   rstrip=_UpperCAmelCase						) if isinstance(_UpperCAmelCase					,   _UpperCAmelCase						) else value
					snake_case__	     =value






			def 						_lowercase  (     self					,   *_UpperCAmelCase					,   **_UpperCAmelCase						)							->							BatchEncoding:
					snake_case__	     =kwargs.get('is_split_into_words'					,   _UpperCAmelCase						)

					if is_split_into_words and not self.add_prefix_space:
							raise ValueError(
							    f"""You need to instantiate {self.__class__.__name__} with add_prefix_space=True """
							    'to use it with pretokenized inputs.'						)

					return super()._batch_encode_plus(*_UpperCAmelCase					,   **_UpperCAmelCase						)






			def 						_lowercase  (     self					,   *_UpperCAmelCase					,   **_UpperCAmelCase						)							->							BatchEncoding:
					snake_case__	     =kwargs.get('is_split_into_words'					,   _UpperCAmelCase						)

					if is_split_into_words and not self.add_prefix_space:
							raise ValueError(
							    f"""You need to instantiate {self.__class__.__name__} with add_prefix_space=True """
							    'to use it with pretokenized inputs.'						)

					return super()._encode_plus(*_UpperCAmelCase					,   **_UpperCAmelCase						)






			def 						_lowercase  (     self					,   _UpperCAmelCase					,   _UpperCAmelCase = None						)							->							Tuple[str]:
					snake_case__	     =self._tokenizer.model.save(_UpperCAmelCase					,   name=_UpperCAmelCase						)
					return tuple(_UpperCAmelCase						)






			def 						_lowercase  (     self					,   _UpperCAmelCase					,   _UpperCAmelCase=None						)							->							Optional[int]:
					snake_case__	     =[self.bos_token_id] + token_ids_a + [self.eos_token_id]
					if token_ids_a is None:
							return output

					return output + [self.eos_token_id] + token_ids_a + [self.eos_token_id]






			def 						_lowercase  (     self					,   _UpperCAmelCase					,   _UpperCAmelCase = None						)							->							List[int]:
					snake_case__	     =[self.sep_token_id]
					snake_case__	     =[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]




