



'''simple docstring'''




import os
import unicodedata
from shutil import copyfile
from typing import Any, Dict, List, Optional, Tuple

import sentencepiece as spm

from ...tokenization_utils import AddedToken, PreTrainedTokenizer
from ...utils import logging


lowerCamelCase							:int				    =					logging.get_logger(__name__)
lowerCamelCase							:int				    =					{'''vocab_file''': '''spiece.model'''}

lowerCamelCase							:List[str]				    =					{
    '''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''',
    }
}

lowerCamelCase							:Optional[int]				    =					{
    '''albert-base-v1''': 5_1_2,
    '''albert-large-v1''': 5_1_2,
    '''albert-xlarge-v1''': 5_1_2,
    '''albert-xxlarge-v1''': 5_1_2,
    '''albert-base-v2''': 5_1_2,
    '''albert-large-v2''': 5_1_2,
    '''albert-xlarge-v2''': 5_1_2,
    '''albert-xxlarge-v2''': 5_1_2,
}

lowerCamelCase							:Any				    =					'''▁'''







class 	_lowerCAmelCase							(			__UpperCAmelCase ):
	__SCREAMING_SNAKE_CASE   :		Optional[Any]					=					VOCAB_FILES_NAMES
	__SCREAMING_SNAKE_CASE   :		Optional[int]					=					PRETRAINED_VOCAB_FILES_MAP
	__SCREAMING_SNAKE_CASE   :		Optional[Any]					=					PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES


	def __init__(self       ,      lowercase       ,      lowercase=True       ,      lowercase=True       ,      lowercase=False       ,      lowercase="[CLS]"       ,      lowercase="[SEP]"       ,      lowercase="<unk>"       ,      lowercase="[SEP]"       ,      lowercase="<pad>"       ,      lowercase="[CLS]"       ,      lowercase="[MASK]"       ,      lowercase = None       ,      **lowercase       ,      ):
				# 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_	:						Tuple	=     (
				    AddedToken(lowercase       ,      lstrip=lowercase       ,      rstrip=lowercase       ,      normalized=lowercase							)
				    if isinstance(lowercase       ,      lowercase							)
				    else mask_token
				)

				A_	:						Union[str, Any]	=     {} if sp_model_kwargs is None else sp_model_kwargs

				super().__init__(
				    do_lower_case=lowercase       ,      remove_space=lowercase       ,      keep_accents=lowercase       ,      bos_token=lowercase       ,      eos_token=lowercase       ,      unk_token=lowercase       ,      sep_token=lowercase       ,      pad_token=lowercase       ,      cls_token=lowercase       ,      mask_token=lowercase       ,      sp_model_kwargs=self.sp_model_kwargs       ,      **lowercase       ,      )

				A_	:						Any	=     do_lower_case
				A_	:						List[str]	=     remove_space
				A_	:						Tuple	=     keep_accents
				A_	:						Dict	=     vocab_file

				A_	:						Any	=     spm.SentencePieceProcessor(**self.sp_model_kwargs							)
				self.sp_model.Load(lowercase							)


	@property
	def     _a	(self							):
				return len(self.sp_model							)


	def     _a	(self							):
				A_	:						List[str]	=     {self.convert_ids_to_tokens(lowercase							): i for i in range(self.vocab_size							)}
				vocab.update(self.added_tokens_encoder							)
				return vocab


	def __getstate__(self							):
				A_	:						Tuple	=     self.__dict__.copy()
				A_	:						Optional[Any]	=     None
				return state


	def __setstate__(self       ,      lowercase							):
				A_	:						Optional[int]	=     d

				# for backward compatibility
				if not hasattr(self       ,      """sp_model_kwargs"""							):
							A_	:						str	=     {}

				A_	:						Union[str, Any]	=     spm.SentencePieceProcessor(**self.sp_model_kwargs							)
				self.sp_model.Load(self.vocab_file							)


	def     _a	(self       ,      lowercase							):
				if self.remove_space:
							A_	:						Any	=     """ """.join(inputs.strip().split()							)
				else:
							A_	:						Tuple	=     inputs
				A_	:						str	=     outputs.replace("""``"""       ,      """\""""							).replace("""''"""       ,      """\""""							)

				if not self.keep_accents:
							A_	:						str	=     unicodedata.normalize("""NFKD"""       ,      lowercase							)
							A_	:						Dict	=     """""".join([c for c in outputs if not unicodedata.combining(lowercase							)]							)
				if self.do_lower_case:
							A_	:						Optional[Any]	=     outputs.lower()

				return outputs


	def     _a	(self       ,      lowercase							):
				A_	:						Optional[Any]	=     self.preprocess_text(lowercase							)
				A_	:						int	=     self.sp_model.encode(lowercase       ,      out_type=lowercase							)
				A_	:						List[Any]	=     []
				for piece in pieces:
							if len(lowercase							) > 1 and piece[-1] == str(""","""							) and piece[-2].isdigit():
										A_	:						Dict	=     self.sp_model.EncodeAsPieces(piece[:-1].replace(lowercase       ,      """"""							)							)
										if piece[0] != SPIECE_UNDERLINE and cur_pieces[0][0] == SPIECE_UNDERLINE:
													if len(cur_pieces[0]							) == 1:
																A_	:						str	=     cur_pieces[1:]
													else:
																A_	:						int	=     cur_pieces[0][1:]
										cur_pieces.append(piece[-1]							)
										new_pieces.extend(lowercase							)
							else:
										new_pieces.append(lowercase							)

				return new_pieces


	def     _a	(self       ,      lowercase							):
				return self.sp_model.PieceToId(lowercase							)


	def     _a	(self       ,      lowercase							):
				return self.sp_model.IdToPiece(lowercase							)


	def     _a	(self       ,      lowercase							):
				A_	:						Any	=     []
				A_	:						Optional[int]	=     """"""
				A_	:						Dict	=     False
				for token in tokens:
							# make sure that special tokens are not decoded using sentencepiece model
							if token in self.all_special_tokens:
										if not prev_is_special:
													out_string += " "
										out_string += self.sp_model.decode(lowercase							) + token
										A_	:						Optional[int]	=     True
										A_	:						int	=     []
							else:
										current_sub_tokens.append(lowercase							)
										A_	:						int	=     False
				out_string += self.sp_model.decode(lowercase							)
				return out_string.strip()


	def     _a	(self       ,      lowercase       ,      lowercase = None							):
				A_	:						List[Any]	=     [self.sep_token_id]
				A_	:						Optional[Any]	=     [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       ,      lowercase       ,      lowercase = None       ,      lowercase = False							):

				if already_has_special_tokens:
							return super().get_special_tokens_mask(
							    token_ids_a=lowercase       ,      token_ids_a=lowercase       ,      already_has_special_tokens=lowercase							)

				if token_ids_a is not None:
							return [1] + ([0] * len(lowercase							)) + [1] + ([0] * len(lowercase							)) + [1]
				return [1] + ([0] * len(lowercase							)) + [1]


	def     _a	(self       ,      lowercase       ,      lowercase = None							):
				A_	:						List[str]	=     [self.sep_token_id]
				A_	:						Optional[int]	=     [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       ,      lowercase       ,      lowercase = None							):
				if not os.path.isdir(lowercase							):
							logger.error(F'Vocabulary path ({save_directory}) should be a directory'							)
							return
				A_	:						List[Any]	=     os.path.join(
				    lowercase       ,      (filename_prefix + """-""" if filename_prefix else """""") + VOCAB_FILES_NAMES["""vocab_file"""]							)

				if os.path.abspath(self.vocab_file							) != os.path.abspath(lowercase							) and os.path.isfile(self.vocab_file							):
							copyfile(self.vocab_file       ,      lowercase							)
				elif not os.path.isfile(self.vocab_file							):
							with open(lowercase       ,      """wb"""							) as fi:
										A_	:						Any	=     self.sp_model.serialized_model_proto()
										fi.write(lowercase							)

				return (out_vocab_file,)




