examples/research_projects/rag/distributed_ray_retriever.py/RayRetriever/retrieve
class RayRetriever:    def retrieve(self, question_hidden_states, n_docs):
        doc_ids, retrieved_doc_embeds = self.retriever._main_retrieve(question_hidden_states, n_docs)
        return doc_ids, retrieved_doc_embeds