



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  __snake_case      (   unittest.TestCase):







							'''simple docstring'''






							def __init__(							self  ,  a_  ,  a_=13  ,  a_=7  ,  a_=True  ,  a_=True  ,  a_=True  ,  a_=True  ,  a_=99  ,  a_=32  ,  a_=5  ,  a_=4  ,  a_=37  ,  a_="gelu"  ,  a_=0.1  ,  a_=0.1  ,  a_=512  ,  a_=16  ,  a_=2  ,  a_=0.02  ,  a_=4  ,  ):
											a__				      =		parent
											a__				      =		batch_size
											a__				      =		seq_length
											a__				      =		is_training
											a__				      =		use_attention_mask
											a__				      =		use_token_type_ids
											a__				      =		use_labels
											a__				      =		vocab_size
											a__				      =		hidden_size
											a__				      =		num_hidden_layers
											a__				      =		num_attention_heads
											a__				      =		intermediate_size
											a__				      =		hidden_act
											a__				      =		hidden_dropout_prob
											a__				      =		attention_probs_dropout_prob
											a__				      =		max_position_embeddings
											a__				      =		type_vocab_size
											a__				      =		type_sequence_label_size
											a__				      =		initializer_range
											a__				      =		num_choices


							def _a      (							self		):
											a__				      =		ids_tensor([self.batch_size, self.seq_length]  ,  self.vocab_size		)

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

											a__				      =		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_=a_  ,  )

											return config, input_ids, attention_mask







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




@require_flax
class  __snake_case      (   SCREAMING_SNAKE_CASE			,unittest.TestCase):







							'''simple docstring'''






							UpperCamelCase__				:      Tuple         =					(
							    (
							        FlaxDistilBertModel,
							        FlaxDistilBertForMaskedLM,
							        FlaxDistilBertForMultipleChoice,
							        FlaxDistilBertForQuestionAnswering,
							        FlaxDistilBertForSequenceClassification,
							        FlaxDistilBertForTokenClassification,
							        FlaxDistilBertForQuestionAnswering,
							    )
							    if is_flax_available()
							    else ()
							)


							def _a      (							self		):
											a__				      =		FlaxDistilBertModelTester(self		)







							@slow
							def _a      (							self		):
											for model_class_name in self.all_model_classes:
															a__				      =		model_class_name.from_pretrained("""distilbert-base-uncased"""		)
															a__				      =		model(np.ones((1, 1)		)		)
															self.assertIsNotNone(a_		)




@require_flax
class  __snake_case      (   unittest.TestCase):







							'''simple docstring'''






							@slow
							def _a      (							self		):
											a__				      =		FlaxDistilBertModel.from_pretrained("""distilbert-base-uncased"""		)
											a__				      =		np.array([[0, 345, 232, 328, 740, 140, 1_695, 69, 6_078, 1_588, 2]]		)
											a__				      =		np.array([[0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]]		)
											a__				      =		model(a_  ,  attention_mask=a_		)[0]
											a__				      =		(1, 11, 768)
											self.assertEqual(output.shape  ,  a_		)
											a__				      =		np.array([[[-0.1_639, 0.3_299, 0.1_648], [-0.1_746, 0.3_289, 0.1_710], [-0.1_884, 0.3_357, 0.1_810]]]		)

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






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_albert import AlbertTokenizer
else:
	UpperCAmelCase							 =					None

UpperCAmelCase							 =					logging.get_logger(__name__)
UpperCAmelCase							 =					{"""vocab_file""": """spiece.model""", """tokenizer_file""": """tokenizer.json"""}

