



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    lowerCamelCase__  (						UpperCAmelCase_):
     """simple docstring"""



     def 						_a      (self    ):






           '''simple docstring'''



           lowerCamelCase								=		tempfile.mkdtemp()
           lowerCamelCase								=		5

           # Realm tok
           lowerCamelCase								=		[
               "[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",
           ]
           lowerCamelCase								=		os.path.join(self.tmpdirname				,			"realm_tokenizer"    )
           os.makedirs(__a				,			exist_ok=__a    )
           lowerCamelCase								=		os.path.join(__a				,			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]    )    )

           lowerCamelCase								=		os.path.join(self.tmpdirname				,			"realm_block_records"    )
           os.makedirs(__a				,			exist_ok=__a    )







     def 						_a      (self    ):






           '''simple docstring'''



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







     def 						_a      (self    ):






           '''simple docstring'''



           shutil.rmtree(self.tmpdirname    )







     def 						_a      (self    ):






           '''simple docstring'''



           lowerCamelCase								=		RealmConfig(num_block_records=self.num_block_records    )
           return config







     def 						_a      (self    ):






           '''simple docstring'''



           lowerCamelCase								=		Dataset.from_dict(
               {
                   "id": ["0", "1"],
                   "question": ["foo", "bar"],
                   "answers": [["Foo", "Bar"], ["Bar"]],
               }    )
           return dataset







     def 						_a      (self    ):






           '''simple docstring'''



           lowerCamelCase								=		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=__a				,			)
           return block_records







     def 						_a      (self    ):






           '''simple docstring'''



           lowerCamelCase								=		RealmRetriever(
               block_records=self.get_dummy_block_records()				,			tokenizer=self.get_tokenizer()				,			)
           return retriever







     def 						_a      (self    ):






           '''simple docstring'''



           lowerCamelCase								=		self.get_config()
           lowerCamelCase								=		self.get_dummy_retriever()
           lowerCamelCase								=		retriever.tokenizer

           lowerCamelCase								=		np.array([0, 3]				,			dtype="long"    )
           lowerCamelCase								=		tokenizer(["Test question"]    ).input_ids
           lowerCamelCase								=		tokenizer(
               ["the fourth"]				,			add_special_tokens=__a				,			return_token_type_ids=__a				,			return_attention_mask=__a				,			).input_ids
           lowerCamelCase								=		config.reader_seq_len

           lowerCamelCase			,							lowerCamelCase			,							lowerCamelCase			,							lowerCamelCase								=		retriever(
               __a				,			__a				,			answer_ids=__a				,			max_length=__a				,			return_tensors="np"    )

           self.assertEqual(len(__a    )				,			2    )
           self.assertEqual(len(__a    )				,			2    )
           self.assertEqual(len(__a    )				,			2    )
           self.assertEqual(concat_inputs.input_ids.shape				,			(2, 10)    )
           self.assertEqual(concat_inputs.attention_mask.shape				,			(2, 10)    )
           self.assertEqual(concat_inputs.token_type_ids.shape				,			(2, 10)    )
           self.assertEqual(concat_inputs.special_tokens_mask.shape				,			(2, 10)    )
           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 						_a      (self    ):






           '''simple docstring'''



           lowerCamelCase								=		self.get_config()
           lowerCamelCase								=		self.get_dummy_retriever()
           lowerCamelCase								=		retriever.tokenizer

           lowerCamelCase								=		np.array([0, 3, 5]				,			dtype="long"    )
           lowerCamelCase								=		tokenizer(["Test question"]    ).input_ids
           lowerCamelCase								=		tokenizer(
               ["the fourth", "longer longer"]				,			add_special_tokens=__a				,			return_token_type_ids=__a				,			return_attention_mask=__a				,			).input_ids
           lowerCamelCase								=		config.reader_seq_len

           lowerCamelCase			,							lowerCamelCase			,							lowerCamelCase			,							lowerCamelCase								=		retriever(
               __a				,			__a				,			answer_ids=__a				,			max_length=__a				,			return_tensors="np"    )

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







     def 						_a      (self    ):






           '''simple docstring'''



           lowerCamelCase								=		self.get_dummy_retriever()
           retriever.save_pretrained(os.path.join(self.tmpdirname				,			"realm_block_records"    )    )

           # Test local path
           lowerCamelCase								=		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:
                 lowerCamelCase								=		os.path.join(
                     os.path.join(self.tmpdirname				,			"realm_block_records"    )				,			_REALM_BLOCK_RECORDS_FILENAME    )
                 lowerCamelCase								=		RealmRetriever.from_pretrained("google/realm-cc-news-pretrained-openqa"    )

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




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_rembert import RemBertTokenizer
else:
     a_						:     Any	     =		None

a_						:     Optional[Any]	     =		logging.get_logger(__name__)
a_						:     int	     =		{'vocab_file': 'sentencepiece.model', 'tokenizer_file': 'tokenizer.json'}

a_						:     Dict	     =		{
    'vocab_file': {
        'google/rembert': 'https://huggingface.co/google/rembert/resolve/main/sentencepiece.model',
    },
    'tokenizer_file': {
        'google/rembert': 'https://huggingface.co/google/rembert/resolve/main/tokenizer.json',
    },
}

a_						:     List[str]	     =		{
    'google/rembert': 2_5_6,
}

a_						:     str	     =		'▁'





class    lowerCamelCase__  (						UpperCAmelCase_):
     """simple docstring"""



     _A       =			VOCAB_FILES_NAMES
     _A       =			PRETRAINED_VOCAB_FILES_MAP
     _A       =			PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
     _A       =			RemBertTokenizer







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






           '''simple docstring'''



           lowerCamelCase								=		AddedToken(__a				,			lstrip=__a				,			rstrip=__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				,			)

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







     def 						_a      (self				,			__a				,			__a = None    ):






           '''simple docstring'''



           lowerCamelCase								=		[self.sep_token_id]
           lowerCamelCase								=		[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 = False    ):






           '''simple docstring'''



           if already_has_special_tokens:
                 if token_ids_a is not None:
                       raise ValueError(
                           "You should not supply a second sequence if the provided sequence of "
                           "ids is already formatted with special tokens for the model."    )
                 return [1 if x in [self.sep_token_id, self.cls_token_id] else 0 for x in token_ids_a]

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







     def 						_a      (self				,			__a				,			__a = None    ):






           '''simple docstring'''



           lowerCamelCase								=		[self.sep_token_id]
           lowerCamelCase								=		[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    ):






           '''simple docstring'''



           if not os.path.isdir(__a    ):
                 logger.error("Vocabulary path ({}) should be a directory".format(__a    )    )
                 return
           lowerCamelCase								=		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,)