search_r1/search/retrieval_server.py

• load_corpus · function · L27-L34 — def load_corpus(corpus_path: str)
• read_jsonl · function · L36-L41 — def read_jsonl(file_path)
• load_docs · function · L43-L45 — def load_docs(corpus, doc_idxs)
• load_model · function · L47-L55 — def load_model(model_path: str, use_fp16: bool = False)
• pooling · function · L57-L71 — def pooling( pooler_output, last_hidden_state, attention_mask = None, pooling_method = "mean" )
• Encoder · class · L73-L128 — class Encoder
• init · method · L74-L82 — def init(self, model_name, model_path, pooling_method, max_length, use_fp16)
• encode · method · L85-L128 — def encode(self, query_list: List[str], is_query=True) -> np.ndarray: # processing query for different encoders
• BaseRetriever · class · L130-L149 — class BaseRetriever
• init · method · L131-L137 — def init(self, config)
• _search · method · L139-L140 — def _search(self, query: str, num: int, return_score: bool)
• _batch_search · method · L142-L143 — def _batch_search(self, query_list: List[str], num: int, return_score: bool)
• search · method · L145-L146 — def search(self, query: str, num: int = None, return_score: bool = False)
• batch_search · method · L148-L149 — def batch_search(self, query_list: List[str], num: int = None, return_score: bool = False)
• BM25Retriever · class · L151-L210 — class BM25Retriever(BaseRetriever)
• init · method · L152-L159 — def init(self, config)
• _check_contain_doc · method · L161-L162 — def _check_contain_doc(self)
• _search · method · L164-L198 — def _search(self, query: str, num: int = None, return_score: bool = False)
• _batch_search · method · L200-L210 — def _batch_search(self, query_list: List[str], num: int = None, return_score: bool = False)
• DenseRetriever · class · L212-L272 — class DenseRetriever(BaseRetriever)
• init · method · L213-L231 — def init(self, config)
• _search · method · L233-L244 — def _search(self, query: str, num: int = None, return_score: bool = False)
• _batch_search · method · L246-L272 — def _batch_search(self, query_list: List[str], num: int = None, return_score: bool = False)
• get_retriever · function · L274-L278 — def get_retriever(config)
• Config · class · L285-L316 — class Config
• init · method · L290-L316 — def init( self, retrieval_method: str = "bm25", retrieval_topk: int = 10, index_path: str = "./index/bm25", corpus_path: str = "./data/corpus.jsonl", dataset_path: str = "./data", data_split: str = "train", faiss_gpu: bool = True, retrieval_model_path: str = "./model", retrieval_pooling_method: str = "mean", retrieval_query_max_length: int = 256, retrieval_use_fp16: bool = False, retrieval_batch_size: int = 128 )
• QueryRequest · class · L319-L322 — class QueryRequest(BaseModel)
• retrieve_endpoint · function · L346-L377 — def retrieve_endpoint(request: QueryRequest)