
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:
						__a         =   None

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

__a         =   {
    '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',
    },
}

__a         =   {
    '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,
}

__a         =   '▁'
class 					lowercase__(							UpperCAmelCase				):



	"""simple docstring"""




	a     :Optional[Any]       	=      VOCAB_FILES_NAMES
	a     :Tuple       	=      PRETRAINED_VOCAB_FILES_MAP
	a     :Dict       	=      PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
	a     :Union[str, Any]       	=      AlbertTokenizer
	def __init__(			self    :     Optional[int]			,				SCREAMING_SNAKE_CASE_    :     Any=None			,				SCREAMING_SNAKE_CASE_    :     List[Any]=None			,				SCREAMING_SNAKE_CASE_    :     List[Any]=True			,				SCREAMING_SNAKE_CASE_    :     Any=True			,				SCREAMING_SNAKE_CASE_    :     Union[str, Any]=False			,				SCREAMING_SNAKE_CASE_    :     List[str]="[CLS]"			,				SCREAMING_SNAKE_CASE_    :     Any="[SEP]"			,				SCREAMING_SNAKE_CASE_    :     str="<unk>"			,				SCREAMING_SNAKE_CASE_    :     Optional[int]="[SEP]"			,				SCREAMING_SNAKE_CASE_    :     str="<pad>"			,				SCREAMING_SNAKE_CASE_    :     Optional[int]="[CLS]"			,				SCREAMING_SNAKE_CASE_    :     Dict="[MASK]"			,				**SCREAMING_SNAKE_CASE_    :     int			,				)			->	Tuple:
					# 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.
					lowercase_	    =   (
					    AddedToken(SCREAMING_SNAKE_CASE_			,				lstrip=SCREAMING_SNAKE_CASE_			,				rstrip=SCREAMING_SNAKE_CASE_			,				normalized=SCREAMING_SNAKE_CASE_			)
					    if isinstance(SCREAMING_SNAKE_CASE_			,				SCREAMING_SNAKE_CASE_			)
					    else mask_token
					)

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

					lowercase_	    =   do_lower_case
					lowercase_	    =   remove_space
					lowercase_	    =   keep_accents
					lowercase_	    =   vocab_file
					lowercase_	    =   False if not self.vocab_file else True
	def 		_lowercase    (			self    :     Any			,				SCREAMING_SNAKE_CASE_    :     List[int]			,				SCREAMING_SNAKE_CASE_    :     Optional[List[int]] = None			)			->	List[int]:
					lowercase_	    =   [self.sep_token_id]
					lowercase_	    =   [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 		_lowercase    (			self    :     Optional[Any]			,				SCREAMING_SNAKE_CASE_    :     List[int]			,				SCREAMING_SNAKE_CASE_    :     Optional[List[int]] = None			)			->	List[int]:
					lowercase_	    =   [self.sep_token_id]
					lowercase_	    =   [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 		_lowercase    (			self    :     Union[str, Any]			,				SCREAMING_SNAKE_CASE_    :     str			,				SCREAMING_SNAKE_CASE_    :     Optional[str] = None			)			->	Tuple[str]:
					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(SCREAMING_SNAKE_CASE_			):
									logger.error(f'''Vocabulary path ({save_directory}) should be a directory'''			)
									return
					lowercase_	    =   os.path.join(
					    SCREAMING_SNAKE_CASE_			,				(filename_prefix + '''-''' if filename_prefix else '''''') + VOCAB_FILES_NAMES['''vocab_file''']			)

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

					return (out_vocab_file,)








import os
import shutil
import tempfile
from unittest import TestCase
from unittest.mock import patch

import numpy as np
from datasets import Dataset

from transformers.models.realm.configuration_realm import RealmConfig
from transformers.models.realm.retrieval_realm import _REALM_BLOCK_RECORDS_FILENAME, RealmRetriever
from transformers.models.realm.tokenization_realm import VOCAB_FILES_NAMES, RealmTokenizer







class 	__lowercase    (__SCREAMING_SNAKE_CASE  ):





							"""simple docstring"""







							def 					UpperCamelCase__     (      self       ):





										"""simple docstring"""




										SCREAMING_SNAKE_CASE_		:     Tuple					=					tempfile.mkdtemp()
										SCREAMING_SNAKE_CASE_		:     Dict					=					5

										# Realm tok
										SCREAMING_SNAKE_CASE_		:     Dict					=					[
										    '[UNK]',
										    '[CLS]',
										    '[SEP]',
										    '[PAD]',
										    '[MASK]',
										    'test',
										    'question',
										    'this',
										    'is',
										    'the',
										    'first',
										    'second',
										    'third',
										    'fourth',
										    'fifth',
										    'record',
										    'want',
										    '##want',
										    '##ed',
										    'wa',
										    'un',
										    'runn',
										    '##ing',
										    ',',
										    'low',
										    'lowest',
										]
										SCREAMING_SNAKE_CASE_		:     Tuple					=					os.path.join(self.tmpdirname			,						'realm_tokenizer'       )
										os.makedirs(lowerCAmelCase__			,						exist_ok=lowerCAmelCase__       )
										SCREAMING_SNAKE_CASE_		:     Tuple					=					os.path.join(lowerCAmelCase__			,						VOCAB_FILES_NAMES['vocab_file']       )
										with open(self.vocab_file			,						'w'			,						encoding='utf-8'       ) as vocab_writer:
													vocab_writer.write(''.join([x + '\n' for x in vocab_tokens]       )       )

										SCREAMING_SNAKE_CASE_		:     str					=					os.path.join(self.tmpdirname			,						'realm_block_records'       )
										os.makedirs(lowerCAmelCase__			,						exist_ok=lowerCAmelCase__       )






							def 					UpperCamelCase__     (      self       ):





										"""simple docstring"""




										return RealmTokenizer.from_pretrained(os.path.join(self.tmpdirname			,						'realm_tokenizer'       )       )






							def 					UpperCamelCase__     (      self       ):





										"""simple docstring"""




										shutil.rmtree(self.tmpdirname       )






							def 					UpperCamelCase__     (      self       ):





										"""simple docstring"""




										SCREAMING_SNAKE_CASE_		:     Union[str, Any]					=					RealmConfig(num_block_records=self.num_block_records       )
										return config






							def 					UpperCamelCase__     (      self       ):





										"""simple docstring"""




										SCREAMING_SNAKE_CASE_		:     Optional[int]					=					Dataset.from_dict(
										    {
										        'id': ['0', '1'],
										        'question': ['foo', 'bar'],
										        'answers': [['Foo', 'Bar'], ['Bar']],
										    }       )
										return dataset






							def 					UpperCamelCase__     (      self       ):





										"""simple docstring"""




										SCREAMING_SNAKE_CASE_		:     Dict					=					np.array(
										    [
										        b'This is the first record',
										        b'This is the second record',
										        b'This is the third record',
										        b'This is the fourth record',
										        b'This is the fifth record',
										        b'This is a longer longer longer record',
										    ]			,						dtype=lowerCAmelCase__			,						)
										return block_records






							def 					UpperCamelCase__     (      self       ):





										"""simple docstring"""




										SCREAMING_SNAKE_CASE_		:     Any					=					RealmRetriever(
										    block_records=self.get_dummy_block_records()			,						tokenizer=self.get_tokenizer()			,						)
										return retriever






							def 					UpperCamelCase__     (      self       ):





										"""simple docstring"""




										SCREAMING_SNAKE_CASE_		:     Any					=					self.get_config()
										SCREAMING_SNAKE_CASE_		:     Optional[Any]					=					self.get_dummy_retriever()
										SCREAMING_SNAKE_CASE_		:     List[str]					=					retriever.tokenizer

										SCREAMING_SNAKE_CASE_		:     Any					=					np.array([0, 3]			,						dtype='long'       )
										SCREAMING_SNAKE_CASE_		:     Union[str, Any]					=					tokenizer(['Test question']       ).input_ids
										SCREAMING_SNAKE_CASE_		:     int					=					tokenizer(
										    ['the fourth']			,						add_special_tokens=lowerCAmelCase__			,						return_token_type_ids=lowerCAmelCase__			,						return_attention_mask=lowerCAmelCase__			,						).input_ids
										SCREAMING_SNAKE_CASE_		:     Any					=					config.reader_seq_len

										SCREAMING_SNAKE_CASE_		,      SCREAMING_SNAKE_CASE_		,      SCREAMING_SNAKE_CASE_		,      SCREAMING_SNAKE_CASE_		:     str					=					retriever(
										    lowerCAmelCase__			,						lowerCAmelCase__			,						answer_ids=lowerCAmelCase__			,						max_length=lowerCAmelCase__			,						return_tensors='np'       )

										self.assertEqual(len(lowerCAmelCase__       )			,						2       )
										self.assertEqual(len(lowerCAmelCase__       )			,						2       )
										self.assertEqual(len(lowerCAmelCase__       )			,						2       )
										self.assertEqual(concat_inputs.input_ids.shape			,						(2, 1_0)       )
										self.assertEqual(concat_inputs.attention_mask.shape			,						(2, 1_0)       )
										self.assertEqual(concat_inputs.token_type_ids.shape			,						(2, 1_0)       )
										self.assertEqual(concat_inputs.special_tokens_mask.shape			,						(2, 1_0)       )
										self.assertEqual(
										    tokenizer.convert_ids_to_tokens(concat_inputs.input_ids[0]       )			,						['[CLS]', 'test', 'question', '[SEP]', 'this', 'is', 'the', 'first', 'record', '[SEP]']			,						)
										self.assertEqual(
										    tokenizer.convert_ids_to_tokens(concat_inputs.input_ids[1]       )			,						['[CLS]', 'test', 'question', '[SEP]', 'this', 'is', 'the', 'fourth', 'record', '[SEP]']			,						)






							def 					UpperCamelCase__     (      self       ):





										"""simple docstring"""




										SCREAMING_SNAKE_CASE_		:     Dict					=					self.get_config()
										SCREAMING_SNAKE_CASE_		:     Union[str, Any]					=					self.get_dummy_retriever()
										SCREAMING_SNAKE_CASE_		:     Tuple					=					retriever.tokenizer

										SCREAMING_SNAKE_CASE_		:     Optional[Any]					=					np.array([0, 3, 5]			,						dtype='long'       )
										SCREAMING_SNAKE_CASE_		:     List[Any]					=					tokenizer(['Test question']       ).input_ids
										SCREAMING_SNAKE_CASE_		:     Optional[int]					=					tokenizer(
										    ['the fourth', 'longer longer']			,						add_special_tokens=lowerCAmelCase__			,						return_token_type_ids=lowerCAmelCase__			,						return_attention_mask=lowerCAmelCase__			,						).input_ids
										SCREAMING_SNAKE_CASE_		:     Union[str, Any]					=					config.reader_seq_len

										SCREAMING_SNAKE_CASE_		,      SCREAMING_SNAKE_CASE_		,      SCREAMING_SNAKE_CASE_		,      SCREAMING_SNAKE_CASE_		:     Tuple					=					retriever(
										    lowerCAmelCase__			,						lowerCAmelCase__			,						answer_ids=lowerCAmelCase__			,						max_length=lowerCAmelCase__			,						return_tensors='np'       )

										self.assertEqual([False, True, True]			,						lowerCAmelCase__       )
										self.assertEqual([[-1, -1, -1], [6, -1, -1], [6, 7, 8]]			,						lowerCAmelCase__       )
										self.assertEqual([[-1, -1, -1], [7, -1, -1], [7, 8, 9]]			,						lowerCAmelCase__       )






							def 					UpperCamelCase__     (      self       ):





										"""simple docstring"""




										SCREAMING_SNAKE_CASE_		:     Dict					=					self.get_dummy_retriever()
										retriever.save_pretrained(os.path.join(self.tmpdirname			,						'realm_block_records'       )       )

										# Test local path
										SCREAMING_SNAKE_CASE_		:     Optional[Any]					=					retriever.from_pretrained(os.path.join(self.tmpdirname			,						'realm_block_records'       )       )
										self.assertEqual(retriever.block_records[0]			,						b'This is the first record'       )

										# Test mocked remote path
										with patch('transformers.models.realm.retrieval_realm.hf_hub_download'       ) as mock_hf_hub_download:
													SCREAMING_SNAKE_CASE_		:     Optional[int]					=					os.path.join(
													    os.path.join(self.tmpdirname			,						'realm_block_records'       )			,						_REALM_BLOCK_RECORDS_FILENAME       )
													SCREAMING_SNAKE_CASE_		:     Optional[Any]					=					RealmRetriever.from_pretrained('google/realm-cc-news-pretrained-openqa'       )

										self.assertEqual(retriever.block_records[0]			,						b'This is the first record'       )





