import unittest

import numpy as np

from transformers import DistilBertConfig, 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.numpy as jnp

							from transformers.models.distilbert.modeling_flax_distilbert import (
							    FlaxDistilBertForMaskedLM,
							    FlaxDistilBertForMultipleChoice,
							    FlaxDistilBertForQuestionAnswering,
							    FlaxDistilBertForSequenceClassification,
							    FlaxDistilBertForTokenClassification,
							    FlaxDistilBertModel,
							)


class   lowercase      (       unittest.TestCase       ):




							def __init__(  self				:   List[str]		,							_UpperCamelCase				:   Dict		,							_UpperCamelCase				:   Optional[Any]=13		,							_UpperCamelCase				:   List[str]=7		,							_UpperCamelCase				:   Tuple=True		,							_UpperCamelCase				:   Dict=True		,							_UpperCamelCase				:   Any=True		,							_UpperCamelCase				:   str=True		,							_UpperCamelCase				:   int=99		,							_UpperCamelCase				:   Dict=32		,							_UpperCamelCase				:   Dict=5		,							_UpperCamelCase				:   str=4		,							_UpperCamelCase				:   int=37		,							_UpperCamelCase				:   List[Any]="gelu"		,							_UpperCamelCase				:   str=0.1		,							_UpperCamelCase				:   Tuple=0.1		,							_UpperCamelCase				:   List[str]=512		,							_UpperCamelCase				:   List[str]=16		,							_UpperCamelCase				:   Union[str, Any]=2		,							_UpperCamelCase				:   int=0.0_2		,							_UpperCamelCase				:   int=4		,							)   ->			Optional[Any]:



												'''simple docstring'''

												SCREAMING_SNAKE_CASE           =    parent
												SCREAMING_SNAKE_CASE           =    batch_size
												SCREAMING_SNAKE_CASE           =    seq_length
												SCREAMING_SNAKE_CASE           =    is_training
												SCREAMING_SNAKE_CASE           =    use_attention_mask
												SCREAMING_SNAKE_CASE           =    use_token_type_ids
												SCREAMING_SNAKE_CASE           =    use_labels
												SCREAMING_SNAKE_CASE           =    vocab_size
												SCREAMING_SNAKE_CASE           =    hidden_size
												SCREAMING_SNAKE_CASE           =    num_hidden_layers
												SCREAMING_SNAKE_CASE           =    num_attention_heads
												SCREAMING_SNAKE_CASE           =    intermediate_size
												SCREAMING_SNAKE_CASE           =    hidden_act
												SCREAMING_SNAKE_CASE           =    hidden_dropout_prob
												SCREAMING_SNAKE_CASE           =    attention_probs_dropout_prob
												SCREAMING_SNAKE_CASE           =    max_position_embeddings
												SCREAMING_SNAKE_CASE           =    type_vocab_size
												SCREAMING_SNAKE_CASE           =    type_sequence_label_size
												SCREAMING_SNAKE_CASE           =    initializer_range
												SCREAMING_SNAKE_CASE           =    num_choices




							def 		__snake_case(  self				:   List[Any] )   ->			Tuple:



												'''simple docstring'''

												SCREAMING_SNAKE_CASE           =    ids_tensor([self.batch_size, self.seq_length]		,							self.vocab_size )

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

												SCREAMING_SNAKE_CASE           =    DistilBertConfig(
												    vocab_size=self.vocab_size		,							dim=self.hidden_size		,							n_layers=self.num_hidden_layers		,							n_heads=self.num_attention_heads		,							hidden_dim=self.intermediate_size		,							hidden_act=self.hidden_act		,							dropout=self.hidden_dropout_prob		,							attention_dropout=self.attention_probs_dropout_prob		,							max_position_embeddings=self.max_position_embeddings		,							initializer_range=self.initializer_range		,							tie_weights_=_UpperCamelCase		,							)

												return config, input_ids, attention_mask





							def 		__snake_case(  self				:   List[str] )   ->			Optional[Any]:



												'''simple docstring'''

												SCREAMING_SNAKE_CASE           =    self.prepare_config_and_inputs()
												SCREAMING_SNAKE_CASE     ,				SCREAMING_SNAKE_CASE     ,				SCREAMING_SNAKE_CASE           =    config_and_inputs
												SCREAMING_SNAKE_CASE           =    {"input_ids": input_ids, "attention_mask": attention_mask}
												return config, inputs_dict