'''simple docstring'''




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


	def __init__(self       ,      lowercase       ,      lowercase=13       ,      lowercase=7       ,      lowercase=True       ,      lowercase=True       ,      lowercase=True       ,      lowercase=True       ,      lowercase=99       ,      lowercase=32       ,      lowercase=5       ,      lowercase=4       ,      lowercase=37       ,      lowercase="gelu"       ,      lowercase=0.1       ,      lowercase=0.1       ,      lowercase=512       ,      lowercase=16       ,      lowercase=2       ,      lowercase=0.02       ,      lowercase=4       ,      ):
				A_	:						List[Any]	=     parent
				A_	:						str	=     batch_size
				A_	:						List[Any]	=     seq_length
				A_	:						Dict	=     is_training
				A_	:						List[Any]	=     use_attention_mask
				A_	:						Any	=     use_token_type_ids
				A_	:						Optional[int]	=     use_labels
				A_	:						Tuple	=     vocab_size
				A_	:						List[str]	=     hidden_size
				A_	:						List[str]	=     num_hidden_layers
				A_	:						Optional[Any]	=     num_attention_heads
				A_	:						int	=     intermediate_size
				A_	:						Optional[Any]	=     hidden_act
				A_	:						List[Any]	=     hidden_dropout_prob
				A_	:						Optional[Any]	=     attention_probs_dropout_prob
				A_	:						Any	=     max_position_embeddings
				A_	:						Union[str, Any]	=     type_vocab_size
				A_	:						int	=     type_sequence_label_size
				A_	:						Any	=     initializer_range
				A_	:						List[str]	=     num_choices


	def     _a	(self							):
				A_	:						List[str]	=     ids_tensor([self.batch_size, self.seq_length]       ,      self.vocab_size							)

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

				A_	:						Union[str, Any]	=     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_=lowercase       ,      )

				return config, input_ids, attention_mask


	def     _a	(self							):
				A_	:						List[str]	=     self.prepare_config_and_inputs()
				A_,    A_,    A_	:						str	=     config_and_inputs
				A_	:						Any	=     {"""input_ids""": input_ids, """attention_mask""": attention_mask}
				return config, inputs_dict







