






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       snake_case__ (				lowercase_):

       '''simple docstring'''



       def        __lowercase    (				self     )			->   int:






            '''simple docstring'''





            __snake_case      :int 							=    tempfile.mkdtemp()
            __snake_case      :int 							=    5

            # Realm tok
            __snake_case      :Optional[Any] 							=    [
                """[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""",
            ]
            __snake_case      :Dict 							=    os.path.join(self.tmpdirname      ,    """realm_tokenizer"""     )
            os.makedirs(a__      ,    exist_ok=a__     )
            __snake_case      :Dict 							=    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]     )     )

            __snake_case      :Optional[int] 							=    os.path.join(self.tmpdirname      ,    """realm_block_records"""     )
            os.makedirs(a__      ,    exist_ok=a__     )



       def        __lowercase    (				self     )			->   RealmTokenizer:






            '''simple docstring'''





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



       def        __lowercase    (				self     )			->   Optional[Any]:






            '''simple docstring'''





            shutil.rmtree(self.tmpdirname     )



       def        __lowercase    (				self     )			->   Optional[int]:






            '''simple docstring'''





            __snake_case      :Any 							=    RealmConfig(num_block_records=self.num_block_records     )
            return config



       def        __lowercase    (				self     )			->   Union[str, Any]:






            '''simple docstring'''





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



       def        __lowercase    (				self     )			->   Optional[int]:






            '''simple docstring'''





            __snake_case      :Union[str, Any] 							=    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        __lowercase    (				self     )			->   Tuple:






            '''simple docstring'''





            __snake_case      :Optional[Any] 							=    RealmRetriever(
                block_records=self.get_dummy_block_records()      ,    tokenizer=self.get_tokenizer()      ,    )
            return retriever



       def        __lowercase    (				self     )			->   Union[str, Any]:






            '''simple docstring'''





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

            __snake_case      :Optional[Any] 							=    np.array([0, 3]      ,    dtype="""long"""     )
            __snake_case      :str 							=    tokenizer(["""Test question"""]     ).input_ids
            __snake_case      :int 							=    tokenizer(
                ["""the fourth"""]      ,    add_special_tokens=a__      ,    return_token_type_ids=a__      ,    return_attention_mask=a__      ,    ).input_ids
            __snake_case      :List[Any] 							=    config.reader_seq_len

            __snake_case							,			__snake_case							,			__snake_case							,			__snake_case      :Dict 							=    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        __lowercase    (				self     )			->   Tuple:






            '''simple docstring'''





            __snake_case      :Any 							=    self.get_config()
            __snake_case      :List[str] 							=    self.get_dummy_retriever()
            __snake_case      :Dict 							=    retriever.tokenizer

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

            __snake_case							,			__snake_case							,			__snake_case							,			__snake_case      :int 							=    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        __lowercase    (				self     )			->   str:






            '''simple docstring'''





            __snake_case      :int 							=    self.get_dummy_retriever()
            retriever.save_pretrained(os.path.join(self.tmpdirname      ,    """realm_block_records"""     )     )

            # Test local path
            __snake_case      :int 							=    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:
                 __snake_case      :List[str] 							=    os.path.join(
                     os.path.join(self.tmpdirname      ,    """realm_block_records"""     )      ,    _REALM_BLOCK_RECORDS_FILENAME     )
                 __snake_case      :List[str] 							=    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

import sentencepiece as spm

from ...tokenization_utils import PreTrainedTokenizer
from ...utils import logging


lowerCamelCase__       		=  logging.get_logger(__name__)

lowerCamelCase__       		=  {"""vocab_file""": """sentencepiece.model"""}

lowerCamelCase__       		=  {
    """vocab_file""": {
        """google/rembert""": """https://huggingface.co/google/rembert/resolve/main/sentencepiece.model""",
    },
}

lowerCamelCase__       		=  {
    """google/rembert""": 256,
}




class       snake_case__ (				lowercase_):

       '''simple docstring'''
       lowerCamelCase				:					str								=			VOCAB_FILES_NAMES
       lowerCamelCase				:					Tuple								=			PRETRAINED_VOCAB_FILES_MAP
       lowerCamelCase				:					Optional[int]								=			PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES



       def __init__(				self      ,    a__      ,    a__=False      ,    a__=True      ,    a__=True      ,    a__="[CLS]"      ,    a__="[SEP]"      ,    a__="[UNK]"      ,    a__="[SEP]"      ,    a__="[PAD]"      ,    a__="[CLS]"      ,    a__="[MASK]"      ,    **a__      ,    )			->   List[str]:






            '''simple docstring'''





            super().__init__(
                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__      ,    )

            __snake_case      :List[Any] 							=    do_lower_case
            __snake_case      :Dict 							=    remove_space
            __snake_case      :int 							=    keep_accents
            __snake_case      :Optional[int] 							=    vocab_file

            __snake_case      :int 							=    spm.SentencePieceProcessor()
            self.sp_model.Load(a__     )



       @property
       def        __lowercase    (				self     )			->   Tuple:






            '''simple docstring'''





            return len(self.sp_model     )



       def        __lowercase    (				self     )			->   Optional[int]:






            '''simple docstring'''





            __snake_case      :int 							=    {self.convert_ids_to_tokens(a__     ): i for i in range(self.vocab_size     )}
            vocab.update(self.added_tokens_encoder     )
            return vocab



       def __getstate__(				self     )			->   str:






            '''simple docstring'''





            __snake_case      :str 							=    self.__dict__.copy()
            __snake_case      :Dict 							=    None
            return state



       def __setstate__(				self      ,    a__     )			->   str:






            '''simple docstring'''





            __snake_case      :Dict 							=    d
            __snake_case      :Dict 							=    spm.SentencePieceProcessor()
            self.sp_model.Load(self.vocab_file     )



       def        __lowercase    (				self      ,    a__      ,    a__=False     )			->   Tuple:






            '''simple docstring'''





            __snake_case      :Union[str, Any] 							=    self.sp_model.EncodeAsPieces(a__     )
            return pieces



       def        __lowercase    (				self      ,    a__     )			->   str:






            '''simple docstring'''





            return self.sp_model.PieceToId(a__     )



       def        __lowercase    (				self      ,    a__     )			->   Any:






            '''simple docstring'''





            return self.sp_model.IdToPiece(a__     )



       def        __lowercase    (				self      ,    a__     )			->   Union[str, Any]:






            '''simple docstring'''





            __snake_case      :List[Any] 							=    self.sp_model.decode_pieces(a__     )
            return out_string



       def        __lowercase    (				self      ,    a__      ,    a__ = None     )			->   List[int]:






            '''simple docstring'''





            __snake_case      :Optional[int] 							=    [self.sep_token_id]
            __snake_case      :List[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        __lowercase    (				self      ,    a__      ,    a__ = None      ,    a__ = False     )			->   List[int]:






            '''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        __lowercase    (				self      ,    a__      ,    a__ = None     )			->   List[int]:






            '''simple docstring'''





            __snake_case      :Any 							=    [self.sep_token_id]
            __snake_case      :Any 							=    [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      ,    a__      ,    a__ = None     )			->   Tuple[str]:






            '''simple docstring'''





            if not os.path.isdir(a__     ):
                 logger.error("""Vocabulary path ({}) should be a directory""".format(a__     )     )
                 return
            __snake_case      :Optional[int] 							=    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,)


