verl/workers/retriever_workers.py

• load_corpus · function · L22-L28 — def load_corpus(corpus_path: str)
• read_jsonl · function · L31-L38 — def read_jsonl(file_path)
• load_docs · function · L41-L44 — def load_docs(corpus, doc_idxs)
• load_model · function · L47-L59 — def load_model( model_path: str, use_fp16: bool = False )
• pooling · function · L62-L76 — def pooling( pooler_output, last_hidden_state, attention_mask = None, pooling_method = "mean" )
• Encoder · class · L79-L135 — class Encoder
• init · method · L80-L88 — def init(self, model_name, model_path, pooling_method, max_length, use_fp16)
• encode · method · L91-L135 — def encode(self, query_list: List[str], is_query=True) -> np.ndarray: # processing query for different encoders
• BaseRetriever · class · L138-L168 — class BaseRetriever
• init · method · L141-L149 — def init(self, config)
• _search · method · L151-L159 — def _search(self, query: str, num: int, return_score:bool) -> List[Dict[str, str]]
• _batch_search · method · L161-L162 — def _batch_search(self, query_list, num, return_score)
• search · method · L164-L165 — def search(self, args, *kwargs)
• batch_search · method · L167-L168 — def batch_search(self, args, *kwargs)
• BM25Retriever · class · L171-L231 — class BM25Retriever(BaseRetriever)
• init · method · L174-L182 — def init(self, config)
• _check_contain_doc · method · L184-L187 — def _check_contain_doc(self)
• _search · method · L189-L217 — def _search(self, query: str, num: int = None, return_score = False) -> List[Dict[str, str]]
• _batch_search · method · L219-L231 — def _batch_search(self, query_list, num: int = None, return_score = False): # TODO: modify batch method
• DenseRetriever · class · L234-L309 — class DenseRetriever(BaseRetriever)
• init · method · L237-L257 — def init(self, config: dict, index)
• _search · method · L259-L272 — def _search(self, query: str, num: int = None, return_score = False)
• _batch_search · method · L274-L309 — def _batch_search(self, query_list: List[str], num: int = None, return_score = False)
• get_retriever · function · L311-L324 — def get_retriever(config, index)
• RetrieveWorker · class · L328-L345 — class RetrieveWorker(Worker)
• init · method · L331-L336 — def init(self, config, faiss_server)
• init_model · method · L339-L341 — def init_model(self)
• batch_search · method · L344-L345 — def batch_search(self, queries)
• FAISSIndexServer · class · L353-L383 — class FAISSIndexServer
• init · method · L356-L360 — def init(self, config)
• load_index · method · L362-L378 — def load_index(self, config)
• batch_search · method · L380-L383 — def batch_search(self, batch_emb, k)