
import collections
import json
import os
import re
from typing import TYPE_CHECKING, List, Optional, Tuple

import numpy as np

from ...tokenization_utils_fast import PreTrainedTokenizer
from ...utils import logging


if TYPE_CHECKING:
 from transformers.pipelines.conversational import Conversation


__lowerCAmelCase					:			Tuple			      =		logging.get_logger(__name__)

__lowerCAmelCase					:			List[str]			      =		{"vocab_file": "vocab.txt", "emoji_file": "emoji.json"}

__lowerCAmelCase					:			Any			      =		{
    "vocab_file": {
        "abeja/gpt-neox-japanese-2.7b": "https://huggingface.co/abeja/gpt-neox-japanese-2.7b/resolve/main/vocab.txt",
    },
    "emoji_file": {
        "abeja/gpt-neox-japanese-2.7b": "https://huggingface.co/abeja/gpt-neox-japanese-2.7b/resolve/main/emoji.json",
    },
}

__lowerCAmelCase					:			Optional[Any]			      =		{
    "abeja/gpt-neox-japanese-2.7b": 2_048,
}


def     UpperCAmelCase_  (						__lowerCAmelCase	,					__lowerCAmelCase	) ->						List[Any]:
    with open(__lowerCAmelCase	,					'''r'''	,					encoding='''utf-8'''	) as f:
        __lowercase		:			List[Any]	    =       json.loads(f.read()	)

    __lowercase		:			Tuple	    =       collections.OrderedDict()
    __lowercase		:			Optional[Any]	    =       collections.OrderedDict()
    __lowercase		:			List[Any]	    =       collections.OrderedDict()
    with open(__lowerCAmelCase	,					'''r'''	,					encoding='''utf-8'''	) as f:
        __lowercase		:			str	    =       f.readlines()
    __lowercase		:			int	    =       [[t.rstrip('''\n'''	)] if (t == ''',''' or ''',''' not in t) else t.rstrip('''\n'''	).split(''','''	) for t in token]
    for idx, b in enumerate(__lowerCAmelCase	):
        __lowercase		:			Tuple	    =       b
        __lowercase		:			str	    =       idx
        for wd in b:
            __lowercase		:			List[Any]	    =       idx

    return vocab, raw_vocab, ids_to_tokens, emoji


