codingutf8 2020 the rag s and the huggingface inc team licensed under the apache license version 2 0 the license you may not use this file except in compliance with the license you may obtain a copy of the license at http www apache orglicenseslicense2 0 unless required by applicable law or agreed to in writing software distributed under the license is distributed on an as is basis without warranties or conditions of any kind either express or implied see the license for the specific language governing permissions and limitations under the license rag retriever model implementation import os import pickle import time from typing import iterable list optional tuple import numpy as np from tokenizationutils import pretrainedtokenizer from tokenizationutilsbase import batchencoding from utils import cachedfile isdatasetsavailable isfaissavailable logging requiresbackends from configurationrag import ragconfig from tokenizationrag import ragtokenizer if isdatasetsavailable from datasets import dataset loaddataset loadfromdisk if isfaissavailable import faiss logger logging getloggername legacyindexpath https storage googleapis comhuggingfacenlpdatasetswikidpr class index def getdocdictsself docids np ndarray listdict raise notimplementederror def gettopdocsself questionhiddenstates np ndarray ndocs5 tuplenp ndarray np ndarray raise notimplementederror def isinitializedself raise notimplementederror def initindexself raise notimplementederror class legacyindexindex indexfilename hfbertbase hnswsq8correctphi128 cindex passagefilename psgsw100 tsv pkl def initself vectorsize indexpath self indexidtodbid self indexpath indexpath self passages self loadpassages self vectorsize vectorsize self index none self indexinitialized false def resolvepathself indexpath filename islocal os path isdirindexpath try load from url or cache if already cached resolvedarchivefile cachedfileindexpath filename except environmenterror msg fcan t load filename make sure that nn f indexpath is a correct remote path to a directory containing a file named filenamenn f or indexpath is the correct path to a directory containing a file named filename nn raise environmenterrormsg if islocal logger infofloading file resolvedarchivefile else logger infofloading file filename from cache at resolvedarchivefile return resolvedarchivefile def loadpassagesself logger infofloading passages from self indexpath passagespath self resolvepathself indexpath self passagefilename with openpassagespath rb as passagesfile passages pickle loadpassagesfile return passages def deserializeindexself logger infofloading index from self indexpath resolvedindexpath self resolvepathself indexpath self indexfilename index dpr self index faiss readindexresolvedindexpath resolvedmetapath self resolvepathself indexpath self indexfilename indexmeta dpr with openresolvedmetapath rb as metadatafile self indexidtodbid pickle loadmetadatafile assert lenself indexidtodbid self index ntotal deserialized indexidtodbid should match faiss index size def isinitializedself return self indexinitialized def initindexself index faiss indexhnswflatself vectorsize 1 512 index hnsw efsearch 128 index hnsw efconstruction 200 self index index self deserializeindex self indexinitialized true def getdocdictsself docids np array doclist for docidsi in docids ids strintdocid for docid in docidsi docs self passagesdocid for docid in ids doclist appenddocs docdicts for docs in doclist docdict docdicttitle doc1 for doc in docs docdicttext doc0 for doc in docs docdicts appenddocdict return docdicts def gettopdocsself questionhiddenstates np ndarray ndocs5 tuplenp ndarray np ndarray auxdim np zeroslenquestionhiddenstates dtypefloat32 reshape1 1 querynhswvectors np hstackquestionhiddenstates auxdim docsids self index searchquerynhswvectors ndocs vectors self index reconstructintdocid 1 for docid in docids for docids in docsids ids intself indexidtodbiddocid for docid in docids for docids in docsids return np arrayids np arrayvectors class hfindexbaseindex def initself vectorsize dataset indexinitializedfalse self vectorsize vectorsize self dataset dataset self indexinitialized indexinitialized self checkdatasetformatwithindexindexinitialized dataset setformatnumpy columnsembeddings outputallcolumnstrue dtypefloat32 def checkdatasetformatself withindex bool if not isinstanceself dataset dataset raise valueerrorfdataset should be a datasets dataset object but got typeself dataset if lentitle text embeddings setself dataset columnnames 0 raise valueerror dataset should be a dataset with the following columns title str text str and embeddings arrays of dimension vectorsize fbut got columns self dataset columnnames if withindex and embeddings not in self dataset listindexes raise valueerror missing faiss index in the dataset make sure you called dataset addfaissindex to compute it or dataset loadfaissindex to load one from the disk def initindexself raise notimplementederror def isinitializedself return self indexinitialized def getdocdictsself docids np ndarray listdict return self datasetdocidsi tolist for i in rangedocids shape0 def gettopdocsself questionhiddenstates np ndarray ndocs5 tuplenp ndarray np ndarray ids self dataset searchbatchembeddings questionhiddenstates ndocs docs self dataseti for i in indices if i 0 for indices in ids vectors docembeddings for doc in docs for i in rangelenvectors if lenvectorsi ndocs vectorsi np vstackvectorsi np zerosndocs lenvectorsi self vectorsize return np arrayids np arrayvectors shapes batchsize ndocs and batchsize ndocs d class canonicalhfindexhfindexbase def init self vectorsize int datasetname str wikidpr datasetsplit str train indexname optionalstr none indexpath optionalstr none usedummydatasetfalse if intindexpath is none intindexname is none 1 raise valueerrorplease provide indexname or indexpath self datasetname datasetname self datasetsplit datasetsplit self indexname indexname self indexpath indexpath self usedummydataset usedummydataset logger infofloading passages from self datasetname dataset loaddataset self datasetname withindexfalse splitself datasetsplit dummyself usedummydataset super initvectorsize dataset indexinitializedfalse def initindexself if self indexpath is not none logger infofloading index from self indexpath self dataset loadfaissindexembeddings fileself indexpath else logger infofloading index from self datasetname with index name self indexname self dataset loaddataset self datasetname withembeddingstrue withindextrue splitself datasetsplit indexnameself indexname dummyself usedummydataset self dataset setformatnumpy columnsembeddings outputallcolumnstrue self indexinitialized true class customhfindexhfindexbase def initself vectorsize int dataset indexpathnone super initvectorsize dataset indexinitializedindexpath is none self indexpath indexpath classmethod def loadfromdiskcls vectorsize datasetpath indexpath logger infofloading passages from datasetpath if datasetpath is none or indexpath is none raise valueerror please provide datasetpath and indexpath after calling dataset savetodiskdatasetpath and dataset getindex embeddings saveindexpath dataset loadfromdiskdatasetpath return clsvectorsizevectorsize datasetdataset indexpathindexpath def initindexself if not self isinitialized logger infofloading index from self indexpath self dataset loadfaissindexembeddings fileself indexpath self indexinitialized true class ragretriever datasets don t support savetodisk with indexes right now retriever initialization function it loads the index into memory def catinputanddocdoctitle doctext inputstring prefix todopatrick if we train more rag models i want to put the input first to take advantage of effortless truncation todopiktus better handling of truncation if doctitle startswith doctitle doctitle1 if doctitle endswith doctitle doctitle 1 if prefix is none prefix out prefix doctitle self config titlesep doctext self config docsep inputstring replace return out raginputstrings catinputanddoc docsititlej docsitextj inputstringsi prefix for i in rangelendocs for j in rangendocs contextualizedinputs self generatortokenizer batchencodeplus raginputstrings maxlengthself config maxcombinedlength returntensorsreturntensors paddingmaxlength truncationtrue return contextualizedinputsinputids contextualizedinputsattentionmask def chunktensorself t iterable chunksize int listiterable return ti i chunksize for i in range0 lent chunksize def mainretrieveself questionhiddenstates np ndarray ndocs int tuplenp ndarray np ndarray questionhiddenstatesbatched self chunktensorquestionhiddenstates self batchsize idsbatched vectorsbatched for questionhiddenstates in questionhiddenstatesbatched starttime time time ids vectors self index gettopdocsquestionhiddenstates ndocs logger debug findex search time time time starttime sec batch size questionhiddenstates shape idsbatched extendids vectorsbatched extendvectors return np arrayidsbatched np arrayvectorsbatched shapes batchsize ndocs and batchsize ndocs d def retrieveself questionhiddenstates np ndarray ndocs int tuplenp ndarray listdict docids retrieveddocembeds self mainretrievequestionhiddenstates ndocs return retrieveddocembeds docids self index getdocdictsdocids def setctxencodertokenizerself ctxencodertokenizer pretrainedtokenizer used in end2end retriever training self ctxencodertokenizer ctxencodertokenizer self returntokenizeddocs true def call self questioninputids listlistint questionhiddenstates np ndarray prefixnone ndocsnone returntensorsnone batchencoding ndocs ndocs if ndocs is not none else self ndocs prefix prefix if prefix is not none else self config generator prefix retrieveddocembeds docids docs self retrievequestionhiddenstates ndocs inputstrings self questionencodertokenizer batchdecodequestioninputids skipspecialtokenstrue contextinputids contextattentionmask self postprocessdocs docs inputstrings prefix ndocs returntensorsreturntensors if self returntokenizeddocs retrieveddoctext retrieveddoctitle for bidx in rangelendocs for docidx in rangendocs retrieveddoctext appenddocsbidxtextdocidx retrieveddoctitle appenddocsbidxtitledocidx tokenizeddocs self ctxencodertokenizer retrieveddoctitle retrieveddoctext truncationtrue paddinglongest returntensorsreturntensors return batchencoding contextinputids contextinputids contextattentionmask contextattentionmask retrieveddocembeds retrieveddocembeds docids docids tokenizeddocids tokenizeddocsinputids tokenizeddocattentionmask tokenizeddocsattentionmask tensortypereturntensors else return batchencoding contextinputids contextinputids contextattentionmask contextattentionmask retrieveddocembeds retrieveddocembeds docids docids tensortypereturntensors coding utf 8 2020 the rag s and the huggingface inc team licensed under the apache license version 2 0 the license you may not use this file except in compliance with the license you may obtain a copy of the license at http www apache org licenses license 2 0 unless required by applicable law or agreed to in writing software distributed under the license is distributed on an as is basis without warranties or conditions of any kind either express or implied see the license for the specific language governing permissions and limitations under the license rag retriever model implementation a base class for the indices encapsulated by the ragretriever returns a list of dictionaries containing titles and text of the retrieved documents args doc_ids np ndarray of shape batch_size n_docs a tensor of document indices for each query in the batch retrieves n_docs documents args question_hidden_states np ndarray of shape batch_size vector_size an array of query vectors n_docs int the number of docs retrieved per query returns np ndarray of shape batch_size n_docs a tensor of indices of retrieved documents np ndarray of shape batch_size vector_size a tensor of vector representations of retrieved documents returns true if index is already initialized a function responsible for loading the index into memory should be called only once per training run of a rag model e g if the model is trained on multiple gpus in a distributed setup only one of the workers will load the index an index which can be deserialized from the files built using https github com facebookresearch dpr we use default faiss index parameters as specified in that repository args vector_size int the dimension of indexed vectors index_path str a path to a directory containing index files compatible with models rag retrieval_rag legacyindex load from url or cache if already cached shapes batch_size n_docs and batch_size n_docs d a wrapper around an instance of datasets datasets if index_path is set to none we load the pre computed index available with the datasets arrow_dataset dataset otherwise we load the index from the indicated path on disk args vector_size int the dimension of the passages embeddings used by the index dataset_name str optional defaults to wiki_dpr a dataset identifier of the indexed dataset on huggingface aws bucket list all available datasets and ids with datasets list_datasets dataset_split str optional defaults to train which split of the dataset to load index_name str optional defaults to train the index_name of the index associated with the dataset the index loaded from index_path will be saved under this name index_path str optional defaults to none the path to the serialized faiss index on disk use_dummy_dataset bool optional defaults to false if true use the dummy configuration of the dataset for tests a wrapper around an instance of datasets datasets the dataset and the index are both loaded from the indicated paths on disk args vector_size int the dimension of the passages embeddings used by the index dataset_path str the path to the serialized dataset on disk the dataset should have 3 columns title str text str and embeddings arrays of dimension vector_size index_path str the path to the serialized faiss index on disk retriever used to get documents from vector queries it retrieves the documents embeddings as well as the documents contents and it formats them to be used with a ragmodel args config ragconfig the configuration of the rag model this retriever is used with contains parameters indicating which index to build you can load your own custom dataset with config index_name custom or use a canonical one default from the datasets library with config index_name wiki_dpr for example question_encoder_tokenizer pretrainedtokenizer the tokenizer that was used to tokenize the question it is used to decode the question and then use the generator_tokenizer generator_tokenizer pretrainedtokenizer the tokenizer used for the generator part of the ragmodel index models rag retrieval_rag index optional defaults to the one defined by the configuration if specified use this index instead of the one built using the configuration examples python to load the default wiki_dpr dataset with 21m passages from wikipedia index name is compressed or exact from transformers import ragretriever retriever ragretriever from_pretrained facebook dpr ctx_encoder single nq base dataset wiki_dpr index_name compressed to load your own indexed dataset built with the datasets library more info on how to build the indexed dataset in examples rag use_own_knowledge_dataset py from transformers import ragretriever dataset dataset must be a datasets datasets object with columns title text and embeddings and it must have a faiss index retriever ragretriever from_pretrained facebook dpr ctx_encoder single nq base indexed_dataset dataset to load your own indexed dataset built with the datasets library that was saved on disk more info in examples rag use_own_knowledge_dataset py from transformers import ragretriever dataset_path path to my dataset dataset saved via dataset save_to_disk index_path path to my index faiss faiss index saved via dataset get_index embeddings save retriever ragretriever from_pretrained facebook dpr ctx_encoder single nq base index_name custom passages_path dataset_path index_path index_path to load the legacy index built originally for rag s paper from transformers import ragretriever retriever ragretriever from_pretrained facebook dpr ctx_encoder single nq base index_name legacy datasets don t support save_to_disk with indexes right now retriever initialization function it loads the index into memory postprocessing retrieved docs and combining them with input_strings args docs dict retrieved documents input_strings str input strings decoded by preprocess_query prefix str prefix added at the beginning of each input typically used with t5 based models return tuple tensors a tuple consisting of two elements contextualized input_ids and a compatible attention_mask todo patrick if we train more rag models i want to put the input first to take advantage of effortless truncation todo piktus better handling of truncation shapes batch_size n_docs and batch_size n_docs d retrieves documents for specified question_hidden_states args question_hidden_states np ndarray of shape batch_size vector_size a batch of query vectors to retrieve with n_docs int the number of docs retrieved per query return tuple np ndarray np ndarray list dict a tuple with the following objects retrieved_doc_embeds np ndarray of shape batch_size n_docs dim the retrieval embeddings of the retrieved docs per query doc_ids np ndarray of shape batch_size n_docs the ids of the documents in the index doc_dicts list dict the retrieved_doc_embeds examples per query used in end2end retriever training retrieves documents for specified question_hidden_states args question_input_ids list list int batch of input ids question_hidden_states np ndarray of shape batch_size vector_size a batch of query vectors to retrieve with prefix str optional the prefix used by the generator s tokenizer n_docs int optional the number of docs retrieved per query return_tensors str or utils tensortype optional defaults to pt if set will return tensors instead of list of python integers acceptable values are tf return tensorflow tf constant objects pt return pytorch torch tensor objects np return numpy np ndarray objects returns batchencoding a batchencoding with the following fields context_input_ids list of token ids to be fed to a model what are input ids glossary input ids context_attention_mask list of indices specifying which tokens should be attended to by the model when return_attention_mask true or if attention_mask is in self model_input_names what are attention masks glossary attention mask retrieved_doc_embeds list of embeddings of the retrieved documents doc_ids list of ids of the retrieved documents
import os
import pickle
import time
from typing import Iterable, List, Optional, Tuple

