



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

import sentencepiece as spm

from ...tokenization_utils import AddedToken, PreTrainedTokenizer
from ...utils import logging


a_			=					logging.get_logger(__name__)
a_			=					{"""vocab_file""": """spiece.model"""}

a_			=					{
    """vocab_file""": {
        """albert-base-v1""": """https://huggingface.co/albert-base-v1/resolve/main/spiece.model""",
        """albert-large-v1""": """https://huggingface.co/albert-large-v1/resolve/main/spiece.model""",
        """albert-xlarge-v1""": """https://huggingface.co/albert-xlarge-v1/resolve/main/spiece.model""",
        """albert-xxlarge-v1""": """https://huggingface.co/albert-xxlarge-v1/resolve/main/spiece.model""",
        """albert-base-v2""": """https://huggingface.co/albert-base-v2/resolve/main/spiece.model""",
        """albert-large-v2""": """https://huggingface.co/albert-large-v2/resolve/main/spiece.model""",
        """albert-xlarge-v2""": """https://huggingface.co/albert-xlarge-v2/resolve/main/spiece.model""",
        """albert-xxlarge-v2""": """https://huggingface.co/albert-xxlarge-v2/resolve/main/spiece.model""",
    }
}

a_			=					{
    """albert-base-v1""": 512,
    """albert-large-v1""": 512,
    """albert-xlarge-v1""": 512,
    """albert-xxlarge-v1""": 512,
    """albert-base-v2""": 512,
    """albert-large-v2""": 512,
    """albert-xlarge-v2""": 512,
    """albert-xxlarge-v2""": 512,
}

a_			=					"""▁"""



