
"""simple docstring"""






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

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


__SCREAMING_SNAKE_CASE							=     logging.get_logger(__name__)

__SCREAMING_SNAKE_CASE							=     '▁'

__SCREAMING_SNAKE_CASE							=     {'vocab_file': 'prophetnet.tokenizer'}

__SCREAMING_SNAKE_CASE							=     {
    'vocab_file': {
        'microsoft/xprophetnet-large-wiki100-cased': (
            'https://huggingface.co/microsoft/xprophetnet-large-wiki100-cased/resolve/main/prophetnet.tokenizer'
        ),
    }
}

__SCREAMING_SNAKE_CASE							=     {
    'microsoft/xprophetnet-large-wiki100-cased': {'do_lower_case': False},
}

__SCREAMING_SNAKE_CASE							=     {
    'microsoft/xprophetnet-large-wiki100-cased': 512,
}




def       A_  (				__lowercase	):
 UpperCamelCase_     :					Union[str, Any]   	=collections.OrderedDict()
 with open(__lowercase						,	'r'						,	encoding='utf-8'	) as reader:
  UpperCamelCase_     :					Optional[Any]   	=reader.readlines()
 for index, token in enumerate(__lowercase	):
  UpperCamelCase_     :					Any   	=token.rstrip('\n'	)
  UpperCamelCase_     :					List[str]   	=index
 return vocab