import numpy as np

from ...tokenization_utils import PreTrainedTokenizer
from ...tokenization_utils_base import BatchEncoding
from ...utils import cached_file, is_datasets_available, is_faiss_available, logging, requires_backends
from .configuration_rag import RagConfig
from .tokenization_rag import RagTokenizer


if is_datasets_available():
    from datasets import Dataset, load_dataset, load_from_disk

if is_faiss_available():
    import faiss


logger = logging.get_logger(__name__)


LEGACY_INDEX_PATH = "https://storage.googleapis.com/huggingface-nlp/datasets/wiki_dpr/"


class Index:
    

    def get_doc_dicts(self, doc_ids: np.ndarray) -> List[dict]:
        
        raise NotImplementedError

    def get_top_docs(self, question_hidden_states: np.ndarray, n_docs=5) -> Tuple[np.ndarray, np.ndarray]:
        
        raise NotImplementedError

    def is_initialized(self):
        
        raise NotImplementedError

    def init_index(self):
        
        raise NotImplementedError


class LegacyIndex(Index):
    

    INDEX_FILENAME = "hf_bert_base.hnswSQ8_correct_phi_128.c_index"
    PASSAGE_FILENAME = "psgs_w100.tsv.pkl"

    def __init__(self, vector_size, index_path):
        self.index_id_to_db_id = []
        self.index_path = index_path
        self.passages = self._load_passages()
        self.vector_size = vector_size
        self.index = None
        self._index_initialized = False

    def _resolve_path(self, index_path, filename):
        is_local = os.path.isdir(index_path)
        try:
            
            resolved_archive_file = cached_file(index_path, filename)
        except EnvironmentError:
            msg = (
                f"Can't load '{filename}'. Make sure that:\n\n"
                f"- '{index_path}' is a correct remote path to a directory containing a file named {filename}\n\n"
                f"- or '{index_path}' is the correct path to a directory containing a file named {filename}.\n\n"
            )
            raise EnvironmentError(msg)
        if is_local:
            logger.info(f"loading file {resolved_archive_file}")
        else:
            logger.info(f"loading file {filename} from cache at {resolved_archive_file}")
        return resolved_archive_file

    def _load_passages(self):
        logger.info(f"Loading passages from {self.index_path}")
        passages_path = self._resolve_path(self.index_path, self.PASSAGE_FILENAME)
        with open(passages_path, "rb") as passages_file:
            passages = pickle.load(passages_file)
        return passages

    def _deserialize_index(self):
        logger.info(f"Loading index from {self.index_path}")
        resolved_index_path = self._resolve_path(self.index_path, self.INDEX_FILENAME + ".index.dpr")
        self.index = faiss.read_index(resolved_index_path)
        resolved_meta_path = self._resolve_path(self.index_path, self.INDEX_FILENAME + ".index_meta.dpr")
        with open(resolved_meta_path, "rb") as metadata_file:
            self.index_id_to_db_id = pickle.load(metadata_file)
        assert (
            len(self.index_id_to_db_id) == self.index.ntotal
        ), "Deserialized index_id_to_db_id should match faiss index size"

    def is_initialized(self):
        return self._index_initialized

    def init_index(self):
        index = faiss.IndexHNSWFlat(self.vector_size + 1, 512)
        index.hnsw.efSearch = 128
        index.hnsw.efConstruction = 200
        self.index = index
        self._deserialize_index()
        self._index_initialized = True

    def get_doc_dicts(self, doc_ids: np.array):
        doc_list = []
        for doc_ids_i in doc_ids:
            ids = [str(int(doc_id)) for doc_id in doc_ids_i]
            docs = [self.passages[doc_id] for doc_id in ids]
            doc_list.append(docs)
        doc_dicts = []
        for docs in doc_list:
            doc_dict = {}
            doc_dict["title"] = [doc[1] for doc in docs]
            doc_dict["text"] = [doc[0] for doc in docs]
            doc_dicts.append(doc_dict)
        return doc_dicts

    def get_top_docs(self, question_hidden_states: np.ndarray, n_docs=5) -> Tuple[np.ndarray, np.ndarray]:
        aux_dim = np.zeros(len(question_hidden_states), dtype="float32").reshape(-1, 1)
        query_nhsw_vectors = np.hstack((question_hidden_states, aux_dim))
        _, docs_ids = self.index.search(query_nhsw_vectors, n_docs)
        vectors = [[self.index.reconstruct(int(doc_id))[:-1] for doc_id in doc_ids] for doc_ids in docs_ids]
        ids = [[int(self.index_id_to_db_id[doc_id]) for doc_id in doc_ids] for doc_ids in docs_ids]
        return np.array(ids), np.array(vectors)