class 						__lowerCAmelCase	(	lowerCAmelCase_      ):



      """simple docstring"""







      A__	:			str        =    VOCAB_FILES_NAMES
      A__	:			Optional[int]        =    PRETRAINED_VOCAB_FILES_MAP
      A__	:			Tuple        =    PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
      A__	:			int        =    ['''input_ids''', '''attention_mask''']




      def __init__(      self		:       List[str]				,       _snake_case		:       str				,       _snake_case		:       Optional[int]				,       _snake_case		:       Union[str, Any]="<|endoftext|>"				,       _snake_case		:       int="<|endoftext|>"				,       _snake_case		:       str="<|startoftext|>"				,       _snake_case		:       List[str]="<|endoftext|>"				,       _snake_case		:       int=False				,       **_snake_case		:       List[Any]				,       ):
          super().__init__(
              unk_token=_snake_case				,       pad_token=_snake_case				,       bos_token=_snake_case				,       eos_token=_snake_case				,       do_clean_text=_snake_case				,       **_snake_case				,       )
          if not os.path.isfile(_snake_case				):
              raise ValueError(
                  F'Can\'t find a vocabulary file at path \'{vocab_file}\'. To load the vocabulary from a Google pretrained'
                  ''' model use `tokenizer = GPTNeoXJapaneseokenizer.from_pretrained(PRETRAINED_MODEL_NAME)`'''				)
          if not os.path.isfile(_snake_case				):
              raise ValueError(
                  F'Can\'t find a emoji file at path \'{emoji_file}\'. To load the emoji information from a Google'
                  ''' pretrained model use `tokenizer = GPTNeoXJapaneseokenizer.from_pretrained(PRETRAINED_MODEL_NAME)`'''				)
          __lowercase		:			Optional[Any]	    =       do_clean_text
          __lowercase    ,      __lowercase    ,      __lowercase    ,      __lowercase		:			Tuple	    =       load_vocab_and_emoji(_snake_case				,       _snake_case				)
          __lowercase		:			Any	    =       SubWordJapaneseTokenizer(
              vocab=self.vocab				,       ids_to_tokens=self.ids_to_tokens				,       emoji=self.emoji				)




      @property
      def 		snake_case_			(      self		:       Optional[Any]				):
          # self.vocab contains support for character fluctuation unique to Japanese, and has a large number of vocab
          return len(self.raw_vocab				)




      def 		snake_case_			(      self		:       Any				):
          return dict(self.raw_vocab				,       **self.added_tokens_encoder				)




      def 		snake_case_			(      self		:       Union[str, Any]				,       _snake_case		:       Tuple				):
          return self.subword_tokenizer.tokenize(_snake_case				,       clean=self.do_clean_text				)




      def 		snake_case_			(      self		:       Tuple				,       _snake_case		:       Tuple				):
          return self.vocab.get(_snake_case				,       self.vocab.get(self.unk_token				)				)




      def 		snake_case_			(      self		:       Tuple				,       _snake_case		:       Any				):
          return self.subword_tokenizer.convert_id_to_token(_snake_case				)




      def 		snake_case_			(      self		:       Optional[Any]				,       _snake_case		:       List[str]				):
          __lowercase		:			List[str]	    =       ''''''.join(_snake_case				).strip()
          return out_string




      def 		snake_case_			(      self		:       List[Any]				,       _snake_case		:       "Conversation"				):
          __lowercase		:			Dict	    =       []
          for is_user, text in conversation.iter_texts():
              input_ids.extend(self.encode(_snake_case				,       add_special_tokens=_snake_case				) + [self.eos_token_id]				)

          if len(_snake_case				) > self.model_max_length:
              __lowercase		:			Union[str, Any]	    =       input_ids[-self.model_max_length :]
          return input_ids




      def 		snake_case_			(      self		:       Optional[int]				,       _snake_case		:       str				,       _snake_case		:       Optional[str] = None				):
          __lowercase		:			List[Any]	    =       0
          if os.path.isdir(_snake_case				):
              __lowercase		:			List[str]	    =       os.path.join(
                  _snake_case				,       (filename_prefix + '''-''' if filename_prefix else '''''') + VOCAB_FILES_NAMES['''vocab_file''']				)
              __lowercase		:			int	    =       os.path.join(
                  _snake_case				,       (filename_prefix + '''-''' if filename_prefix else '''''') + VOCAB_FILES_NAMES['''emoji_file''']				)
          else:
              __lowercase		:			Union[str, Any]	    =       (
                  (filename_prefix + '''-''' if filename_prefix else '''''') + save_directory + VOCAB_FILES_NAMES['''vocab_file''']
              )
              __lowercase		:			Optional[int]	    =       (
                  (filename_prefix + '''-''' if filename_prefix else '''''') + save_directory + VOCAB_FILES_NAMES['''emoji_file''']
              )
          with open(_snake_case				,       '''w'''				,       encoding='''utf-8'''				) as writer:
              for token_index, token in self.ids_to_tokens.items():
                  if index != token_index:
                      logger.warning(
                          F'Saving vocabulary to {vocab_file}: vocabulary indices are not consecutive.'
                          ''' Please check that the vocabulary is not corrupted!'''				)
                      __lowercase		:			List[str]	    =       token_index
                  writer.write(''','''.join(_snake_case				) + '''\n'''				)
                  index += 1
          with open(_snake_case				,       '''w'''				,       encoding='''utf-8'''				) as writer:
              json.dump(self.emoji				,       _snake_case				)
          return vocab_file, emoji_file