UpperCAmelCase							 =					{
    """vocab_file""": {
        """albert-base-v1""": """https://huggingface.co/albert-base-v1/resolve/main/spiece.model""",
        """albert-large-v1""": """https://huggingface.co/albert-large-v1/resolve/main/spiece.model""",
        """albert-xlarge-v1""": """https://huggingface.co/albert-xlarge-v1/resolve/main/spiece.model""",
        """albert-xxlarge-v1""": """https://huggingface.co/albert-xxlarge-v1/resolve/main/spiece.model""",
        """albert-base-v2""": """https://huggingface.co/albert-base-v2/resolve/main/spiece.model""",
        """albert-large-v2""": """https://huggingface.co/albert-large-v2/resolve/main/spiece.model""",
        """albert-xlarge-v2""": """https://huggingface.co/albert-xlarge-v2/resolve/main/spiece.model""",
        """albert-xxlarge-v2""": """https://huggingface.co/albert-xxlarge-v2/resolve/main/spiece.model""",
    },
    """tokenizer_file""": {
        """albert-base-v1""": """https://huggingface.co/albert-base-v1/resolve/main/tokenizer.json""",
        """albert-large-v1""": """https://huggingface.co/albert-large-v1/resolve/main/tokenizer.json""",
        """albert-xlarge-v1""": """https://huggingface.co/albert-xlarge-v1/resolve/main/tokenizer.json""",
        """albert-xxlarge-v1""": """https://huggingface.co/albert-xxlarge-v1/resolve/main/tokenizer.json""",
        """albert-base-v2""": """https://huggingface.co/albert-base-v2/resolve/main/tokenizer.json""",
        """albert-large-v2""": """https://huggingface.co/albert-large-v2/resolve/main/tokenizer.json""",
        """albert-xlarge-v2""": """https://huggingface.co/albert-xlarge-v2/resolve/main/tokenizer.json""",
        """albert-xxlarge-v2""": """https://huggingface.co/albert-xxlarge-v2/resolve/main/tokenizer.json""",
    },
}

UpperCAmelCase							 =					{
    """albert-base-v1""": 512,
    """albert-large-v1""": 512,
    """albert-xlarge-v1""": 512,
    """albert-xxlarge-v1""": 512,
    """albert-base-v2""": 512,
    """albert-large-v2""": 512,
    """albert-xlarge-v2""": 512,
    """albert-xxlarge-v2""": 512,
}

UpperCAmelCase							 =					"""▁"""

class  __snake_case      (   SCREAMING_SNAKE_CASE):







							'''simple docstring'''






							UpperCamelCase__				:      Tuple         =					VOCAB_FILES_NAMES
							UpperCamelCase__				:      Union[str, Any]         =					PRETRAINED_VOCAB_FILES_MAP
							UpperCamelCase__				:      List[Any]         =					PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
							UpperCamelCase__				:      List[str]         =					AlbertTokenizer


							def __init__(							self  ,  a_=None  ,  a_=None  ,  a_=True  ,  a_=True  ,  a_=False  ,  a_="[CLS]"  ,  a_="[SEP]"  ,  a_="<unk>"  ,  a_="[SEP]"  ,  a_="<pad>"  ,  a_="[CLS]"  ,  a_="[MASK]"  ,  **a_  ,  ):
											# Mask token behave like a normal word, i.e. include the space before it and
											# is included in the raw text, there should be a match in a non-normalized sentence.
											a__				      =		(
											    AddedToken(a_  ,  lstrip=a_  ,  rstrip=a_  ,  normalized=a_		)
											    if isinstance(a_  ,  a_		)
											    else mask_token
											)

											super().__init__(
											    a_  ,  tokenizer_file=a_  ,  do_lower_case=a_  ,  remove_space=a_  ,  keep_accents=a_  ,  bos_token=a_  ,  eos_token=a_  ,  unk_token=a_  ,  sep_token=a_  ,  pad_token=a_  ,  cls_token=a_  ,  mask_token=a_  ,  **a_  ,  )

											a__				      =		do_lower_case
											a__				      =		remove_space
											a__				      =		keep_accents
											a__				      =		vocab_file
											a__				      =		False if not self.vocab_file else True


							def _a      (							self  ,  a_  ,  a_ = None		):
											a__				      =		[self.sep_token_id]
											a__				      =		[self.cls_token_id]
											if token_ids_a is None:
															return cls + token_ids_a + sep
											return cls + token_ids_a + sep + token_ids_a + sep


							def _a      (							self  ,  a_  ,  a_ = None		):
											a__				      =		[self.sep_token_id]
											a__				      =		[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		) * [0] + len(token_ids_a + sep		) * [1]







							def _a      (							self  ,  a_  ,  a_ = 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(a_		):
															logger.error(F'''Vocabulary path ({save_directory}) should be a directory'''		)
															return
											a__				      =		os.path.join(
											    a_  ,  (filename_prefix + """-""" if filename_prefix else """""") + VOCAB_FILES_NAMES["""vocab_file"""]		)

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

											return (out_vocab_file,)