class 				a__					(						A__       ):
     UpperCAmelCase__         =      VOCAB_FILES_NAMES
     UpperCAmelCase__         =      PRETRAINED_VOCAB_FILES_MAP
     UpperCAmelCase__         =      PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
     UpperCAmelCase__         =      ['''input_ids''', '''attention_mask''']



     def __init__(							self :int      ,  _lowerCamelCase :int      ,  _lowerCamelCase :Optional[Any]="[SEP]"      ,  _lowerCamelCase :Dict="[SEP]"      ,  _lowerCamelCase :Optional[Any]="[SEP]"      ,  _lowerCamelCase :Tuple="[UNK]"      ,  _lowerCamelCase :List[Any]="[PAD]"      ,  _lowerCamelCase :Any="[CLS]"      ,  _lowerCamelCase :Optional[Any]="[MASK]"      ,  _lowerCamelCase :Optional[Dict[str, Any]] = None      ,  **_lowerCamelCase :Tuple      ,  ):







      '''simple docstring'''



      UpperCamelCase_     :					int   	={} if sp_model_kwargs is None else sp_model_kwargs

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

      try:
       import sentencepiece as spm
      except ImportError:
       logger.warning(
           'You need to install SentencePiece to use XLMRobertaTokenizer: https://github.com/google/sentencepiece'
           ' pip install sentencepiece'   )
       raise

      UpperCamelCase_     :					str   	=spm.SentencePieceProcessor(**self.sp_model_kwargs   )
      self.sp_model.Load(str(_lowerCamelCase   )   )
      UpperCamelCase_     :					Tuple   	=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'

      # put special tokens and [unused] tokens into the vocab
      UpperCamelCase_     :					int   	={'[PAD]': 0, '[CLS]': 1, '[SEP]': 2, '[UNK]': 3, '[MASK]': 4}

      for i in range(10   ):
       UpperCamelCase_     :					Optional[Any]   	=f'''[unused{i}]'''
       UpperCamelCase_     :					List[str]   	=5 + i

      # The first "real" token "," has position 15 in the embedding vocab and position 3 in the spm vocab
      UpperCamelCase_     :					int   	=12
      UpperCamelCase_     :					Union[str, Any]   	={v: k for k, v in self.fairseq_tokens_to_ids.items()}
      for k in self.fairseq_tokens_to_ids.keys():
       self.unique_no_split_tokens.append(_lowerCamelCase   )



     def __getstate__(							self :Tuple   ):







      '''simple docstring'''



      UpperCamelCase_     :					List[Any]   	=self.__dict__.copy()
      UpperCamelCase_     :					Any   	=None
      return state



     def __setstate__(							self :List[str]      ,  _lowerCamelCase :Union[str, Any]   ):







      '''simple docstring'''



      UpperCamelCase_     :					Optional[Any]   	=d
      try:
       import sentencepiece as spm
      except ImportError:
       logger.warning(
           'You need to install SentencePiece to use XLMRobertaTokenizer: https://github.com/google/sentencepiece'
           ' pip install sentencepiece'   )
       raise

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

      UpperCamelCase_     :					str   	=spm.SentencePieceProcessor(**self.sp_model_kwargs   )
      self.sp_model.Load(self.vocab_file   )



     def  lowerCamelCase_						(							self :Union[str, Any]      ,  _lowerCamelCase :List[int]      ,  _lowerCamelCase :Optional[List[int]] = None      ,  _lowerCamelCase :bool = 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 ([0] * len(_lowerCamelCase   )) + [1]
      return ([0] * len(_lowerCamelCase   )) + [1] + ([0] * len(_lowerCamelCase   )) + [1]



     def  lowerCamelCase_						(							self :Any      ,  _lowerCamelCase :List[int]      ,  _lowerCamelCase :Optional[List[int]] = None   ):







      '''simple docstring'''



      UpperCamelCase_     :					Tuple   	=[self.sep_token_id]

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



     @property
     def  lowerCamelCase_						(							self :Dict   ):







      '''simple docstring'''



      return len(self.sp_model   ) + self.fairseq_offset



     def  lowerCamelCase_						(							self :List[Any]   ):







      '''simple docstring'''



      UpperCamelCase_     :					Union[str, Any]   	={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 :str      ,  _lowerCamelCase :str   ):







      '''simple docstring'''



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



     def  lowerCamelCase_						(							self :Optional[int]      ,  _lowerCamelCase :Union[str, Any]   ):







      '''simple docstring'''



      if token in self.fairseq_tokens_to_ids:
       return self.fairseq_tokens_to_ids[token]
      UpperCamelCase_     :					List[Any]   	=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 :Optional[int]      ,  _lowerCamelCase :List[str]   ):







      '''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 :Tuple      ,  _lowerCamelCase :Optional[Any]   ):







      '''simple docstring'''



      UpperCamelCase_     :					str   	=''.join(_lowerCamelCase   ).replace(_lowerCamelCase      ,  ' '   ).strip()
      return out_string



     def  lowerCamelCase_						(							self :Any      ,  _lowerCamelCase :str      ,  _lowerCamelCase :Optional[str] = None   ):







      '''simple docstring'''



      if not os.path.isdir(_lowerCamelCase   ):
       logger.error(f'''Vocabulary path ({save_directory}) should be a directory'''   )
       return
      UpperCamelCase_     :					Optional[int]   	=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:
        UpperCamelCase_     :					List[str]   	=self.sp_model.serialized_model_proto()
        fi.write(_lowerCamelCase   )

      return (out_vocab_file,)




     def  lowerCamelCase_						(							self :List[str]      ,  _lowerCamelCase :List[int]      ,  _lowerCamelCase :Optional[List[int]] = None   ):







      '''simple docstring'''



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


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



     def  lowerCamelCase_						(							self :List[str]   ):







      '''simple docstring'''



      UpperCamelCase_     :					Optional[Any]   	=tempfile.mkdtemp()
      UpperCamelCase_     :					str   	=5

      # Realm tok
      UpperCamelCase_     :					List[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',
      ]
      UpperCamelCase_     :					List[str]   	=os.path.join(self.tmpdirname      ,  'realm_tokenizer'   )
      os.makedirs(_lowerCamelCase      ,  exist_ok=_lowerCamelCase   )
      UpperCamelCase_     :					int   	=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]   )   )

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



     def  lowerCamelCase_						(							self :Union[str, Any]   ):







      '''simple docstring'''



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



     def  lowerCamelCase_						(							self :Dict   ):







      '''simple docstring'''



      shutil.rmtree(self.tmpdirname   )



     def  lowerCamelCase_						(							self :Dict   ):







      '''simple docstring'''



      UpperCamelCase_     :					List[str]   	=RealmConfig(num_block_records=self.num_block_records   )
      return config



     def  lowerCamelCase_						(							self :Optional[int]   ):







      '''simple docstring'''



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



     def  lowerCamelCase_						(							self :Optional[Any]   ):







      '''simple docstring'''



      UpperCamelCase_     :					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=_lowerCamelCase      ,  )
      return block_records



     def  lowerCamelCase_						(							self :Any   ):







      '''simple docstring'''



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



     def  lowerCamelCase_						(							self :int   ):







      '''simple docstring'''



      UpperCamelCase_     :					Any   	=self.get_config()
      UpperCamelCase_     :					Tuple   	=self.get_dummy_retriever()
      UpperCamelCase_     :					Tuple   	=retriever.tokenizer

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

      UpperCamelCase_	,     UpperCamelCase_	,     UpperCamelCase_	,     UpperCamelCase_     :					Any   	=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, 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  lowerCamelCase_						(							self :Optional[int]   ):







      '''simple docstring'''



      UpperCamelCase_     :					Optional[int]   	=self.get_config()
      UpperCamelCase_     :					int   	=self.get_dummy_retriever()
      UpperCamelCase_     :					Any   	=retriever.tokenizer

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

      UpperCamelCase_	,     UpperCamelCase_	,     UpperCamelCase_	,     UpperCamelCase_     :					Union[str, Any]   	=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  lowerCamelCase_						(							self :List[Any]   ):







      '''simple docstring'''



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

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

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