class 						__lowerCAmelCase	(	lowerCAmelCase_      ):



      """simple docstring"""







      def __init__(      self		:       Any				,       _snake_case		:       Tuple				,       _snake_case		:       Union[str, Any]				,       _snake_case		:       List[Any]				):
          __lowercase		:			int	    =       vocab  # same as swe
          __lowercase		:			List[Any]	    =       ids_to_tokens  # same as bpe
          __lowercase		:			List[str]	    =       emoji
          __lowercase		:			Optional[Any]	    =       np.max([len(_snake_case				) for w in self.vocab.keys()]				)
          __lowercase		:			str	    =       re.compile(r'''(https?|ftp)(:\/\/[-_\.!~*\'()a-zA-Z0-9;\/?:\@&=\+$,%#]+)'''				)
          __lowercase		:			Optional[int]	    =       re.compile(r'''[A-Za-z0-9\._+]*@[\-_0-9A-Za-z]+(\.[A-Za-z]+)*'''				)
          __lowercase		:			Optional[Any]	    =       re.compile(r'''[\(]{0,1}[0-9]{2,4}[\)\-\(]{0,1}[0-9]{2,4}[\)\-]{0,1}[0-9]{3,4}'''				)
          __lowercase		:			Union[str, Any]	    =       re.compile(
              r'''([12]\d{3}[/\-年])*(0?[1-9]|1[0-2])[/\-月]((0?[1-9]|[12][0-9]|3[01])日?)*(\d{1,2}|:|\d{1,2}時|\d{1,2}分|\(日\)|\(月\)|\(火\)|\(水\)|\(木\)|\(金\)|\(土\)|㈰|㈪|㈫|㈬|㈭|㈮|㈯)*'''				)
          __lowercase		:			Dict	    =       re.compile(
              r'''(明治|大正|昭和|平成|令和|㍾|㍽|㍼|㍻|\u32ff)\d{1,2}年(0?[1-9]|1[0-2])月(0?[1-9]|[12][0-9]|3[01])日(\d{1,2}|:|\d{1,2}時|\d{1,2}分|\(日\)|\(月\)|\(火\)|\(水\)|\(木\)|\(金\)|\(土\)|㈰|㈪|㈫|㈬|㈭|㈮|㈯)*'''				)
          __lowercase		:			Optional[Any]	    =       re.compile(
              r'''((0|[1-9]\d*|[1-9]\d{0,2}(,\d{3})+)*億)*((0|[1-9]\d*|[1-9]\d{0,2}(,\d{3})+)*万)*((0|[1-9]\d*|[1-9]\d{0,2}(,\d{3})+)*千)*(0|[1-9]\d*|[1-9]\d{0,2}(,\d{3})+)*(千円|万円|千万円|円|千ドル|万ドル|千万ドル|ドル|千ユーロ|万ユーロ|千万ユーロ|ユーロ)+(\(税込\)|\(税抜\)|\+tax)*'''				)
          __lowercase		:			List[Any]	    =       '''─━│┃┄┅┆┇┈┉┊┋┌┍┎┏┐┑┒┓└┕┖┗┘┙┚┛├┝┞┟┠┡┢┣┤┥┦┧┨┩┪┫┬┭┮┯┰┱┲┳┴┵┶┷┸┹┺┻┼┽┾┿╀╁╂╃╄╅╆╇╈╉╊╋╌╍╎╏═║╒╓╔╕╖╗╘╙╚╛╜╝╞╟╠╡╢╣╤╥╦╧╨╩╪╫╬╭╮╯╰╱╲╳╴╵╶╷╸╹╺╻╼╽╾╿'''
          __lowercase		:			str	    =       '''▀▁▂▃▄▅▆▇█▉▊▋▌▍▎▏▐░▒▓▔▕▖▗▘▙▚▛▜▝▞▟'''
          __lowercase		:			Dict	    =       str.maketrans({k: '''<BLOCK>''' for k in keisen + blocks}				)




      def __len__(      self		:       List[Any]				):
          return len(self.ids_to_tokens				)




      def 		snake_case_			(      self		:       List[str]				,       _snake_case		:       Optional[int]				):
          __lowercase		:			List[Any]	    =       self.content_repattera.sub('''<URL>'''				,       _snake_case				)
          __lowercase		:			Optional[int]	    =       self.content_repattera.sub('''<EMAIL>'''				,       _snake_case				)
          __lowercase		:			List[Any]	    =       self.content_repattera.sub('''<TEL>'''				,       _snake_case				)
          __lowercase		:			Optional[int]	    =       self.content_repattera.sub('''<DATE>'''				,       _snake_case				)
          __lowercase		:			Optional[Any]	    =       self.content_repattera.sub('''<DATE>'''				,       _snake_case				)
          __lowercase		:			Union[str, Any]	    =       self.content_repattera.sub('''<PRICE>'''				,       _snake_case				)
          __lowercase		:			int	    =       content.translate(self.content_transa				)
          while "<BLOCK><BLOCK>" in content:
              __lowercase		:			List[str]	    =       content.replace('''<BLOCK><BLOCK>'''				,       '''<BLOCK>'''				)
          return content




      def 		snake_case_			(      self		:       Tuple				,       _snake_case		:       int				,       _snake_case		:       Tuple=False				):
          __lowercase		:			Union[str, Any]	    =       text.replace(''' '''				,       '''<SP>'''				)
          __lowercase		:			Optional[Any]	    =       text.replace('''　'''				,       '''<SP>'''				)
          __lowercase		:			int	    =       text.replace('''\r\n'''				,       '''<BR>'''				)
          __lowercase		:			Any	    =       text.replace('''\n'''				,       '''<BR>'''				)
          __lowercase		:			Optional[Any]	    =       text.replace('''\r'''				,       '''<BR>'''				)
          __lowercase		:			Union[str, Any]	    =       text.replace('''\t'''				,       '''<TAB>'''				)
          __lowercase		:			int	    =       text.replace('''—'''				,       '''ー'''				)
          __lowercase		:			Optional[Any]	    =       text.replace('''−'''				,       '''ー'''				)
          for k, v in self.emoji["emoji"].items():
              if k in text:
                  __lowercase		:			int	    =       text.replace(_snake_case				,       _snake_case				)
          if clean:
              __lowercase		:			List[Any]	    =       self.clean_text(_snake_case				)

          def check_simbol(_snake_case		:       Dict				):
              __lowercase		:			List[Any]	    =       x.encode()
              if len(_snake_case				) == 1 and len(_snake_case				) == 2:
                  __lowercase		:			Optional[Any]	    =       (int(e[0]				) << 8) + int(e[1]				)
                  if (
                      (c >= 0XC_2_A_1 and c <= 0XC_2_B_F)
                      or (c >= 0XC_7_8_0 and c <= 0XC_7_8_3)
                      or (c >= 0XC_A_B_9 and c <= 0XC_B_B_F)
                      or (c >= 0XC_C_8_0 and c <= 0XC_D_A_2)
                  ):
                      return True
              return False

          def checkuae(_snake_case		:       List[str]				):
              __lowercase		:			Optional[int]	    =       x.encode()
              if len(_snake_case				) == 1 and len(_snake_case				) == 3:
                  __lowercase		:			Optional[Any]	    =       (int(e[0]				) << 16) + (int(e[1]				) << 8) + int(e[2]				)
                  if c >= 0XE_2_8_0_8_0 and c <= 0XE_2_B_0_7_F:
                      return True
              return False

          __lowercase		:			str	    =       0
          __lowercase		:			Dict	    =       []
          while pos < len(_snake_case				):
              __lowercase		:			Optional[Any]	    =       min(len(_snake_case				)				,       pos + self.maxlen + 1				) if text[pos] == '''<''' else pos + 3
              __lowercase		:			List[Any]	    =       []  # (token_id, token, pos)
              for e in range(_snake_case				,       _snake_case				,       -1				):
                  __lowercase		:			Tuple	    =       text[pos:e]
                  if wd in self.vocab:
                      if wd[0] == "<" and len(_snake_case				) > 2:
                          __lowercase		:			Dict	    =       [(self.vocab[wd], wd, e)]
                          break
                      else:
                          candidates.append((self.vocab[wd], wd, e)				)
              if len(_snake_case				) > 0:
                  # the smallest token_id is adopted
                  __lowercase    ,      __lowercase    ,      __lowercase		:			Any	    =       sorted(_snake_case				,       key=lambda _snake_case				: x[0]				)[0]
                  result.append(_snake_case				)
                  __lowercase		:			int	    =       e
              else:
                  __lowercase		:			List[str]	    =       pos + 1
                  __lowercase		:			str	    =       text[pos:end]
                  if check_simbol(_snake_case				):
                      result.append('''<KIGOU>'''				)
                  elif checkuae(_snake_case				):
                      result.append('''<U2000U2BFF>'''				)
                  else:
                      for i in wd.encode('''utf-8'''				):
                          result.append('''<|byte%d|>''' % i				)
                  __lowercase		:			Optional[Any]	    =       end
          return result




      def 		snake_case_			(      self		:       Optional[int]				,       _snake_case		:       int				,       _snake_case		:       List[Any]="\n"				):
          __lowercase		:			Optional[int]	    =       []
          __lowercase		:			Union[str, Any]	    =       []
          __lowercase		:			Dict	    =       self.ids_to_tokens[index][0]
          if word[:6] == "<|byte" and word[-2:] == "|>":
              byte_tokens.append(int(word[6:-2]				)				)
          else:
              if len(_snake_case				) > 0:
                  words.append(bytearray(_snake_case				).decode('''utf-8'''				,       errors='''replace'''				)				)
                  __lowercase		:			int	    =       []
              if word[:7] == "<|emoji" and word[-2:] == "|>":
                  words.append(self.emoji['''emoji_inv'''][word]				)
              elif word == "<SP>":
                  words.append(''' '''				)
              elif word == "<BR>":
                  words.append(_snake_case				)
              elif word == "<TAB>":
                  words.append('''\t'''				)
              elif word == "<BLOCK>":
                  words.append('''▀'''				)
              elif word == "<KIGOU>":
                  words.append('''ǀ'''				)
              elif word == "<U2000U2BFF>":
                  words.append('''‖'''				)
              else:
                  words.append(_snake_case				)
          if len(_snake_case				) > 0:
              words.append(bytearray(_snake_case				).decode('''utf-8'''				,       errors='''replace'''				)				)
          __lowercase		:			Tuple	    =       ''''''.join(_snake_case				)
          return text




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

