delete the old self retriever object before assigning the new index a distributed retriever built on top of the ray api a library for building distributed applications https docs ray ioenmaster package during training all training workers initialize their own instance of a ragraydistributedretriever and each instance of this distributed retriever shares a common set of retrieval ray actors https docs ray ioenmasterwalkthrough htmlremote classesactors that load the index on separate processes ray handles the communication between the ragraydistributedretriever instances and the remote ray actors if training is done in a nondistributed setup the index will simply be loaded in the same process as the training worker and ray will not be used args config class transformers ragconfig the configuration of the rag model this retriever is used with contains parameters indicating which index to build questionencodertokenizer class transformers pretrainedtokenizer the tokenizer that was used to tokenize the question it is used to decode the question and then use the generatortokenizer generatortokenizer class transformers pretrainedtokenizer the tokenizer used for the generator part of the ragmodel retrievalworkers obj listray actorclassrayretriever a list of already initialized rayretriever actors these actor classes run on remote processes and are responsible for performing the index lookup index class transformers retrievalrag index optional defaults to the one defined by the configuration if specified use this index instead of the one built using the configuration retriever initialization function needs to be called from the training process this function triggers retrieval initialization for all retrieval actors if using distributed setting or loads index into current process if training is not distributed nondistributed training load index into this same process retrieves documents for specified questionhiddenstates if running training with multiple workers a random retrieval actor is selected to perform the index lookup and return the result args questionhiddenstates obj np ndarray of shape obj batchsize vectorsize a batch of query vectors to retrieve with ndocs obj int the number of docs retrieved per query output retrieveddocembeds obj np ndarray of shape obj batchsize ndocs dim the retrieval embeddings of the retrieved docs per query docids obj np ndarray of shape obj batchsize ndocs the ids of the documents in the index docdicts obj listdict the retrieveddocembeds examples per query select a random retrieval actor access from the training loop build the index object again delete the old self retriever object before assigning the new index a distributed retriever built on top of the ray api a library for building distributed applications https docs ray io en master package during training all training workers initialize their own instance of a ragraydistributedretriever and each instance of this distributed retriever shares a common set of retrieval ray actors https docs ray io en master walkthrough html remote classes actors that load the index on separate processes ray handles the communication between the ragraydistributedretriever instances and the remote ray actors if training is done in a non distributed setup the index will simply be loaded in the same process as the training worker and ray will not be used args config class transformers ragconfig the configuration of the rag model this retriever is used with contains parameters indicating which index to build question_encoder_tokenizer class transformers 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 class transformers pretrainedtokenizer the tokenizer used for the generator part of the ragmodel retrieval_workers obj list ray actorclass rayretriever a list of already initialized rayretriever actors these actor classes run on remote processes and are responsible for performing the index lookup index class transformers 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 retriever initialization function needs to be called from the training process this function triggers retrieval initialization for all retrieval actors if using distributed setting or loads index into current process if training is not distributed non distributed training load index into this same process retrieves documents for specified question_hidden_states if running training with multiple workers a random retrieval actor is selected to perform the index lookup and return the result args question_hidden_states obj np ndarray of shape obj batch_size vector_size a batch of query vectors to retrieve with n_docs obj int the number of docs retrieved per query output retrieved_doc_embeds obj np ndarray of shape obj batch_size n_docs dim the retrieval embeddings of the retrieved docs per query doc_ids obj np ndarray of shape obj batch_size n_docs the ids of the documents in the index doc_dicts obj list dict the retrieved_doc_embeds examples per query select a random retrieval actor access from the training loop build the index object again
import logging
import random

import ray

from transformers import RagConfig, RagRetriever, RagTokenizer
from transformers.models.rag.retrieval_rag import CustomHFIndex


logger = logging.getLogger(__name__)


class RayRetriever:
    def __init__(self):
        self.initialized = False

    def create_rag_retriever(self, config, question_encoder_tokenizer, generator_tokenizer, index):
        if not self.initialized:
            self.retriever = RagRetriever(
                config,
                question_encoder_tokenizer=question_encoder_tokenizer,
                generator_tokenizer=generator_tokenizer,
                index=index,
                init_retrieval=False,
            )
            self.initialized = True

    def init_retrieval(self):
        self.retriever.index.init_index()

    def clear_object(self):
        
        del self.retriever
        self.initialized = False

    def retrieve(self, question_hidden_states, n_docs):
        doc_ids, retrieved_doc_embeds = self.retriever._main_retrieve(question_hidden_states, n_docs)
        doc_dicts = self.retriever.index.get_doc_dicts(doc_ids)
        return doc_ids, retrieved_doc_embeds, doc_dicts


class RagRayDistributedRetriever(RagRetriever):
    

    def __init__(self, config, question_encoder_tokenizer, generator_tokenizer, retrieval_workers, index=None):
        if index is not None and index.is_initialized() and len(retrieval_workers) > 0:
            raise ValueError(
                "When using Ray for distributed fine-tuning, "
                "you'll need to provide the paths instead, "
                "as the dataset and the index are loaded "
                "separately. More info in examples/rag/use_own_knowledge_dataset.py "
            )

        super().__init__(
            config,
            question_encoder_tokenizer=question_encoder_tokenizer,
            generator_tokenizer=generator_tokenizer,
            index=index,
            init_retrieval=False,
        )

        self.retrieval_workers = retrieval_workers
        self.question_encoder_tokenizer = question_encoder_tokenizer
        self.generator_tokenizer = generator_tokenizer
        if len(self.retrieval_workers) > 0:
            ray.get(
                [
                    worker.create_rag_retriever.remote(config, question_encoder_tokenizer, generator_tokenizer, index)
                    for worker in self.retrieval_workers
                ]
            )

    def init_retrieval(self):
        
        logger.info("initializing retrieval")

        if len(self.retrieval_workers) > 0:
            ray.get([worker.init_retrieval.remote() for worker in self.retrieval_workers])
        else:
            
            self.index.init_index()

    def retrieve(self, question_hidden_states, n_docs):
        
        if len(self.retrieval_workers) > 0:
            
            random_worker = self.retrieval_workers[random.randint(0, len(self.retrieval_workers) - 1)]
            doc_ids, retrieved_doc_embeds, doc_dicts = ray.get(
                random_worker.retrieve.remote(question_hidden_states, n_docs)
            )
        else:
            doc_ids, retrieved_doc_embeds = self._main_retrieve(question_hidden_states, n_docs)
            doc_dicts = self.index.get_doc_dicts(doc_ids)
        return retrieved_doc_embeds, doc_ids, doc_dicts

    @classmethod
    def get_tokenizers(cls, retriever_name_or_path, indexed_dataset=None, **kwargs):
        return super(RagRayDistributedRetriever, cls).get_tokenizers(retriever_name_or_path, indexed_dataset, **kwargs)

    @classmethod
    def from_pretrained(cls, retriever_name_or_path, actor_handles, indexed_dataset=None, **kwargs):
        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,
            retrieval_workers=actor_handles,
            index=index,
        )

    def re_load(self):
        logger.info("re-loading the new dataset with embeddings")
        

        ray.get([worker.clear_object.remote() for worker in self.retrieval_workers])

        
        index = self._build_index(self.config)

        ray.get(
            [
                worker.create_rag_retriever.remote(
                    self.config, self.question_encoder_tokenizer, self.generator_tokenizer, index
                )
                for worker in self.retrieval_workers
            ]
        )