@require_flax
class   lowercase      (       a			,    unittest.TestCase       ):
							lowercase__  :    Any									=							(
							    (
							        FlaxDistilBertModel,
							        FlaxDistilBertForMaskedLM,
							        FlaxDistilBertForMultipleChoice,
							        FlaxDistilBertForQuestionAnswering,
							        FlaxDistilBertForSequenceClassification,
							        FlaxDistilBertForTokenClassification,
							        FlaxDistilBertForQuestionAnswering,
							    )
							    if is_flax_available()
							    else ()
							)




							def 		__snake_case(  self				:   Optional[int] )   ->			Union[str, Any]:



												'''simple docstring'''

												SCREAMING_SNAKE_CASE           =    FlaxDistilBertModelTester(self )





							@slow
							def 		__snake_case(  self				:   List[str] )   ->			List[str]:



												'''simple docstring'''

												for model_class_name in self.all_model_classes:
																	SCREAMING_SNAKE_CASE           =    model_class_name.from_pretrained("distilbert-base-uncased" )
																	SCREAMING_SNAKE_CASE           =    model(np.ones((1, 1) ) )
																	self.assertIsNotNone(_UpperCamelCase )



@require_flax
class   lowercase      (       unittest.TestCase       ):




							@slow
							def 		__snake_case(  self				:   List[Any] )   ->			str:



												'''simple docstring'''

												SCREAMING_SNAKE_CASE           =    FlaxDistilBertModel.from_pretrained("distilbert-base-uncased" )
												SCREAMING_SNAKE_CASE           =    np.array([[0, 345, 232, 328, 740, 140, 1_695, 69, 6_078, 1_588, 2]] )
												SCREAMING_SNAKE_CASE           =    np.array([[0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]] )
												SCREAMING_SNAKE_CASE           =    model(_UpperCamelCase		,							attention_mask=_UpperCamelCase )[0]
												SCREAMING_SNAKE_CASE           =    (1, 11, 768)
												self.assertEqual(output.shape		,							_UpperCamelCase )
												SCREAMING_SNAKE_CASE           =    np.array([[[-0.1_6_3_9, 0.3_2_9_9, 0.1_6_4_8], [-0.1_7_4_6, 0.3_2_8_9, 0.1_7_1_0], [-0.1_8_8_4, 0.3_3_5_7, 0.1_8_1_0]]] )

												self.assertTrue(jnp.allclose(output[:, 1:4, 1:4]		,							_UpperCamelCase		,							atol=1e-4 ) )





import json
from typing import Dict, List, Optional, Tuple, Union

from tokenizers import pre_tokenizers, processors

from ...tokenization_utils_base import AddedToken, BatchEncoding, EncodedInput
from ...tokenization_utils_fast import PreTrainedTokenizerFast
from ...utils import PaddingStrategy, logging
from .tokenization_led import LEDTokenizer


_lowerCamelCase   :				Union[str, Any]     			=						logging.get_logger(__name__)


_lowerCamelCase   :				str     			=						{'''vocab_file''': '''vocab.json''', '''merges_file''': '''merges.txt''', '''tokenizer_file''': '''tokenizer.json'''}

_lowerCamelCase   :				int     			=						{
    '''vocab_file''': {
        '''allenai/led-base-16384''': '''https://huggingface.co/allenai/led-base-16384/resolve/main/vocab.json''',
    },
    '''merges_file''': {
        '''allenai/led-base-16384''': '''https://huggingface.co/allenai/led-base-16384/resolve/main/merges.txt''',
    },
    '''tokenizer_file''': {
        '''allenai/led-base-16384''': '''https://huggingface.co/allenai/led-base-16384/resolve/main/tokenizer.json''',
    },
}