import numpy as np
from datasets import Dataset

from transformers import is_faiss_available
from transformers.models.bart.configuration_bart import BartConfig
from transformers.models.bart.tokenization_bart import BartTokenizer
from transformers.models.bert.tokenization_bert import VOCAB_FILES_NAMES as DPR_VOCAB_FILES_NAMES
from transformers.models.dpr.configuration_dpr import DPRConfig
from transformers.models.dpr.tokenization_dpr import DPRContextEncoderTokenizer, DPRQuestionEncoderTokenizer
from transformers.models.rag.configuration_rag import RagConfig
from transformers.models.rag.retrieval_rag import CustomHFIndex, RagRetriever
from transformers.models.roberta.tokenization_roberta import VOCAB_FILES_NAMES as BART_VOCAB_FILES_NAMES
from transformers.testing_utils import require_faiss, require_sentencepiece, require_tokenizers, require_torch


if is_faiss_available():
 import faiss


@require_faiss
class 						__lowerCAmelCase	(	lowerCAmelCase_      ):



      """simple docstring"""







      def 		snake_case_			(      self		:       Optional[int]				):
          __lowercase		:			Dict	    =       tempfile.mkdtemp()
          __lowercase		:			Tuple	    =       8

          # DPR tok
          __lowercase		:			Optional[int]	    =       [
              '''[UNK]''',
              '''[CLS]''',
              '''[SEP]''',
              '''[PAD]''',
              '''[MASK]''',
              '''want''',
              '''##want''',
              '''##ed''',
              '''wa''',
              '''un''',
              '''runn''',
              '''##ing''',
              ''',''',
              '''low''',
              '''lowest''',
          ]
          __lowercase		:			Optional[Any]	    =       os.path.join(self.tmpdirname				,       '''dpr_tokenizer'''				)
          os.makedirs(_snake_case				,       exist_ok=_snake_case				)
          __lowercase		:			str	    =       os.path.join(_snake_case				,       DPR_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]				)				)

          # BART tok
          __lowercase		:			List[str]	    =       [
              '''l''',
              '''o''',
              '''w''',
              '''e''',
              '''r''',
              '''s''',
              '''t''',
              '''i''',
              '''d''',
              '''n''',
              '''\u0120''',
              '''\u0120l''',
              '''\u0120n''',
              '''\u0120lo''',
              '''\u0120low''',
              '''er''',
              '''\u0120lowest''',
              '''\u0120newer''',
              '''\u0120wider''',
              '''<unk>''',
          ]
          __lowercase		:			Optional[Any]	    =       dict(zip(_snake_case				,       range(len(_snake_case				)				)				)				)
          __lowercase		:			int	    =       ['''#version: 0.2''', '''\u0120 l''', '''\u0120l o''', '''\u0120lo w''', '''e r''', '''''']
          __lowercase		:			Any	    =       {'''unk_token''': '''<unk>'''}

          __lowercase		:			int	    =       os.path.join(self.tmpdirname				,       '''bart_tokenizer'''				)
          os.makedirs(_snake_case				,       exist_ok=_snake_case				)
          __lowercase		:			List[Any]	    =       os.path.join(_snake_case				,       BART_VOCAB_FILES_NAMES['''vocab_file''']				)
          __lowercase		:			Tuple	    =       os.path.join(_snake_case				,       BART_VOCAB_FILES_NAMES['''merges_file''']				)
          with open(self.vocab_file				,       '''w'''				,       encoding='''utf-8'''				) as fp:
              fp.write(json.dumps(_snake_case				) + '''\n'''				)
          with open(self.merges_file				,       '''w'''				,       encoding='''utf-8'''				) as fp:
              fp.write('''\n'''.join(_snake_case				)				)




      def 		snake_case_			(      self		:       List[Any]				):
          return DPRQuestionEncoderTokenizer.from_pretrained(os.path.join(self.tmpdirname				,       '''dpr_tokenizer'''				)				)




      def 		snake_case_			(      self		:       Dict				):
          return DPRContextEncoderTokenizer.from_pretrained(os.path.join(self.tmpdirname				,       '''dpr_tokenizer'''				)				)




      def 		snake_case_			(      self		:       Optional[int]				):
          return BartTokenizer.from_pretrained(os.path.join(self.tmpdirname				,       '''bart_tokenizer'''				)				)




      def 		snake_case_			(      self		:       str				):
          shutil.rmtree(self.tmpdirname				)




      def 		snake_case_			(      self		:       Tuple				):
          __lowercase		:			Dict	    =       Dataset.from_dict(
              {
                  '''id''': ['''0''', '''1'''],
                  '''text''': ['''foo''', '''bar'''],
                  '''title''': ['''Foo''', '''Bar'''],
                  '''embeddings''': [np.ones(self.retrieval_vector_size				), 2 * np.ones(self.retrieval_vector_size				)],
              }				)
          dataset.add_faiss_index('''embeddings'''				,       string_factory='''Flat'''				,       metric_type=faiss.METRIC_INNER_PRODUCT				)
          return dataset




      def 		snake_case_			(      self		:       Union[str, Any]				):
          __lowercase		:			Union[str, Any]	    =       self.get_dummy_dataset()
          __lowercase		:			Optional[int]	    =       RagConfig(
              retrieval_vector_size=self.retrieval_vector_size				,       question_encoder=DPRConfig().to_dict()				,       generator=BartConfig().to_dict()				,       )
          with patch('''transformers.models.rag.retrieval_rag.load_dataset'''				) as mock_load_dataset:
              __lowercase		:			List[Any]	    =       dataset
              __lowercase		:			str	    =       RagRetriever(
                  _snake_case				,       question_encoder_tokenizer=self.get_dpr_tokenizer()				,       generator_tokenizer=self.get_bart_tokenizer()				,       )
          return retriever




      def 		snake_case_			(      self		:       int				,       _snake_case		:       bool				):
          __lowercase		:			Dict	    =       self.get_dummy_dataset()
          __lowercase		:			List[Any]	    =       RagConfig(
              retrieval_vector_size=self.retrieval_vector_size				,       question_encoder=DPRConfig().to_dict()				,       generator=BartConfig().to_dict()				,       index_name='''custom'''				,       )
          if from_disk:
              __lowercase		:			List[Any]	    =       os.path.join(self.tmpdirname				,       '''dataset'''				)
              __lowercase		:			Union[str, Any]	    =       os.path.join(self.tmpdirname				,       '''index.faiss'''				)
              dataset.get_index('''embeddings'''				).save(os.path.join(self.tmpdirname				,       '''index.faiss'''				)				)
              dataset.drop_index('''embeddings'''				)
              dataset.save_to_disk(os.path.join(self.tmpdirname				,       '''dataset'''				)				)
              del dataset
              __lowercase		:			Optional[Any]	    =       RagRetriever(
                  _snake_case				,       question_encoder_tokenizer=self.get_dpr_tokenizer()				,       generator_tokenizer=self.get_bart_tokenizer()				,       )
          else:
              __lowercase		:			Optional[Any]	    =       RagRetriever(
                  _snake_case				,       question_encoder_tokenizer=self.get_dpr_tokenizer()				,       generator_tokenizer=self.get_bart_tokenizer()				,       index=CustomHFIndex(config.retrieval_vector_size				,       _snake_case				)				,       )
          return retriever




      def 		snake_case_			(      self		:       str				):
          __lowercase		:			List[str]	    =       Dataset.from_dict(
              {
                  '''id''': ['''0''', '''1'''],
                  '''text''': ['''foo''', '''bar'''],
                  '''title''': ['''Foo''', '''Bar'''],
                  '''embeddings''': [np.ones(self.retrieval_vector_size + 1				), 2 * np.ones(self.retrieval_vector_size + 1				)],
              }				)
          dataset.add_faiss_index('''embeddings'''				,       string_factory='''Flat'''				,       metric_type=faiss.METRIC_INNER_PRODUCT				)

          __lowercase		:			Optional[int]	    =       os.path.join(self.tmpdirname				,       '''hf_bert_base.hnswSQ8_correct_phi_128.c_index'''				)
          dataset.save_faiss_index('''embeddings'''				,       index_file_name + '''.index.dpr'''				)
          pickle.dump(dataset['''id''']				,       open(index_file_name + '''.index_meta.dpr'''				,       '''wb'''				)				)

          __lowercase		:			str	    =       os.path.join(self.tmpdirname				,       '''psgs_w100.tsv.pkl'''				)
          __lowercase		:			Dict	    =       {sample['''id''']: [sample['''text'''], sample['''title''']] for sample in dataset}
          pickle.dump(_snake_case				,       open(_snake_case				,       '''wb'''				)				)

          __lowercase		:			List[str]	    =       RagConfig(
              retrieval_vector_size=self.retrieval_vector_size				,       question_encoder=DPRConfig().to_dict()				,       generator=BartConfig().to_dict()				,       index_name='''legacy'''				,       index_path=self.tmpdirname				,       )
          __lowercase		:			Optional[int]	    =       RagRetriever(
              _snake_case				,       question_encoder_tokenizer=self.get_dpr_tokenizer()				,       generator_tokenizer=self.get_bart_tokenizer()				)
          return retriever




      def 		snake_case_			(      self		:       Optional[Any]				):
          __lowercase		:			List[str]	    =       1
          __lowercase		:			Tuple	    =       self.get_dummy_canonical_hf_index_retriever()
          __lowercase		:			List[Any]	    =       np.array(
              [np.ones(self.retrieval_vector_size				), -np.ones(self.retrieval_vector_size				)]				,       dtype=np.floataa				)
          __lowercase    ,      __lowercase    ,      __lowercase		:			str	    =       retriever.retrieve(_snake_case				,       n_docs=_snake_case				)
          self.assertEqual(retrieved_doc_embeds.shape				,       (2, n_docs, self.retrieval_vector_size)				)
          self.assertEqual(len(_snake_case				)				,       2				)
          self.assertEqual(sorted(doc_dicts[0]				)				,       ['''embeddings''', '''id''', '''text''', '''title''']				)
          self.assertEqual(len(doc_dicts[0]['''id''']				)				,       _snake_case				)
          self.assertEqual(doc_dicts[0]['''id'''][0]				,       '''1'''				)  # max inner product is reached with second doc
          self.assertEqual(doc_dicts[1]['''id'''][0]				,       '''0'''				)  # max inner product is reached with first doc
          self.assertListEqual(doc_ids.tolist()				,       [[1], [0]]				)




      def 		snake_case_			(      self		:       int				):
          __lowercase		:			int	    =       self.get_dummy_canonical_hf_index_retriever()
          with tempfile.TemporaryDirectory() as tmp_dirname:
              with patch('''transformers.models.rag.retrieval_rag.load_dataset'''				) as mock_load_dataset:
                  __lowercase		:			Optional[Any]	    =       self.get_dummy_dataset()
                  retriever.save_pretrained(_snake_case				)
                  __lowercase		:			str	    =       RagRetriever.from_pretrained(_snake_case				)
                  self.assertIsInstance(_snake_case				,       _snake_case				)
                  __lowercase		:			List[str]	    =       np.array(
                      [np.ones(self.retrieval_vector_size				), -np.ones(self.retrieval_vector_size				)]				,       dtype=np.floataa				)
                  __lowercase		:			Union[str, Any]	    =       retriever.retrieve(_snake_case				,       n_docs=1				)
                  self.assertTrue(out is not None				)




      def 		snake_case_			(      self		:       str				):
          __lowercase		:			List[str]	    =       1
          __lowercase		:			Dict	    =       self.get_dummy_custom_hf_index_retriever(from_disk=_snake_case				)
          __lowercase		:			Optional[Any]	    =       np.array(
              [np.ones(self.retrieval_vector_size				), -np.ones(self.retrieval_vector_size				)]				,       dtype=np.floataa				)
          __lowercase    ,      __lowercase    ,      __lowercase		:			Tuple	    =       retriever.retrieve(_snake_case				,       n_docs=_snake_case				)
          self.assertEqual(retrieved_doc_embeds.shape				,       (2, n_docs, self.retrieval_vector_size)				)
          self.assertEqual(len(_snake_case				)				,       2				)
          self.assertEqual(sorted(doc_dicts[0]				)				,       ['''embeddings''', '''id''', '''text''', '''title''']				)
          self.assertEqual(len(doc_dicts[0]['''id''']				)				,       _snake_case				)
          self.assertEqual(doc_dicts[0]['''id'''][0]				,       '''1'''				)  # max inner product is reached with second doc
          self.assertEqual(doc_dicts[1]['''id'''][0]				,       '''0'''				)  # max inner product is reached with first doc
          self.assertListEqual(doc_ids.tolist()				,       [[1], [0]]				)




      def 		snake_case_			(      self		:       Any				):
          __lowercase		:			List[str]	    =       self.get_dummy_custom_hf_index_retriever(from_disk=_snake_case				)
          with tempfile.TemporaryDirectory() as tmp_dirname:
              retriever.save_pretrained(_snake_case				)
              __lowercase		:			Optional[Any]	    =       RagRetriever.from_pretrained(_snake_case				)
              self.assertIsInstance(_snake_case				,       _snake_case				)
              __lowercase		:			Optional[Any]	    =       np.array(
                  [np.ones(self.retrieval_vector_size				), -np.ones(self.retrieval_vector_size				)]				,       dtype=np.floataa				)
              __lowercase		:			List[Any]	    =       retriever.retrieve(_snake_case				,       n_docs=1				)
              self.assertTrue(out is not None				)




      def 		snake_case_			(      self		:       List[Any]				):
          __lowercase		:			Any	    =       1
          __lowercase		:			List[Any]	    =       self.get_dummy_custom_hf_index_retriever(from_disk=_snake_case				)
          __lowercase		:			Any	    =       np.array(
              [np.ones(self.retrieval_vector_size				), -np.ones(self.retrieval_vector_size				)]				,       dtype=np.floataa				)
          __lowercase    ,      __lowercase    ,      __lowercase		:			Union[str, Any]	    =       retriever.retrieve(_snake_case				,       n_docs=_snake_case				)
          self.assertEqual(retrieved_doc_embeds.shape				,       (2, n_docs, self.retrieval_vector_size)				)
          self.assertEqual(len(_snake_case				)				,       2				)
          self.assertEqual(sorted(doc_dicts[0]				)				,       ['''embeddings''', '''id''', '''text''', '''title''']				)
          self.assertEqual(len(doc_dicts[0]['''id''']				)				,       _snake_case				)
          self.assertEqual(doc_dicts[0]['''id'''][0]				,       '''1'''				)  # max inner product is reached with second doc
          self.assertEqual(doc_dicts[1]['''id'''][0]				,       '''0'''				)  # max inner product is reached with first doc
          self.assertListEqual(doc_ids.tolist()				,       [[1], [0]]				)




      def 		snake_case_			(      self		:       Any				):
          __lowercase		:			str	    =       self.get_dummy_custom_hf_index_retriever(from_disk=_snake_case				)
          with tempfile.TemporaryDirectory() as tmp_dirname:
              retriever.save_pretrained(_snake_case				)
              __lowercase		:			Optional[int]	    =       RagRetriever.from_pretrained(_snake_case				)
              self.assertIsInstance(_snake_case				,       _snake_case				)
              __lowercase		:			Union[str, Any]	    =       np.array(
                  [np.ones(self.retrieval_vector_size				), -np.ones(self.retrieval_vector_size				)]				,       dtype=np.floataa				)
              __lowercase		:			Tuple	    =       retriever.retrieve(_snake_case				,       n_docs=1				)
              self.assertTrue(out is not None				)




      def 		snake_case_			(      self		:       Tuple				):
          __lowercase		:			Optional[int]	    =       1
          __lowercase		:			str	    =       self.get_dummy_legacy_index_retriever()
          __lowercase		:			Tuple	    =       np.array(
              [np.ones(self.retrieval_vector_size				), -np.ones(self.retrieval_vector_size				)]				,       dtype=np.floataa				)
          __lowercase    ,      __lowercase    ,      __lowercase		:			Tuple	    =       retriever.retrieve(_snake_case				,       n_docs=_snake_case				)
          self.assertEqual(retrieved_doc_embeds.shape				,       (2, n_docs, self.retrieval_vector_size)				)
          self.assertEqual(len(_snake_case				)				,       2				)
          self.assertEqual(sorted(doc_dicts[0]				)				,       ['''text''', '''title''']				)
          self.assertEqual(len(doc_dicts[0]['''text''']				)				,       _snake_case				)
          self.assertEqual(doc_dicts[0]['''text'''][0]				,       '''bar'''				)  # max inner product is reached with second doc
          self.assertEqual(doc_dicts[1]['''text'''][0]				,       '''foo'''				)  # max inner product is reached with first doc
          self.assertListEqual(doc_ids.tolist()				,       [[1], [0]]				)




      def 		snake_case_			(      self		:       Union[str, Any]				):
          __lowercase		:			Tuple	    =       self.get_dummy_legacy_index_retriever()
          with tempfile.TemporaryDirectory() as tmp_dirname:
              retriever.save_pretrained(_snake_case				)
              __lowercase		:			Tuple	    =       RagRetriever.from_pretrained(_snake_case				)
              self.assertIsInstance(_snake_case				,       _snake_case				)
              __lowercase		:			Optional[int]	    =       np.array(
                  [np.ones(self.retrieval_vector_size				), -np.ones(self.retrieval_vector_size				)]				,       dtype=np.floataa				)
              __lowercase		:			Tuple	    =       retriever.retrieve(_snake_case				,       n_docs=1				)
              self.assertTrue(out is not None				)




      @require_torch
      @require_tokenizers
      @require_sentencepiece
      def 		snake_case_			(      self		:       Optional[Any]				):
          import torch

          __lowercase		:			Tuple	    =       1
          __lowercase		:			Any	    =       self.get_dummy_canonical_hf_index_retriever()
          __lowercase		:			str	    =       [[5, 7], [10, 11]]
          __lowercase		:			List[Any]	    =       np.array(
              [np.ones(self.retrieval_vector_size				), -np.ones(self.retrieval_vector_size				)]				,       dtype=np.floataa				)
          __lowercase		:			Any	    =       retriever(_snake_case				,       _snake_case				,       prefix=retriever.config.generator.prefix				,       n_docs=_snake_case				)
          __lowercase    ,      __lowercase    ,      __lowercase		:			Tuple	    =       (
              out['''context_input_ids'''],
              out['''context_attention_mask'''],
              out['''retrieved_doc_embeds'''],
          )
          self.assertEqual(retrieved_doc_embeds.shape				,       (2, n_docs, self.retrieval_vector_size)				)
          self.assertIsInstance(_snake_case				,       _snake_case				)
          self.assertIsInstance(_snake_case				,       _snake_case				)
          self.assertIsInstance(_snake_case				,       np.ndarray				)

          __lowercase		:			Optional[Any]	    =       retriever(
              _snake_case				,       _snake_case				,       prefix=retriever.config.generator.prefix				,       n_docs=_snake_case				,       return_tensors='''pt'''				,       )
          __lowercase    ,      __lowercase    ,      __lowercase    ,      __lowercase		:			Optional[Any]	    =       (  # noqa: F841
              out['''context_input_ids'''],
              out['''context_attention_mask'''],
              out['''retrieved_doc_embeds'''],
              out['''doc_ids'''],
          )
          self.assertEqual(retrieved_doc_embeds.shape				,       (2, n_docs, self.retrieval_vector_size)				)
          self.assertIsInstance(_snake_case				,       torch.Tensor				)
          self.assertIsInstance(_snake_case				,       torch.Tensor				)
          self.assertIsInstance(_snake_case				,       torch.Tensor				)




      @require_torch
      @require_tokenizers
      @require_sentencepiece
      def 		snake_case_			(      self		:       List[Any]				):
          __lowercase		:			Tuple	    =       self.get_dpr_ctx_encoder_tokenizer()
          __lowercase		:			str	    =       1
          __lowercase		:			int	    =       self.get_dummy_custom_hf_index_retriever(from_disk=_snake_case				)
          retriever.set_ctx_encoder_tokenizer(_snake_case				)

          __lowercase		:			Tuple	    =       [[5, 7], [10, 11]]
          __lowercase		:			int	    =       np.array(
              [np.ones(self.retrieval_vector_size				), -np.ones(self.retrieval_vector_size				)]				,       dtype=np.floataa				)
          __lowercase		:			Any	    =       retriever(_snake_case				,       _snake_case				,       prefix=retriever.config.generator.prefix				,       n_docs=_snake_case				)

          self.assertEqual(
              len(_snake_case				)				,       6				)  # check whether the retriever output consist of 6 attributes including tokenized docs
          self.assertEqual(
              all(k in out for k in ('''tokenized_doc_ids''', '''tokenized_doc_attention_mask''')				)				,       _snake_case				)  # check for doc token related keys in dictionary.


