search_r1/search/retrieval.py

• load_corpus · function · L16-L22 — def load_corpus(corpus_path: str)
• read_jsonl · function · L25-L32 — def read_jsonl(file_path)
• load_docs · function · L35-L38 — def load_docs(corpus, doc_idxs)
• load_model · function · L41-L53 — def load_model( model_path: str, use_fp16: bool = False )
• pooling · function · L56-L70 — def pooling( pooler_output, last_hidden_state, attention_mask = None, pooling_method = "mean" )
• Encoder · class · L73-L129 — class Encoder
• init · method · L74-L82 — def init(self, model_name, model_path, pooling_method, max_length, use_fp16)
• encode · method · L85-L129 — def encode(self, query_list: List[str], is_query=True) -> np.ndarray: # processing query for different encoders
• BaseRetriever · class · L132-L162 — class BaseRetriever
• init · method · L135-L143 — def init(self, config)
• _search · method · L145-L153 — def _search(self, query: str, num: int, return_score:bool) -> List[Dict[str, str]]
• _batch_search · method · L155-L156 — def _batch_search(self, query_list, num, return_score)
• search · method · L158-L159 — def search(self, args, *kwargs)
• batch_search · method · L161-L162 — def batch_search(self, args, *kwargs)
• BM25Retriever · class · L165-L224 — class BM25Retriever(BaseRetriever)
• init · method · L168-L175 — def init(self, config)
• _check_contain_doc · method · L177-L180 — def _check_contain_doc(self)
• _search · method · L182-L210 — def _search(self, query: str, num: int = None, return_score = False) -> List[Dict[str, str]]
• _batch_search · method · L212-L224 — def _batch_search(self, query_list, num: int = None, return_score = False): # TODO: modify batch method
• get_available_gpu_memory · function · L226-L233 — def get_available_gpu_memory()
• DenseRetriever · class · L236-L309 — class DenseRetriever(BaseRetriever)
• init · method · L239-L258 — def init(self, config: dict)
• _search · method · L260-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-L323 — def get_retriever(config)
• get_dataset · function · L326-L330 — def get_dataset(config)