_lowerCamelCase   :				List[Any]     			=						{
    '''allenai/led-base-16384''': 1_63_84,
}


class   lowercase      (       a       ):
							lowercase__  :    Tuple									=							VOCAB_FILES_NAMES
							lowercase__  :    Any									=							PRETRAINED_VOCAB_FILES_MAP
							lowercase__  :    Optional[Any]									=							PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
							lowercase__  :    Optional[Any]									=							LEDTokenizer
							lowercase__  :    str									=							["""input_ids""", """attention_mask"""]




							def __init__(  self				:   Union[str, Any]		,							_UpperCamelCase				:   Tuple=None		,							_UpperCamelCase				:   str=None		,							_UpperCamelCase				:   Tuple=None		,							_UpperCamelCase				:   List[str]="replace"		,							_UpperCamelCase				:   str="<s>"		,							_UpperCamelCase				:   List[Any]="</s>"		,							_UpperCamelCase				:   List[Any]="</s>"		,							_UpperCamelCase				:   List[str]="<s>"		,							_UpperCamelCase				:   Tuple="<unk>"		,							_UpperCamelCase				:   List[Any]="<pad>"		,							_UpperCamelCase				:   Tuple="<mask>"		,							_UpperCamelCase				:   List[str]=False		,							_UpperCamelCase				:   List[Any]=True		,							**_UpperCamelCase				:   Optional[Any]		,							)   ->			Tuple:



												'''simple docstring'''

												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		,							)

												SCREAMING_SNAKE_CASE           =    json.loads(self.backend_tokenizer.pre_tokenizer.__getstate__() )
												if pre_tok_state.get("add_prefix_space"		,							_UpperCamelCase ) != add_prefix_space:
																	SCREAMING_SNAKE_CASE           =    getattr(_UpperCamelCase		,							pre_tok_state.pop("type" ) )
																	SCREAMING_SNAKE_CASE           =    add_prefix_space
																	SCREAMING_SNAKE_CASE           =    pre_tok_class(**_UpperCamelCase )

												SCREAMING_SNAKE_CASE           =    add_prefix_space

												# the pre_tokenizer is already updated in the GPT2TokenizerFast `__init__`
												SCREAMING_SNAKE_CASE           =    "post_processor"
												SCREAMING_SNAKE_CASE           =    getattr(self.backend_tokenizer		,							_UpperCamelCase		,							_UpperCamelCase )
												if tokenizer_component_instance:
																	SCREAMING_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:
																						SCREAMING_SNAKE_CASE           =    tuple(state["sep"] )
																	if "cls" in state:
																						SCREAMING_SNAKE_CASE           =    tuple(state["cls"] )

																	SCREAMING_SNAKE_CASE           =    False

																	if state.get("add_prefix_space"		,							_UpperCamelCase ) != add_prefix_space:
																						SCREAMING_SNAKE_CASE           =    add_prefix_space
																						SCREAMING_SNAKE_CASE           =    True

																	if state.get("trim_offsets"		,							_UpperCamelCase ) != trim_offsets:
																						SCREAMING_SNAKE_CASE           =    trim_offsets
																						SCREAMING_SNAKE_CASE           =    True

																	if changes_to_apply:
																						SCREAMING_SNAKE_CASE           =    getattr(_UpperCamelCase		,							state.pop("type" ) )
																						SCREAMING_SNAKE_CASE           =    component_class(**_UpperCamelCase )
																						setattr(self.backend_tokenizer		,							_UpperCamelCase		,							_UpperCamelCase )




							@property
							# Copied from transformers.models.bart.tokenization_bart_fast.BartTokenizerFast.mask_token with BART->LED
							def 		__snake_case(  self				:   int )   ->			str:



												'''simple docstring'''

												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 		__snake_case(  self				:   Optional[int]		,							_UpperCamelCase				:   Optional[int] )   ->			Dict:



												'''simple docstring'''

												SCREAMING_SNAKE_CASE           =    AddedToken(_UpperCamelCase		,							lstrip=_UpperCamelCase		,							rstrip=_UpperCamelCase ) if isinstance(_UpperCamelCase		,							_UpperCamelCase ) else value
												SCREAMING_SNAKE_CASE           =    value




							def 		__snake_case(  self				:   List[Any]		,							*_UpperCamelCase				:   Optional[int]		,							**_UpperCamelCase				:   Union[str, Any] )   ->			BatchEncoding:



												'''simple docstring'''

												SCREAMING_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 		__snake_case(  self				:   int		,							*_UpperCamelCase				:   Dict		,							**_UpperCamelCase				:   Tuple )   ->			BatchEncoding:



												'''simple docstring'''

												SCREAMING_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 		__snake_case(  self				:   List[Any]		,							_UpperCamelCase				:   str		,							_UpperCamelCase				:   Optional[str] = None )   ->			Tuple[str]:



												'''simple docstring'''

												SCREAMING_SNAKE_CASE           =    self._tokenizer.model.save(_UpperCamelCase		,							name=_UpperCamelCase )
												return tuple(_UpperCamelCase )




							def 		__snake_case(  self				:   List[Any]		,							_UpperCamelCase				:   List[Any]		,							_UpperCamelCase				:   int=None )   ->			List[Any]:



												'''simple docstring'''

												SCREAMING_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 		__snake_case(  self				:   Dict		,							_UpperCamelCase				:   List[int]		,							_UpperCamelCase				:   Optional[List[int]] = None )   ->			List[int]:



												'''simple docstring'''

												SCREAMING_SNAKE_CASE           =    [self.sep_token_id]
												SCREAMING_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]





							def 		__snake_case(  self				:   Optional[Any]		,							_UpperCamelCase				:   Union[Dict[str, EncodedInput], BatchEncoding]		,							_UpperCamelCase				:   Optional[int] = None		,							_UpperCamelCase				:   PaddingStrategy = PaddingStrategy.DO_NOT_PAD		,							_UpperCamelCase				:   Optional[int] = None		,							_UpperCamelCase				:   Optional[bool] = None		,							)   ->			dict:



												'''simple docstring'''

												SCREAMING_SNAKE_CASE           =    super()._pad(
												    encoded_inputs=_UpperCamelCase		,							max_length=_UpperCamelCase		,							padding_strategy=_UpperCamelCase		,							pad_to_multiple_of=_UpperCamelCase		,							return_attention_mask=_UpperCamelCase		,							)

												# Load from model defaults
												if return_attention_mask is None:
																	SCREAMING_SNAKE_CASE           =    "attention_mask" in self.model_input_names

												if return_attention_mask and "global_attention_mask" in encoded_inputs:
																	SCREAMING_SNAKE_CASE           =    encoded_inputs[self.model_input_names[0]]
																	# `global_attention_mask` need to have the same length as other (sequential) inputs.
																	SCREAMING_SNAKE_CASE           =    len(encoded_inputs["global_attention_mask"] ) != len(_UpperCamelCase )

																	if needs_to_be_padded:
																						SCREAMING_SNAKE_CASE           =    len(_UpperCamelCase ) - len(encoded_inputs["global_attention_mask"] )

																						if self.padding_side == "right":
																											# Use `-1` since `0` in `global_attention_mask` means `local attention` instead of `not to attend`
																											SCREAMING_SNAKE_CASE           =    (
																											    encoded_inputs["global_attention_mask"] + [-1] * difference
																											)
																						elif self.padding_side == "left":
																											SCREAMING_SNAKE_CASE           =    [-1] * difference + encoded_inputs[
																											    "global_attention_mask"
																											]
																						else:
																											raise ValueError("Invalid padding strategy:" + str(self.padding_side ) )

												return encoded_inputs