class HFIndexBase(Index):
    def __init__(self, vector_size, dataset, index_initialized=False):
        self.vector_size = vector_size
        self.dataset = dataset
        self._index_initialized = index_initialized
        self._check_dataset_format(with_index=index_initialized)
        dataset.set_format("numpy", columns=["embeddings"], output_all_columns=True, dtype="float32")

    def _check_dataset_format(self, with_index: bool):
        if not isinstance(self.dataset, Dataset):
            raise ValueError(f"Dataset should be a datasets.Dataset object, but got {type(self.dataset)}")
        if len({"title", "text", "embeddings"} - set(self.dataset.column_names)) > 0:
            raise ValueError(
                "Dataset should be a dataset with the following columns: "
                "title (str), text (str) and embeddings (arrays of dimension vector_size), "
                f"but got columns {self.dataset.column_names}"
            )
        if with_index and "embeddings" not in self.dataset.list_indexes():
            raise ValueError(
                "Missing faiss index in the dataset. Make sure you called `dataset.add_faiss_index` to compute it "
                "or `dataset.load_faiss_index` to load one from the disk."
            )

    def init_index(self):
        raise NotImplementedError()

    def is_initialized(self):
        return self._index_initialized

    def get_doc_dicts(self, doc_ids: np.ndarray) -> List[dict]:
        return [self.dataset[doc_ids[i].tolist()] for i in range(doc_ids.shape[0])]

    def get_top_docs(self, question_hidden_states: np.ndarray, n_docs=5) -> Tuple[np.ndarray, np.ndarray]:
        _, ids = self.dataset.search_batch("embeddings", question_hidden_states, n_docs)
        docs = [self.dataset[[i for i in indices if i >= 0]] for indices in ids]
        vectors = [doc["embeddings"] for doc in docs]
        for i in range(len(vectors)):
            if len(vectors[i]) < n_docs:
                vectors[i] = np.vstack([vectors[i], np.zeros((n_docs - len(vectors[i]), self.vector_size))])
        return np.array(ids), np.array(vectors)  