@require_flax
class 	_lowerCAmelCase							(			__UpperCAmelCase      ,			unittest.TestCase ):
	__SCREAMING_SNAKE_CASE   :		Optional[Any]					=					(
	    (
	        FlaxDistilBertModel,
	        FlaxDistilBertForMaskedLM,
	        FlaxDistilBertForMultipleChoice,
	        FlaxDistilBertForQuestionAnswering,
	        FlaxDistilBertForSequenceClassification,
	        FlaxDistilBertForTokenClassification,
	        FlaxDistilBertForQuestionAnswering,
	    )
	    if is_flax_available()
	    else ()
	)


	def     _a	(self							):
				A_	:						Tuple	=     FlaxDistilBertModelTester(self							)


	@slow
	def     _a	(self							):
				for model_class_name in self.all_model_classes:
							A_	:						Union[str, Any]	=     model_class_name.from_pretrained("""distilbert-base-uncased"""							)
							A_	:						Any	=     model(np.ones((1, 1)							)							)
							self.assertIsNotNone(lowercase							)







@require_flax
class 	_lowerCAmelCase							(			unittest.TestCase ):


	@slow
	def     _a	(self							):
				A_	:						List[str]	=     FlaxDistilBertModel.from_pretrained("""distilbert-base-uncased"""							)
				A_	:						Optional[Any]	=     np.array([[0, 345, 232, 328, 740, 140, 1695, 69, 6078, 1588, 2]]							)
				A_	:						int	=     np.array([[0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]]							)
				A_	:						Optional[int]	=     model(lowercase       ,      attention_mask=lowercase							)[0]
				A_	:						Optional[Any]	=     (1, 11, 768)
				self.assertEqual(output.shape       ,      lowercase							)
				A_	:						Union[str, Any]	=     np.array([[[-0.16_39, 0.32_99, 0.16_48], [-0.17_46, 0.32_89, 0.17_10], [-0.18_84, 0.33_57, 0.18_10]]]							)

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