class  __lowerCAmelCase     (	lowerCAmelCase__   ):
 lowerCAmelCase__       					=  VOCAB_FILES_NAMES
 lowerCAmelCase__       					=  PRETRAINED_VOCAB_FILES_MAP
 lowerCAmelCase__       					=  PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES




 def __init__(						self	,				__UpperCAmelCase	,				__UpperCAmelCase=True	,				__UpperCAmelCase=True	,				__UpperCAmelCase=False	,				__UpperCAmelCase="[CLS]"	,				__UpperCAmelCase="[SEP]"	,				__UpperCAmelCase="<unk>"	,				__UpperCAmelCase="[SEP]"	,				__UpperCAmelCase="<pad>"	,				__UpperCAmelCase="[CLS]"	,				__UpperCAmelCase="[MASK]"	,				__UpperCAmelCase = None	,				**__UpperCAmelCase	,				):



    '''simple docstring'''
    # Mask token behave like a normal word, i.e. include the space before it and
    # is included in the raw text, there should be a match in a non-normalized sentence.
    __lowerCamelCase													=    (
        AddedToken(__UpperCAmelCase	,				lstrip=__UpperCAmelCase	,				rstrip=__UpperCAmelCase	,				normalized=__UpperCAmelCase				)
        if isinstance(__UpperCAmelCase	,				__UpperCAmelCase				)
        else mask_token
    )

    __lowerCamelCase													=    {} if sp_model_kwargs is None else sp_model_kwargs

    super().__init__(
        do_lower_case=__UpperCAmelCase	,				remove_space=__UpperCAmelCase	,				keep_accents=__UpperCAmelCase	,				bos_token=__UpperCAmelCase	,				eos_token=__UpperCAmelCase	,				unk_token=__UpperCAmelCase	,				sep_token=__UpperCAmelCase	,				pad_token=__UpperCAmelCase	,				cls_token=__UpperCAmelCase	,				mask_token=__UpperCAmelCase	,				sp_model_kwargs=self.sp_model_kwargs	,				**__UpperCAmelCase	,				)

    __lowerCamelCase													=    do_lower_case
    __lowerCamelCase													=    remove_space
    __lowerCamelCase													=    keep_accents
    __lowerCamelCase													=    vocab_file

    __lowerCamelCase													=    spm.SentencePieceProcessor(**self.sp_model_kwargs				)
    self.sp_model.Load(__UpperCAmelCase				)




 @property
 def        lowerCamelCase      (						self				):



    '''simple docstring'''
    return len(self.sp_model				)




 def        lowerCamelCase      (						self				):



    '''simple docstring'''
    __lowerCamelCase													=    {self.convert_ids_to_tokens(__UpperCAmelCase				): i for i in range(self.vocab_size				)}
    vocab.update(self.added_tokens_encoder				)
    return vocab




 def __getstate__(						self				):



    '''simple docstring'''
    __lowerCamelCase													=    self.__dict__.copy()
    __lowerCamelCase													=    None
    return state




 def __setstate__(						self	,				__UpperCAmelCase				):



    '''simple docstring'''
    __lowerCamelCase													=    d

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

    __lowerCamelCase													=    spm.SentencePieceProcessor(**self.sp_model_kwargs				)
    self.sp_model.Load(self.vocab_file				)




 def        lowerCamelCase      (						self	,				__UpperCAmelCase				):



    '''simple docstring'''
    if self.remove_space:
       __lowerCamelCase													=    ''' '''.join(inputs.strip().split()				)
    else:
       __lowerCamelCase													=    inputs
    __lowerCamelCase													=    outputs.replace('''``'''	,				'''"'''				).replace('''\'\''''	,				'''"'''				)

    if not self.keep_accents:
       __lowerCamelCase													=    unicodedata.normalize('''NFKD'''	,				__UpperCAmelCase				)
       __lowerCamelCase													=    ''''''.join([c for c in outputs if not unicodedata.combining(__UpperCAmelCase				)]				)
    if self.do_lower_case:
       __lowerCamelCase													=    outputs.lower()

    return outputs




 def        lowerCamelCase      (						self	,				__UpperCAmelCase				):



    '''simple docstring'''
    __lowerCamelCase													=    self.preprocess_text(__UpperCAmelCase				)
    __lowerCamelCase													=    self.sp_model.encode(__UpperCAmelCase	,				out_type=__UpperCAmelCase				)
    __lowerCamelCase													=    []
    for piece in pieces:
       if len(__UpperCAmelCase				) > 1 and piece[-1] == str(''','''				) and piece[-2].isdigit():
          __lowerCamelCase													=    self.sp_model.EncodeAsPieces(piece[:-1].replace(__UpperCAmelCase	,				''''''				)				)
          if piece[0] != SPIECE_UNDERLINE and cur_pieces[0][0] == SPIECE_UNDERLINE:
             if len(cur_pieces[0]				) == 1:
                __lowerCamelCase													=    cur_pieces[1:]
             else:
                __lowerCamelCase													=    cur_pieces[0][1:]
          cur_pieces.append(piece[-1]				)
          new_pieces.extend(__UpperCAmelCase				)
       else:
          new_pieces.append(__UpperCAmelCase				)

    return new_pieces




 def        lowerCamelCase      (						self	,				__UpperCAmelCase				):



    '''simple docstring'''
    return self.sp_model.PieceToId(__UpperCAmelCase				)




 def        lowerCamelCase      (						self	,				__UpperCAmelCase				):



    '''simple docstring'''
    return self.sp_model.IdToPiece(__UpperCAmelCase				)




 def        lowerCamelCase      (						self	,				__UpperCAmelCase				):



    '''simple docstring'''
    __lowerCamelCase													=    []
    __lowerCamelCase													=    ''''''
    __lowerCamelCase													=    False
    for token in tokens:
       # make sure that special tokens are not decoded using sentencepiece model
       if token in self.all_special_tokens:
          if not prev_is_special:
             out_string += " "
          out_string += self.sp_model.decode(__UpperCAmelCase				) + token
          __lowerCamelCase													=    True
          __lowerCamelCase													=    []
       else:
          current_sub_tokens.append(__UpperCAmelCase				)
          __lowerCamelCase													=    False
    out_string += self.sp_model.decode(__UpperCAmelCase				)
    return out_string.strip()




 def        lowerCamelCase      (						self	,				__UpperCAmelCase	,				__UpperCAmelCase = 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        lowerCamelCase      (						self	,				__UpperCAmelCase	,				__UpperCAmelCase = None	,				__UpperCAmelCase = False				):



    '''simple docstring'''

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

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




 def        lowerCamelCase      (						self	,				__UpperCAmelCase	,				__UpperCAmelCase = 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        lowerCamelCase      (						self	,				__UpperCAmelCase	,				__UpperCAmelCase = None				):



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

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

    return (out_vocab_file,)











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




 def        lowerCamelCase      (						self				):



    '''simple docstring'''
    __lowerCamelCase													=    tempfile.mkdtemp()
    __lowerCamelCase													=    8

    # DPR tok
    __lowerCamelCase													=    [
        '''[UNK]''',
        '''[CLS]''',
        '''[SEP]''',
        '''[PAD]''',
        '''[MASK]''',
        '''want''',
        '''##want''',
        '''##ed''',
        '''wa''',
        '''un''',
        '''runn''',
        '''##ing''',
        ''',''',
        '''low''',
        '''lowest''',
    ]
    __lowerCamelCase													=    os.path.join(self.tmpdirname	,				'''dpr_tokenizer'''				)
    os.makedirs(__UpperCAmelCase	,				exist_ok=__UpperCAmelCase				)
    __lowerCamelCase													=    os.path.join(__UpperCAmelCase	,				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
    __lowerCamelCase													=    [
        '''l''',
        '''o''',
        '''w''',
        '''e''',
        '''r''',
        '''s''',
        '''t''',
        '''i''',
        '''d''',
        '''n''',
        '''\u0120''',
        '''\u0120l''',
        '''\u0120n''',
        '''\u0120lo''',
        '''\u0120low''',
        '''er''',
        '''\u0120lowest''',
        '''\u0120newer''',
        '''\u0120wider''',
        '''<unk>''',
    ]
    __lowerCamelCase													=    dict(zip(__UpperCAmelCase	,				range(len(__UpperCAmelCase				)				)				)				)
    __lowerCamelCase													=    ['''#version: 0.2''', '''\u0120 l''', '''\u0120l o''', '''\u0120lo w''', '''e r''', '''''']
    __lowerCamelCase													=    {'''unk_token''': '''<unk>'''}

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




 def        lowerCamelCase      (						self				):



    '''simple docstring'''
    return DPRQuestionEncoderTokenizer.from_pretrained(os.path.join(self.tmpdirname	,				'''dpr_tokenizer'''				)				)




 def        lowerCamelCase      (						self				):



    '''simple docstring'''
    return DPRContextEncoderTokenizer.from_pretrained(os.path.join(self.tmpdirname	,				'''dpr_tokenizer'''				)				)




 def        lowerCamelCase      (						self				):



    '''simple docstring'''
    return BartTokenizer.from_pretrained(os.path.join(self.tmpdirname	,				'''bart_tokenizer'''				)				)




 def        lowerCamelCase      (						self				):



    '''simple docstring'''
    shutil.rmtree(self.tmpdirname				)




 def        lowerCamelCase      (						self				):



    '''simple docstring'''
    __lowerCamelCase													=    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        lowerCamelCase      (						self				):



    '''simple docstring'''
    __lowerCamelCase													=    self.get_dummy_dataset()
    __lowerCamelCase													=    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:
       __lowerCamelCase													=    dataset
       __lowerCamelCase													=    RagRetriever(
           __UpperCAmelCase	,				question_encoder_tokenizer=self.get_dpr_tokenizer()	,				generator_tokenizer=self.get_bart_tokenizer()	,				)
    return retriever




 def        lowerCamelCase      (						self	,				__UpperCAmelCase				):



    '''simple docstring'''
    __lowerCamelCase													=    self.get_dummy_dataset()
    __lowerCamelCase													=    RagConfig(
        retrieval_vector_size=self.retrieval_vector_size	,				question_encoder=DPRConfig().to_dict()	,				generator=BartConfig().to_dict()	,				index_name='''custom'''	,				)
    if from_disk:
       __lowerCamelCase													=    os.path.join(self.tmpdirname	,				'''dataset'''				)
       __lowerCamelCase													=    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
       __lowerCamelCase													=    RagRetriever(
           __UpperCAmelCase	,				question_encoder_tokenizer=self.get_dpr_tokenizer()	,				generator_tokenizer=self.get_bart_tokenizer()	,				)
    else:
       __lowerCamelCase													=    RagRetriever(
           __UpperCAmelCase	,				question_encoder_tokenizer=self.get_dpr_tokenizer()	,				generator_tokenizer=self.get_bart_tokenizer()	,				index=CustomHFIndex(config.retrieval_vector_size	,				__UpperCAmelCase				)	,				)
    return retriever




 def        lowerCamelCase      (						self				):



    '''simple docstring'''
    __lowerCamelCase													=    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				)

    __lowerCamelCase													=    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'''				)				)

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

    __lowerCamelCase													=    RagConfig(
        retrieval_vector_size=self.retrieval_vector_size	,				question_encoder=DPRConfig().to_dict()	,				generator=BartConfig().to_dict()	,				index_name='''legacy'''	,				index_path=self.tmpdirname	,				)
    __lowerCamelCase													=    RagRetriever(
        __UpperCAmelCase	,				question_encoder_tokenizer=self.get_dpr_tokenizer()	,				generator_tokenizer=self.get_bart_tokenizer()				)
    return retriever




 def        lowerCamelCase      (						self				):



    '''simple docstring'''
    __lowerCamelCase													=    1
    __lowerCamelCase													=    self.get_dummy_canonical_hf_index_retriever()
    __lowerCamelCase													=    np.array(
        [np.ones(self.retrieval_vector_size				), -np.ones(self.retrieval_vector_size				)]	,				dtype=np.floataa				)
    __lowerCamelCase							,__lowerCamelCase							,__lowerCamelCase													=    retriever.retrieve(__UpperCAmelCase	,				n_docs=__UpperCAmelCase				)
    self.assertEqual(retrieved_doc_embeds.shape	,				(2, n_docs, self.retrieval_vector_size)				)
    self.assertEqual(len(__UpperCAmelCase				)	,				2				)
    self.assertEqual(sorted(doc_dicts[0]				)	,				['''embeddings''', '''id''', '''text''', '''title''']				)
    self.assertEqual(len(doc_dicts[0]['''id''']				)	,				__UpperCAmelCase				)
    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        lowerCamelCase      (						self				):



    '''simple docstring'''
    __lowerCamelCase													=    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:
          __lowerCamelCase													=    self.get_dummy_dataset()
          retriever.save_pretrained(__UpperCAmelCase				)
          __lowerCamelCase													=    RagRetriever.from_pretrained(__UpperCAmelCase				)
          self.assertIsInstance(__UpperCAmelCase	,				__UpperCAmelCase				)
          __lowerCamelCase													=    np.array(
              [np.ones(self.retrieval_vector_size				), -np.ones(self.retrieval_vector_size				)]	,				dtype=np.floataa				)
          __lowerCamelCase													=    retriever.retrieve(__UpperCAmelCase	,				n_docs=1				)
          self.assertTrue(out is not None				)




 def        lowerCamelCase      (						self				):



    '''simple docstring'''
    __lowerCamelCase													=    1
    __lowerCamelCase													=    self.get_dummy_custom_hf_index_retriever(from_disk=__UpperCAmelCase				)
    __lowerCamelCase													=    np.array(
        [np.ones(self.retrieval_vector_size				), -np.ones(self.retrieval_vector_size				)]	,				dtype=np.floataa				)
    __lowerCamelCase							,__lowerCamelCase							,__lowerCamelCase													=    retriever.retrieve(__UpperCAmelCase	,				n_docs=__UpperCAmelCase				)
    self.assertEqual(retrieved_doc_embeds.shape	,				(2, n_docs, self.retrieval_vector_size)				)
    self.assertEqual(len(__UpperCAmelCase				)	,				2				)
    self.assertEqual(sorted(doc_dicts[0]				)	,				['''embeddings''', '''id''', '''text''', '''title''']				)
    self.assertEqual(len(doc_dicts[0]['''id''']				)	,				__UpperCAmelCase				)
    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        lowerCamelCase      (						self				):



    '''simple docstring'''
    __lowerCamelCase													=    self.get_dummy_custom_hf_index_retriever(from_disk=__UpperCAmelCase				)
    with tempfile.TemporaryDirectory() as tmp_dirname:
       retriever.save_pretrained(__UpperCAmelCase				)
       __lowerCamelCase													=    RagRetriever.from_pretrained(__UpperCAmelCase				)
       self.assertIsInstance(__UpperCAmelCase	,				__UpperCAmelCase				)
       __lowerCamelCase													=    np.array(
           [np.ones(self.retrieval_vector_size				), -np.ones(self.retrieval_vector_size				)]	,				dtype=np.floataa				)
       __lowerCamelCase													=    retriever.retrieve(__UpperCAmelCase	,				n_docs=1				)
       self.assertTrue(out is not None				)




 def        lowerCamelCase      (						self				):



    '''simple docstring'''
    __lowerCamelCase													=    1
    __lowerCamelCase													=    self.get_dummy_custom_hf_index_retriever(from_disk=__UpperCAmelCase				)
    __lowerCamelCase													=    np.array(
        [np.ones(self.retrieval_vector_size				), -np.ones(self.retrieval_vector_size				)]	,				dtype=np.floataa				)
    __lowerCamelCase							,__lowerCamelCase							,__lowerCamelCase													=    retriever.retrieve(__UpperCAmelCase	,				n_docs=__UpperCAmelCase				)
    self.assertEqual(retrieved_doc_embeds.shape	,				(2, n_docs, self.retrieval_vector_size)				)
    self.assertEqual(len(__UpperCAmelCase				)	,				2				)
    self.assertEqual(sorted(doc_dicts[0]				)	,				['''embeddings''', '''id''', '''text''', '''title''']				)
    self.assertEqual(len(doc_dicts[0]['''id''']				)	,				__UpperCAmelCase				)
    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        lowerCamelCase      (						self				):



    '''simple docstring'''
    __lowerCamelCase													=    self.get_dummy_custom_hf_index_retriever(from_disk=__UpperCAmelCase				)
    with tempfile.TemporaryDirectory() as tmp_dirname:
       retriever.save_pretrained(__UpperCAmelCase				)
       __lowerCamelCase													=    RagRetriever.from_pretrained(__UpperCAmelCase				)
       self.assertIsInstance(__UpperCAmelCase	,				__UpperCAmelCase				)
       __lowerCamelCase													=    np.array(
           [np.ones(self.retrieval_vector_size				), -np.ones(self.retrieval_vector_size				)]	,				dtype=np.floataa				)
       __lowerCamelCase													=    retriever.retrieve(__UpperCAmelCase	,				n_docs=1				)
       self.assertTrue(out is not None				)




 def        lowerCamelCase      (						self				):



    '''simple docstring'''
    __lowerCamelCase													=    1
    __lowerCamelCase													=    self.get_dummy_legacy_index_retriever()
    __lowerCamelCase													=    np.array(
        [np.ones(self.retrieval_vector_size				), -np.ones(self.retrieval_vector_size				)]	,				dtype=np.floataa				)
    __lowerCamelCase							,__lowerCamelCase							,__lowerCamelCase													=    retriever.retrieve(__UpperCAmelCase	,				n_docs=__UpperCAmelCase				)
    self.assertEqual(retrieved_doc_embeds.shape	,				(2, n_docs, self.retrieval_vector_size)				)
    self.assertEqual(len(__UpperCAmelCase				)	,				2				)
    self.assertEqual(sorted(doc_dicts[0]				)	,				['''text''', '''title''']				)
    self.assertEqual(len(doc_dicts[0]['''text''']				)	,				__UpperCAmelCase				)
    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        lowerCamelCase      (						self				):



    '''simple docstring'''
    __lowerCamelCase													=    self.get_dummy_legacy_index_retriever()
    with tempfile.TemporaryDirectory() as tmp_dirname:
       retriever.save_pretrained(__UpperCAmelCase				)
       __lowerCamelCase													=    RagRetriever.from_pretrained(__UpperCAmelCase				)
       self.assertIsInstance(__UpperCAmelCase	,				__UpperCAmelCase				)
       __lowerCamelCase													=    np.array(
           [np.ones(self.retrieval_vector_size				), -np.ones(self.retrieval_vector_size				)]	,				dtype=np.floataa				)
       __lowerCamelCase													=    retriever.retrieve(__UpperCAmelCase	,				n_docs=1				)
       self.assertTrue(out is not None				)




 @require_torch
 @require_tokenizers
 @require_sentencepiece
 def        lowerCamelCase      (						self				):



    '''simple docstring'''
    import torch

    __lowerCamelCase													=    1
    __lowerCamelCase													=    self.get_dummy_canonical_hf_index_retriever()
    __lowerCamelCase													=    [[5, 7], [10, 11]]
    __lowerCamelCase													=    np.array(
        [np.ones(self.retrieval_vector_size				), -np.ones(self.retrieval_vector_size				)]	,				dtype=np.floataa				)
    __lowerCamelCase													=    retriever(__UpperCAmelCase	,				__UpperCAmelCase	,				prefix=retriever.config.generator.prefix	,				n_docs=__UpperCAmelCase				)
    __lowerCamelCase							,__lowerCamelCase							,__lowerCamelCase													=    (
        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(__UpperCAmelCase	,				__UpperCAmelCase				)
    self.assertIsInstance(__UpperCAmelCase	,				__UpperCAmelCase				)
    self.assertIsInstance(__UpperCAmelCase	,				np.ndarray				)

    __lowerCamelCase													=    retriever(
        __UpperCAmelCase	,				__UpperCAmelCase	,				prefix=retriever.config.generator.prefix	,				n_docs=__UpperCAmelCase	,				return_tensors='''pt'''	,				)
    __lowerCamelCase							,__lowerCamelCase							,__lowerCamelCase							,__lowerCamelCase													=    (  # 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(__UpperCAmelCase	,				torch.Tensor				)
    self.assertIsInstance(__UpperCAmelCase	,				torch.Tensor				)
    self.assertIsInstance(__UpperCAmelCase	,				torch.Tensor				)




 @require_torch
 @require_tokenizers
 @require_sentencepiece
 def        lowerCamelCase      (						self				):



    '''simple docstring'''
    __lowerCamelCase													=    self.get_dpr_ctx_encoder_tokenizer()
    __lowerCamelCase													=    1
    __lowerCamelCase													=    self.get_dummy_custom_hf_index_retriever(from_disk=__UpperCAmelCase				)
    retriever.set_ctx_encoder_tokenizer(__UpperCAmelCase				)

    __lowerCamelCase													=    [[5, 7], [10, 11]]
    __lowerCamelCase													=    np.array(
        [np.ones(self.retrieval_vector_size				), -np.ones(self.retrieval_vector_size				)]	,				dtype=np.floataa				)
    __lowerCamelCase													=    retriever(__UpperCAmelCase	,				__UpperCAmelCase	,				prefix=retriever.config.generator.prefix	,				n_docs=__UpperCAmelCase				)

    self.assertEqual(
        len(__UpperCAmelCase				)	,				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''')				)	,				__UpperCAmelCase				)  # check for doc token related keys in dictionary.