class CanonicalHFIndex(HFIndexBase):
    

    def __init__(
        self,
        vector_size: int,
        dataset_name: str = "wiki_dpr",
        dataset_split: str = "train",
        index_name: Optional[str] = None,
        index_path: Optional[str] = None,
        use_dummy_dataset=False,
    ):
        if int(index_path is None) + int(index_name is None) != 1:
            raise ValueError("Please provide `index_name` or `index_path`.")
        self.dataset_name = dataset_name
        self.dataset_split = dataset_split
        self.index_name = index_name
        self.index_path = index_path
        self.use_dummy_dataset = use_dummy_dataset
        logger.info(f"Loading passages from {self.dataset_name}")
        dataset = load_dataset(
            self.dataset_name, with_index=False, split=self.dataset_split, dummy=self.use_dummy_dataset
        )
        super().__init__(vector_size, dataset, index_initialized=False)

    def init_index(self):
        if self.index_path is not None:
            logger.info(f"Loading index from {self.index_path}")
            self.dataset.load_faiss_index("embeddings", file=self.index_path)
        else:
            logger.info(f"Loading index from {self.dataset_name} with index name {self.index_name}")
            self.dataset = load_dataset(
                self.dataset_name,
                with_embeddings=True,
                with_index=True,
                split=self.dataset_split,
                index_name=self.index_name,
                dummy=self.use_dummy_dataset,
            )
            self.dataset.set_format("numpy", columns=["embeddings"], output_all_columns=True)
        self._index_initialized = True


