
'''simple docstring'''







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

import sentencepiece as spm

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


lowercase    :							int      =					logging.get_logger(__name__)

lowercase    :							Optional[int]      =					"▁"

lowercase    :							Any      =					{"vocab_file": "sentencepiece.bpe.model"}

lowercase    :							List[Any]      =					{
    "vocab_file": {
        "facebook/xglm-564M": "https://huggingface.co/facebook/xglm-564M/resolve/main/sentencepiece.bpe.model",
    }
}

lowercase    :							Tuple      =					{
    "facebook/xglm-564M": 2048,
}


class 						__UpperCAmelCase      (					_lowerCamelCase						):
							__lowercase      		=       VOCAB_FILES_NAMES
							__lowercase      		=       PRETRAINED_VOCAB_FILES_MAP
							__lowercase      		=       PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
							__lowercase      		=       ["""input_ids""", """attention_mask"""]


							def __init__(       self    ,   lowerCAmelCase_    ,   lowerCAmelCase_="<s>"    ,   lowerCAmelCase_="</s>"    ,   lowerCAmelCase_="</s>"    ,   lowerCAmelCase_="<s>"    ,   lowerCAmelCase_="<unk>"    ,   lowerCAmelCase_="<pad>"    ,   lowerCAmelCase_ = None    ,   **lowerCAmelCase_    ,   ):





								"""simple docstring"""






								_snake_case   						=		{} if sp_model_kwargs is None else sp_model_kwargs

								# Compatibility with the original tokenizer
								_snake_case   						=		7
								_snake_case   						=		[F'<madeupword{i}>' for i in range(self.num_madeup_words  )]

								_snake_case   						=		kwargs.get('additional_special_tokens'    ,   []  )
								kwargs["additional_special_tokens"] += [
								    word for word in madeup_words if word not in kwargs["additional_special_tokens"]
								]

								super().__init__(
								    bos_token=lowerCAmelCase_    ,   eos_token=lowerCAmelCase_    ,   unk_token=lowerCAmelCase_    ,   sep_token=lowerCAmelCase_    ,   cls_token=lowerCAmelCase_    ,   pad_token=lowerCAmelCase_    ,   sp_model_kwargs=self.sp_model_kwargs    ,   **lowerCAmelCase_    ,   )

								_snake_case   						=		spm.SentencePieceProcessor(**self.sp_model_kwargs  )
								self.sp_model.Load(str(lowerCAmelCase_  )  )
								_snake_case   						=		vocab_file

								# Original fairseq vocab and spm vocab must be "aligned":
								# Vocab    |    0    |    1    |   2    |    3    |  4  |  5  |  6  |   7   |   8   |  9
								# -------- | ------- | ------- | ------ | ------- | --- | --- | --- | ----- | ----- | ----
								# fairseq  | '<s>'   | '<pad>' | '</s>' | '<unk>' | ',' | '.' | '▁' | 's'   | '▁de' | '-'
								# spm      | '<unk>' | '<s>'   | '</s>' | ','     | '.' | '▁' | 's' | '▁de' | '-'   | '▁a'

								# The first "real" token "," has position 4 in the original fairseq vocab and position 3 in the spm vocab
								_snake_case   						=		1

								# Mimic fairseq token-to-id alignment for the first 4 token
								_snake_case   						=		{'<s>': 0, '<pad>': 1, '</s>': 2, '<unk>': 3}

								_snake_case   						=		len(self.sp_model  )
								_snake_case   						=		{F'<madeupword{i}>': sp_size + i + self.fairseq_offset for i in range(self.num_madeup_words  )}
								self.fairseq_tokens_to_ids.update(lowerCAmelCase_  )

								_snake_case   						=		{v: k for k, v in self.fairseq_tokens_to_ids.items()}


							def __getstate__(       self  ):





								"""simple docstring"""






								_snake_case   						=		self.__dict__.copy()
								_snake_case   						=		None
								_snake_case   						=		self.sp_model.serialized_model_proto()
								return state


							def __setstate__(       self    ,   lowerCAmelCase_  ):





								"""simple docstring"""






								_snake_case   						=		d

								# for backward compatibility
								if not hasattr(self    ,   'sp_model_kwargs'  ):
									_snake_case   						=		{}

								_snake_case   						=		spm.SentencePieceProcessor(**self.sp_model_kwargs  )
								self.sp_model.LoadFromSerializedProto(self.sp_model_proto  )


							def  lowerCamelCase				(       self    ,   lowerCAmelCase_    ,   lowerCAmelCase_ = None  ):





								"""simple docstring"""






								if token_ids_a is None:
									return [self.sep_token_id] + token_ids_a
								_snake_case   						=		[self.sep_token_id]
								return sep + token_ids_a + sep + sep + token_ids_a


							def  lowerCamelCase				(       self    ,   lowerCAmelCase_    ,   lowerCAmelCase_ = None    ,   lowerCAmelCase_ = False  ):





								"""simple docstring"""






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

								if token_ids_a is None:
									return [1] + ([0] * len(lowerCAmelCase_  ))
								return [1] + ([0] * len(lowerCAmelCase_  )) + [1, 1] + ([0] * len(lowerCAmelCase_  ))


							def  lowerCamelCase				(       self    ,   lowerCAmelCase_    ,   lowerCAmelCase_ = None  ):





								"""simple docstring"""






								_snake_case   						=		[self.sep_token_id]

								if token_ids_a is None:
									return len(sep + token_ids_a  ) * [0]
								return len(sep + token_ids_a + sep + sep + token_ids_a  ) * [0]


							@property
							def  lowerCamelCase				(       self  ):





								"""simple docstring"""






								return len(self.sp_model  ) + self.fairseq_offset + self.num_madeup_words


							def  lowerCamelCase				(       self  ):





								"""simple docstring"""






								_snake_case   						=		{self.convert_ids_to_tokens(lowerCAmelCase_  ): i for i in range(self.vocab_size  )}
								vocab.update(self.added_tokens_encoder  )
								return vocab


							def  lowerCamelCase				(       self    ,   lowerCAmelCase_  ):





								"""simple docstring"""






								return self.sp_model.encode(lowerCAmelCase_    ,   out_type=lowerCAmelCase_  )


							def  lowerCamelCase				(       self    ,   lowerCAmelCase_  ):





								"""simple docstring"""






								if token in self.fairseq_tokens_to_ids:
									return self.fairseq_tokens_to_ids[token]
								_snake_case   						=		self.sp_model.PieceToId(lowerCAmelCase_  )

								# Need to return unknown token if the SP model returned 0
								return spm_id + self.fairseq_offset if spm_id else self.unk_token_id


							def  lowerCamelCase				(       self    ,   lowerCAmelCase_  ):





								"""simple docstring"""






								if index in self.fairseq_ids_to_tokens:
									return self.fairseq_ids_to_tokens[index]
								return self.sp_model.IdToPiece(index - self.fairseq_offset  )


							def  lowerCamelCase				(       self    ,   lowerCAmelCase_  ):





								"""simple docstring"""






								_snake_case   						=		''.join(lowerCAmelCase_  ).replace(lowerCAmelCase_    ,   ' '  ).strip()
								return out_string




							def  lowerCamelCase				(       self    ,   lowerCAmelCase_    ,   lowerCAmelCase_ = None  ):





								"""simple docstring"""






								if not os.path.isdir(lowerCAmelCase_  ):
									logger.error(F'Vocabulary path ({save_directory}) should be a directory'  )
									return
								_snake_case   						=		os.path.join(
								    lowerCAmelCase_    ,   (filename_prefix + '-' if filename_prefix else '') + VOCAB_FILES_NAMES['vocab_file']  )

								if os.path.abspath(self.vocab_file  ) != os.path.abspath(lowerCAmelCase_  ) and os.path.isfile(self.vocab_file  ):
									copyfile(self.vocab_file    ,   lowerCAmelCase_  )
								elif not os.path.isfile(self.vocab_file  ):
									with open(lowerCAmelCase_    ,   'wb'  ) as fi:
										_snake_case   						=		self.sp_model.serialized_model_proto()
										fi.write(lowerCAmelCase_  )

								return (out_vocab_file,)