class CustomHFIndex(HFIndexBase):
    

    def __init__(self, vector_size: int, dataset, index_path=None):
        super().__init__(vector_size, dataset, index_initialized=index_path is None)
        self.index_path = index_path

    @classmethod
    def load_from_disk(cls, vector_size, dataset_path, index_path):
        logger.info(f"Loading passages from {dataset_path}")
        if dataset_path is None or index_path is None:
            raise ValueError(
                "Please provide `dataset_path` and `index_path` after calling `dataset.save_to_disk(dataset_path)` "
                "and `dataset.get_index('embeddings').save(index_path)`."
            )
        dataset = load_from_disk(dataset_path)
        return cls(vector_size=vector_size, dataset=dataset, index_path=index_path)

    def init_index(self):
        if not self.is_initialized():
            logger.info(f"Loading index from {self.index_path}")
            self.dataset.load_faiss_index("embeddings", file=self.index_path)
            self._index_initialized = True


class RagRetriever:
    

    def __init__(self, config, question_encoder_tokenizer, generator_tokenizer, index=None, init_retrieval=True):
        self._init_retrieval = init_retrieval
        requires_backends(self, ["datasets", "faiss"])
        super().__init__()
        self.index = index or self._build_index(config)
        self.generator_tokenizer = generator_tokenizer
        self.question_encoder_tokenizer = question_encoder_tokenizer

        self.n_docs = config.n_docs
        self.batch_size = config.retrieval_batch_size

        self.config = config
        if self._init_retrieval:
            self.init_retrieval()

        self.ctx_encoder_tokenizer = None
        self.return_tokenized_docs = False

    @staticmethod
    def _build_index(config):
        if config.index_name == "legacy":
            return LegacyIndex(
                config.retrieval_vector_size,
                config.index_path or LEGACY_INDEX_PATH,
            )
        elif config.index_name == "custom":
            return CustomHFIndex.load_from_disk(
                vector_size=config.retrieval_vector_size,
                dataset_path=config.passages_path,
                index_path=config.index_path,
            )
        else:
            return CanonicalHFIndex(
                vector_size=config.retrieval_vector_size,
                dataset_name=config.dataset,
                dataset_split=config.dataset_split,
                index_name=config.index_name,
                index_path=config.index_path,
                use_dummy_dataset=config.use_dummy_dataset,
            )

    @classmethod
    def from_pretrained(cls, retriever_name_or_path, indexed_dataset=None, **kwargs):
        requires_backends(cls, ["datasets", "faiss"])
        config = kwargs.pop("config", None) or RagConfig.from_pretrained(retriever_name_or_path, **kwargs)
        rag_tokenizer = RagTokenizer.from_pretrained(retriever_name_or_path, config=config)
        question_encoder_tokenizer = rag_tokenizer.question_encoder
        generator_tokenizer = rag_tokenizer.generator
        if indexed_dataset is not None:
            config.index_name = "custom"
            index = CustomHFIndex(config.retrieval_vector_size, indexed_dataset)
        else:
            index = cls._build_index(config)
        return cls(
            config,
            question_encoder_tokenizer=question_encoder_tokenizer,
            generator_tokenizer=generator_tokenizer,
            index=index,
        )

    def save_pretrained(self, save_directory):
        if isinstance(self.index, CustomHFIndex):
            if self.config.index_path is None:
                index_path = os.path.join(save_directory, "hf_dataset_index.faiss")
                self.index.dataset.get_index("embeddings").save(index_path)
                self.config.index_path = index_path
            if self.config.passages_path is None:
                passages_path = os.path.join(save_directory, "hf_dataset")
                
                faiss_index = self.index.dataset._indexes.pop("embeddings")
                self.index.dataset.save_to_disk(passages_path)
                self.index.dataset._indexes["embeddings"] = faiss_index
                self.config.passages_path = passages_path
        self.config.save_pretrained(save_directory)
        rag_tokenizer = RagTokenizer(
            question_encoder=self.question_encoder_tokenizer,
            generator=self.generator_tokenizer,
        )
        rag_tokenizer.save_pretrained(save_directory)

    def init_retrieval(self):
        

        logger.info("initializing retrieval")
        self.index.init_index()

    def postprocess_docs(self, docs, input_strings, prefix, n_docs, return_tensors=None):
        r

        def cat_input_and_doc(doc_title, doc_text, input_string, prefix):
            
            
            if doc_title.startswith('"'):
                doc_title = doc_title[1:]
            if doc_title.endswith('"'):
                doc_title = doc_title[:-1]
            if prefix is None:
                prefix = ""
            out = (prefix + doc_title + self.config.title_sep + doc_text + self.config.doc_sep + input_string).replace(
                "  ", " "
            )
            return out

        rag_input_strings = [
            cat_input_and_doc(
                docs[i]["title"][j],
                docs[i]["text"][j],
                input_strings[i],
                prefix,
            )
            for i in range(len(docs))
            for j in range(n_docs)
        ]

        contextualized_inputs = self.generator_tokenizer.batch_encode_plus(
            rag_input_strings,
            max_length=self.config.max_combined_length,
            return_tensors=return_tensors,
            padding="max_length",
            truncation=True,
        )

        return contextualized_inputs["input_ids"], contextualized_inputs["attention_mask"]

    def _chunk_tensor(self, t: Iterable, chunk_size: int) -> List[Iterable]:
        return [t[i : i + chunk_size] for i in range(0, len(t), chunk_size)]

    def _main_retrieve(self, question_hidden_states: np.ndarray, n_docs: int) -> Tuple[np.ndarray, np.ndarray]:
        question_hidden_states_batched = self._chunk_tensor(question_hidden_states, self.batch_size)
        ids_batched = []
        vectors_batched = []
        for question_hidden_states in question_hidden_states_batched:
            start_time = time.time()
            ids, vectors = self.index.get_top_docs(question_hidden_states, n_docs)
            logger.debug(
                f"index search time: {time.time() - start_time} sec, batch size {question_hidden_states.shape}"
            )
            ids_batched.extend(ids)
            vectors_batched.extend(vectors)
        return (
            np.array(ids_batched),
            np.array(vectors_batched),
        )  

    def retrieve(self, question_hidden_states: np.ndarray, n_docs: int) -> Tuple[np.ndarray, List[dict]]:
        

        doc_ids, retrieved_doc_embeds = self._main_retrieve(question_hidden_states, n_docs)
        return retrieved_doc_embeds, doc_ids, self.index.get_doc_dicts(doc_ids)

    def set_ctx_encoder_tokenizer(self, ctx_encoder_tokenizer: PreTrainedTokenizer):
        
        self.ctx_encoder_tokenizer = ctx_encoder_tokenizer
        self.return_tokenized_docs = True

    def __call__(
        self,
        question_input_ids: List[List[int]],
        question_hidden_states: np.ndarray,
        prefix=None,
        n_docs=None,
        return_tensors=None,
    ) -> BatchEncoding:
        

        n_docs = n_docs if n_docs is not None else self.n_docs
        prefix = prefix if prefix is not None else self.config.generator.prefix
        retrieved_doc_embeds, doc_ids, docs = self.retrieve(question_hidden_states, n_docs)

        input_strings = self.question_encoder_tokenizer.batch_decode(question_input_ids, skip_special_tokens=True)
        context_input_ids, context_attention_mask = self.postprocess_docs(
            docs, input_strings, prefix, n_docs, return_tensors=return_tensors
        )

        if self.return_tokenized_docs:
            retrieved_doc_text = []
            retrieved_doc_title = []

            for b_idx in range(len(docs)):
                for doc_idx in range(n_docs):
                    retrieved_doc_text.append(docs[b_idx]["text"][doc_idx])
                    retrieved_doc_title.append(docs[b_idx]["title"][doc_idx])

            tokenized_docs = self.ctx_encoder_tokenizer(
                retrieved_doc_title,
                retrieved_doc_text,
                truncation=True,
                padding="longest",
                return_tensors=return_tensors,
            )

            return BatchEncoding(
                {
                    "context_input_ids": context_input_ids,
                    "context_attention_mask": context_attention_mask,
                    "retrieved_doc_embeds": retrieved_doc_embeds,
                    "doc_ids": doc_ids,
                    "tokenized_doc_ids": tokenized_docs["input_ids"],
                    "tokenized_doc_attention_mask": tokenized_docs["attention_mask"],
                },
                tensor_type=return_tensors,
            )

        else:
            return BatchEncoding(
                {
                    "context_input_ids": context_input_ids,
                    "context_attention_mask": context_attention_mask,
                    "retrieved_doc_embeds": retrieved_doc_embeds,
                    "doc_ids": doc_ids,
                },
                tensor_type=return_tensors,
            )