'''simple docstring'''




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  SCREAMING_SNAKE_CASE    (a__							):


     def 							SCREAMING_SNAKE_CASE  (						self):

          '''simple docstring'''






          __A						: Any      =       tempfile.mkdtemp()
          __A						: str      =       5

          # Realm tok
          __A						: Union[str, 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',
          ]
          __A						: Union[str, Any]      =       os.path.join(self.tmpdirname						,	'realm_tokenizer')
          os.makedirs(_UpperCAmelCase						,	exist_ok=_UpperCAmelCase)
          __A						: Tuple      =       os.path.join(_UpperCAmelCase						,	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]))

          __A						: int      =       os.path.join(self.tmpdirname						,	'realm_block_records')
          os.makedirs(_UpperCAmelCase						,	exist_ok=_UpperCAmelCase)


     def 							SCREAMING_SNAKE_CASE  (						self):

          '''simple docstring'''






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


     def 							SCREAMING_SNAKE_CASE  (						self):

          '''simple docstring'''






          shutil.rmtree(self.tmpdirname)


     def 							SCREAMING_SNAKE_CASE  (						self):

          '''simple docstring'''






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


     def 							SCREAMING_SNAKE_CASE  (						self):

          '''simple docstring'''






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


     def 							SCREAMING_SNAKE_CASE  (						self):

          '''simple docstring'''






          __A						: 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=_UpperCAmelCase						,	)
          return block_records


     def 							SCREAMING_SNAKE_CASE  (						self):

          '''simple docstring'''






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


     def 							SCREAMING_SNAKE_CASE  (						self):

          '''simple docstring'''






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

          __A						: Dict      =       np.array([0, 3]						,	dtype='long')
          __A						: Dict      =       tokenizer(['Test question']).input_ids
          __A						: Optional[Any]      =       tokenizer(
              ['the fourth']						,	add_special_tokens=_UpperCAmelCase						,	return_token_type_ids=_UpperCAmelCase						,	return_attention_mask=_UpperCAmelCase						,	).input_ids
          __A						: str      =       config.reader_seq_len

          __A ,__A ,__A ,__A						: List[str]      =       retriever(
              _UpperCAmelCase						,	_UpperCAmelCase						,	answer_ids=_UpperCAmelCase						,	max_length=_UpperCAmelCase						,	return_tensors='np')

          self.assertEqual(len(_UpperCAmelCase)						,	2)
          self.assertEqual(len(_UpperCAmelCase)						,	2)
          self.assertEqual(len(_UpperCAmelCase)						,	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 							SCREAMING_SNAKE_CASE  (						self):

          '''simple docstring'''






          __A						: int      =       self.get_config()
          __A						: Any      =       self.get_dummy_retriever()
          __A						: str      =       retriever.tokenizer

          __A						: Dict      =       np.array([0, 3, 5]						,	dtype='long')
          __A						: Tuple      =       tokenizer(['Test question']).input_ids
          __A						: Union[str, Any]      =       tokenizer(
              ['the fourth', 'longer longer']						,	add_special_tokens=_UpperCAmelCase						,	return_token_type_ids=_UpperCAmelCase						,	return_attention_mask=_UpperCAmelCase						,	).input_ids
          __A						: Dict      =       config.reader_seq_len

          __A ,__A ,__A ,__A						: str      =       retriever(
              _UpperCAmelCase						,	_UpperCAmelCase						,	answer_ids=_UpperCAmelCase						,	max_length=_UpperCAmelCase						,	return_tensors='np')

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


     def 							SCREAMING_SNAKE_CASE  (						self):

          '''simple docstring'''






          __A						: Optional[Any]      =       self.get_dummy_retriever()
          retriever.save_pretrained(os.path.join(self.tmpdirname						,	'realm_block_records'))

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

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