chore: import upstream snapshot with attribution

simlm/src/inference/encode_main.py

@@ -0,0 +1,119 @@
+import os
+import tqdm
+import torch
+
+from contextlib import nullcontext
+from torch.utils.data import DataLoader
+from functools import partial
+from datasets import load_dataset
+from typing import Dict, List
+from transformers.file_utils import PaddingStrategy
+from transformers import (
+    AutoTokenizer,
+    PreTrainedTokenizerFast,
+    DataCollatorWithPadding,
+    HfArgumentParser,
+    BatchEncoding
+)
+
+from config import Arguments
+from logger_config import logger
+from utils import move_to_cuda
+from models import BiencoderModelForInference, BiencoderOutput
+
+parser = HfArgumentParser((Arguments,))
+args: Arguments = parser.parse_args_into_dataclasses()[0]
+
+
+def _psg_transform_func(tokenizer: PreTrainedTokenizerFast,
+                        examples: Dict[str, List]) -> BatchEncoding:
+    batch_dict = tokenizer(examples['title'],
+                           text_pair=examples['contents'],
+                           max_length=args.p_max_len,
+                           padding=PaddingStrategy.DO_NOT_PAD,
+                           truncation=True)
+    # for co-Condenser reproduction purpose only
+    if args.model_name_or_path.startswith('Luyu/'):
+        del batch_dict['token_type_ids']
+
+    return batch_dict
+
+
+@torch.no_grad()
+def _worker_encode_passages(gpu_idx: int):
+    def _get_out_path(shard_idx: int = 0) -> str:
+        return '{}/shard_{}_{}'.format(args.encode_save_dir, gpu_idx, shard_idx)
+
+    if os.path.exists(_get_out_path(0)):
+        logger.error('{} already exists, will skip encoding'.format(_get_out_path(0)))
+        return
+
+    dataset = load_dataset('json', data_files=args.encode_in_path)['train']
+    if args.dry_run:
+        dataset = dataset.select(range(4096))
+    dataset = dataset.shard(num_shards=torch.cuda.device_count(),
+                            index=gpu_idx,
+                            contiguous=True)
+
+    logger.info('GPU {} needs to process {} examples'.format(gpu_idx, len(dataset)))
+    torch.cuda.set_device(gpu_idx)
+
+    tokenizer: PreTrainedTokenizerFast = AutoTokenizer.from_pretrained(args.model_name_or_path)
+    model: BiencoderModelForInference = BiencoderModelForInference.build(args)
+    model.eval()
+    model.cuda()
+
+    dataset.set_transform(partial(_psg_transform_func, tokenizer))
+
+    data_collator = DataCollatorWithPadding(tokenizer, pad_to_multiple_of=8 if args.fp16 else None)
+    data_loader = DataLoader(
+        dataset,
+        batch_size=args.encode_batch_size,
+        shuffle=False,
+        drop_last=False,
+        num_workers=args.dataloader_num_workers,
+        collate_fn=data_collator,
+        pin_memory=True)
+
+    num_encoded_docs, encoded_embeds, cur_shard_idx = 0, [], 0
+    for batch_dict in tqdm.tqdm(data_loader, desc='passage encoding', mininterval=8):
+        batch_dict = move_to_cuda(batch_dict)
+
+        with torch.cuda.amp.autocast() if args.fp16 else nullcontext():
+            outputs: BiencoderOutput = model(query=None, passage=batch_dict)
+        encoded_embeds.append(outputs.p_reps.cpu())
+        num_encoded_docs += outputs.p_reps.shape[0]
+
+        if num_encoded_docs >= args.encode_shard_size:
+            out_path = _get_out_path(cur_shard_idx)
+            concat_embeds = torch.cat(encoded_embeds, dim=0)
+            logger.info('GPU {} save {} embeds to {}'.format(gpu_idx, concat_embeds.shape[0], out_path))
+            torch.save(concat_embeds, out_path)
+
+            cur_shard_idx += 1
+            num_encoded_docs = 0
+            encoded_embeds.clear()
+
+    if num_encoded_docs > 0:
+        out_path = _get_out_path(cur_shard_idx)
+        concat_embeds = torch.cat(encoded_embeds, dim=0)
+        logger.info('GPU {} save {} embeds to {}'.format(gpu_idx, concat_embeds.shape[0], out_path))
+        torch.save(concat_embeds, out_path)
+
+    logger.info('Done computing score for worker {}'.format(gpu_idx))
+
+
+def _batch_encode_passages():
+    logger.info('Args={}'.format(str(args)))
+    gpu_count = torch.cuda.device_count()
+    if gpu_count == 0:
+        logger.error('No gpu available')
+        return
+
+    logger.info('Use {} gpus'.format(gpu_count))
+    torch.multiprocessing.spawn(_worker_encode_passages, args=(), nprocs=gpu_count)
+    logger.info('Done batch encode passages')
+
+
+if __name__ == '__main__':
+    _batch_encode_passages()

simlm/src/inference/gen_teacher_scores.py

@@ -0,0 +1,174 @@
+import os
+import tqdm
+import torch
+
+from contextlib import nullcontext
+from torch.utils.data import DataLoader
+from functools import partial
+from datasets import Dataset, load_dataset
+from typing import Dict, List
+from transformers.file_utils import PaddingStrategy
+from transformers.modeling_outputs import SequenceClassifierOutput
+from transformers import (
+    AutoTokenizer,
+    PreTrainedTokenizerFast,
+    DataCollatorWithPadding,
+    HfArgumentParser,
+    BatchEncoding
+)
+
+from config import Arguments
+from logger_config import logger
+from utils import move_to_cuda
+from models import RerankerForInference
+from data_utils import load_corpus, load_queries, save_to_readable_format
+
+parser = HfArgumentParser((Arguments,))
+args: Arguments = parser.parse_args_into_dataclasses()[0]
+kd_gen_score_in_path = os.path.join(args.data_dir, '{}.jsonl'.format(args.kd_gen_score_split))
+kd_gen_score_out_path = os.path.join(args.data_dir, 'kd_{}.jsonl'.format(args.kd_gen_score_split))
+
+
+def _kd_gen_score_transform_func(tokenizer: PreTrainedTokenizerFast,
+                                 corpus: Dataset,
+                                 queries: Dict[str, str],
+                                 examples: Dict[str, List]) -> BatchEncoding:
+    input_docs: List[str] = []
+
+    # ATTENTION: this code should be consistent with CrossEncoderDataLoader
+    for doc_id in examples['doc_id']:
+        doc_id = int(doc_id)
+        prefix = ''
+        if corpus[doc_id].get('title', ''):
+            prefix = corpus[doc_id]['title'] + ': '
+        input_docs.append(prefix + corpus[doc_id]['contents'])
+
+    input_queries = [queries[query_id] for query_id in examples['query_id']]
+    batch_dict = tokenizer(input_queries,
+                           text_pair=input_docs,
+                           max_length=args.rerank_max_length,
+                           padding=PaddingStrategy.DO_NOT_PAD,
+                           truncation=True)
+
+    return batch_dict
+
+
+def _get_shard_path(worker_idx: int) -> str:
+    return '{}_shard_{}'.format(kd_gen_score_in_path, worker_idx)
+
+
+@torch.no_grad()
+def _worker_gen_teacher_score(gpu_idx: int):
+    dataset = load_dataset('json', data_files=kd_gen_score_in_path)['train']
+    if args.dry_run:
+        dataset = dataset.select(range(100))
+    dataset = dataset.shard(num_shards=torch.cuda.device_count(),
+                            index=gpu_idx,
+                            contiguous=True)
+
+    qid_pids = []
+    for ex in tqdm.tqdm(dataset, desc='get qid-pid pairs', mininterval=3):
+        for pos_doc_id in ex['positives']['doc_id']:
+            qid_pids.append((ex['query_id'], pos_doc_id))
+        for neg_doc_id in ex['negatives']['doc_id'][:args.kd_gen_score_n_neg]:
+            qid_pids.append((ex['query_id'], neg_doc_id))
+
+    dataset = Dataset.from_dict({'query_id': [t[0] for t in qid_pids],
+                                 'doc_id': [t[1] for t in qid_pids]})
+
+    query_ids, doc_ids = dataset['query_id'], dataset['doc_id']
+
+    logger.info('GPU {} needs to process {} examples'.format(gpu_idx, len(dataset)))
+    torch.cuda.set_device(gpu_idx)
+
+    tokenizer: PreTrainedTokenizerFast = AutoTokenizer.from_pretrained(args.model_name_or_path)
+    model: RerankerForInference = RerankerForInference.from_pretrained(args.model_name_or_path)
+    model.eval()
+    model.cuda()
+
+    corpus: Dataset = load_corpus(path=os.path.join(args.data_dir, 'passages.jsonl.gz'))
+    queries = load_queries(path='{}/{}_queries.tsv'.format(args.data_dir, args.kd_gen_score_split),
+                           task_type=args.task_type)
+    dataset.set_transform(partial(_kd_gen_score_transform_func, tokenizer, corpus, queries))
+
+    data_collator = DataCollatorWithPadding(tokenizer, pad_to_multiple_of=8 if args.fp16 else None)
+    data_loader = DataLoader(
+        dataset,
+        batch_size=args.kd_gen_score_batch_size,
+        shuffle=False,
+        drop_last=False,
+        num_workers=args.dataloader_num_workers,
+        collate_fn=data_collator,
+        pin_memory=True)
+
+    scores = []
+    for batch_dict in tqdm.tqdm(data_loader, desc='generate teacher score', mininterval=5):
+        batch_dict = move_to_cuda(batch_dict)
+
+        with torch.cuda.amp.autocast() if args.fp16 else nullcontext():
+            outputs: SequenceClassifierOutput = model(batch_dict)
+        scores.append(outputs.logits.squeeze(dim=-1).cpu())
+        assert len(scores[-1].shape) == 1
+
+    all_scores = torch.cat(scores, dim=-1)
+    assert all_scores.shape[0] == len(dataset), '{} != {}'
+    all_scores = all_scores.tolist()
+
+    with open(_get_shard_path(gpu_idx), 'w', encoding='utf-8') as writer:
+        for idx in range(len(query_ids)):
+            writer.write('{}\t{}\t{}\n'.format(query_ids[idx], doc_ids[idx], round(all_scores[idx], 5)))
+
+    logger.info('Done computing teacher score for worker {}'.format(gpu_idx))
+
+
+def _merge_teacher_scores(worker_cnt: int):
+    qid_to_pid_to_score = {}
+    for worker_idx in range(worker_cnt):
+        shard_path = _get_shard_path(worker_idx)
+        for line in tqdm.tqdm(open(shard_path, 'r', encoding='utf-8'),
+                              desc='Load shard {} score'.format(worker_idx), mininterval=3):
+            fs = line.strip().split('\t')
+            assert len(fs) == 3
+            qid, pid, score = fs
+            if qid not in qid_to_pid_to_score:
+                qid_to_pid_to_score[qid] = {}
+            qid_to_pid_to_score[qid][pid] = float(score)
+        os.remove(shard_path)
+
+    dataset = load_dataset('json', data_files=kd_gen_score_in_path)['train']
+    if args.dry_run:
+        dataset = dataset.select(range(100))
+
+    def _update_score(ex: Dict) -> Dict:
+        query_id = ex['query_id']
+        pid_to_score = qid_to_pid_to_score[query_id]
+        ex['negatives']['doc_id'] = [neg_doc_id for neg_doc_id in ex['negatives']['doc_id'] if neg_doc_id in pid_to_score]
+        ex['positives']['score'] = [pid_to_score[pos_doc_id] for pos_doc_id in ex['positives']['doc_id']]
+        ex['negatives']['score'] = [pid_to_score[neg_doc_id] for neg_doc_id in ex['negatives']['doc_id']]
+        return ex
+
+    dataset = dataset.map(_update_score, num_proc=4)
+    logger.info('Writing teacher score to {}'.format(kd_gen_score_out_path))
+    dataset.to_json(kd_gen_score_out_path, force_ascii=False, lines=True)
+
+
+def _batch_compute_teacher_score():
+    logger.info('Args={}'.format(str(args)))
+    gpu_count = torch.cuda.device_count()
+    if gpu_count == 0:
+        logger.error('No gpu available')
+        return
+
+    logger.info('Use {} gpus'.format(gpu_count))
+    torch.multiprocessing.spawn(_worker_gen_teacher_score, args=(), nprocs=gpu_count)
+    logger.info('Done batch generate teacher score')
+
+    _merge_teacher_scores(gpu_count)
+    logger.info('Done merge results')
+
+    corpus = load_corpus(path=os.path.join(args.data_dir, 'passages.jsonl.gz'))
+    save_to_readable_format(in_path=kd_gen_score_out_path, corpus=corpus)
+
+
+if __name__ == '__main__':
+    _batch_compute_teacher_score()

simlm/src/inference/rerank_main.py

@@ -0,0 +1,133 @@
+import os
+import tqdm
+import torch
+
+from contextlib import nullcontext
+from torch.utils.data import DataLoader
+from functools import partial
+from datasets import Dataset
+from typing import Dict, List
+from transformers.file_utils import PaddingStrategy
+from transformers.modeling_outputs import SequenceClassifierOutput
+from transformers import (
+    AutoTokenizer,
+    PreTrainedTokenizerFast,
+    DataCollatorWithPadding,
+    HfArgumentParser,
+    BatchEncoding
+)
+
+from config import Arguments
+from logger_config import logger
+from utils import move_to_cuda
+from models import RerankerForInference
+from data_utils import load_msmarco_predictions, load_corpus, load_queries, \
+    merge_rerank_predictions, get_rerank_shard_path
+
+parser = HfArgumentParser((Arguments,))
+args: Arguments = parser.parse_args_into_dataclasses()[0]
+
+
+def _rerank_transform_func(tokenizer: PreTrainedTokenizerFast,
+                           corpus: Dataset,
+                           queries: Dict[str, str],
+                           examples: Dict[str, List]) -> BatchEncoding:
+    input_docs: List[str] = []
+
+    # ATTENTION: this code should be consistent with RerankDataLoader
+    for doc_id in examples['doc_id']:
+        doc_id = int(doc_id)
+        prefix = ''
+        if corpus[doc_id].get('title', ''):
+            prefix = corpus[doc_id]['title'] + ': '
+        input_docs.append(prefix + corpus[doc_id]['contents'])
+
+    input_queries = [queries[query_id] for query_id in examples['query_id']]
+    batch_dict = tokenizer(input_queries,
+                           text_pair=input_docs,
+                           max_length=args.rerank_max_length,
+                           padding=PaddingStrategy.DO_NOT_PAD,
+                           truncation=True)
+
+    return batch_dict
+
+
+@torch.no_grad()
+def _worker_compute_reranker_score(gpu_idx: int):
+    preds = load_msmarco_predictions(args.rerank_in_path)
+    query_ids = sorted(list(preds.keys()))
+    qid_pid = []
+    for query_id in tqdm.tqdm(query_ids, desc='load qid-pid', mininterval=2):
+        qid_pid += [(scored_doc.qid, scored_doc.pid) for scored_doc in preds[query_id]
+                    if scored_doc.rank <= args.rerank_depth]
+
+    dataset = Dataset.from_dict({'query_id': [t[0] for t in qid_pid],
+                                 'doc_id': [t[1] for t in qid_pid]})
+    dataset = dataset.shard(num_shards=torch.cuda.device_count(),
+                            index=gpu_idx,
+                            contiguous=True)
+
+    logger.info('GPU {} needs to process {} examples'.format(gpu_idx, len(dataset)))
+    torch.cuda.set_device(gpu_idx)
+
+    query_ids, doc_ids = dataset['query_id'], dataset['doc_id']
+    assert len(dataset) == len(query_ids)
+
+    tokenizer: PreTrainedTokenizerFast = AutoTokenizer.from_pretrained(args.model_name_or_path)
+    model: RerankerForInference = RerankerForInference.from_pretrained(args.model_name_or_path)
+    model.eval()
+    model.cuda()
+
+    corpus: Dataset = load_corpus(path=os.path.join(args.data_dir, 'passages.jsonl.gz'))
+    queries = load_queries(path='{}/{}_queries.tsv'.format(args.data_dir, args.rerank_split),
+                           task_type=args.task_type)
+    dataset.set_transform(partial(_rerank_transform_func, tokenizer, corpus, queries))
+
+    data_collator = DataCollatorWithPadding(tokenizer, pad_to_multiple_of=8 if args.fp16 else None)
+    data_loader = DataLoader(
+        dataset,
+        batch_size=args.rerank_batch_size,
+        shuffle=False,
+        drop_last=False,
+        num_workers=args.dataloader_num_workers,
+        collate_fn=data_collator,
+        pin_memory=True)
+
+    scores = []
+    for batch_dict in tqdm.tqdm(data_loader, desc='passage rerank', mininterval=5):
+        batch_dict = move_to_cuda(batch_dict)
+
+        with torch.cuda.amp.autocast() if args.fp16 else nullcontext():
+            outputs: SequenceClassifierOutput = model(batch_dict)
+        scores.append(outputs.logits.squeeze(dim=-1).cpu())
+        assert len(scores[-1].shape) == 1
+
+    all_scores = torch.cat(scores, dim=-1)
+    assert all_scores.shape[0] == len(query_ids), '{} != {}'.format(all_scores.shape[0], len(query_ids))
+    all_scores = all_scores.tolist()
+
+    with open(get_rerank_shard_path(args, gpu_idx), 'w', encoding='utf-8') as writer:
+        for idx in range(len(query_ids)):
+            # dummy rank, since a query may be split across different workers
+            writer.write('{}\t{}\t{}\t{}\n'.format(query_ids[idx], doc_ids[idx], -1, round(all_scores[idx], 5)))
+
+    logger.info('Done computing rerank score for worker {}'.format(gpu_idx))
+
+
+def _batch_compute_reranker_score():
+    logger.info('Args={}'.format(str(args)))
+    gpu_count = torch.cuda.device_count()
+    if gpu_count == 0:
+        logger.error('No gpu available')
+        return
+
+    logger.info('Use {} gpus'.format(gpu_count))
+    torch.multiprocessing.spawn(_worker_compute_reranker_score, args=(), nprocs=gpu_count)
+    logger.info('Done batch compute rerank score')
+
+    merge_rerank_predictions(args, gpu_count)
+    logger.info('Done merge results')
+
+
+if __name__ == '__main__':
+    _batch_compute_reranker_score()

simlm/src/inference/search_main.py

@@ -0,0 +1,193 @@
+import json
+import os
+import glob
+import tqdm
+import torch
+
+from contextlib import nullcontext
+from torch.utils.data import DataLoader
+from functools import partial
+from collections import defaultdict
+from datasets import Dataset
+from typing import Dict, List, Tuple
+from transformers.file_utils import PaddingStrategy
+from transformers import (
+    AutoTokenizer,
+    PreTrainedTokenizerFast,
+    DataCollatorWithPadding,
+    HfArgumentParser,
+    BatchEncoding
+)
+
+from config import Arguments
+from logger_config import logger
+from utils import move_to_cuda, save_json_to_file
+from metrics import compute_mrr, trec_eval, ScoredDoc
+from data_utils import load_queries, load_qrels, load_msmarco_predictions, save_preds_to_msmarco_format
+from models import BiencoderModelForInference, BiencoderOutput
+
+parser = HfArgumentParser((Arguments,))
+args: Arguments = parser.parse_args_into_dataclasses()[0]
+assert os.path.exists(args.encode_save_dir)
+
+
+def _get_all_shards_path() -> List[str]:
+    path_list = glob.glob('{}/shard_*_*'.format(args.encode_save_dir))
+    assert len(path_list) > 0
+
+    def _parse_worker_idx_shard_idx(p: str) -> Tuple:
+        worker_idx, shard_idx = [int(f) for f in os.path.basename(p).split('_')[-2:]]
+        return worker_idx, shard_idx
+
+    path_list = sorted(path_list, key=lambda path: _parse_worker_idx_shard_idx(path))
+    logger.info('Embeddings path list: {}'.format(path_list))
+    return path_list
+
+
+def _get_topk_result_save_path(worker_idx: int) -> str:
+    return '{}/top{}_{}_{}.txt'.format(args.search_out_dir, args.search_topk, args.search_split, worker_idx)
+
+
+def _query_transform_func(tokenizer: PreTrainedTokenizerFast,
+                          examples: Dict[str, List]) -> BatchEncoding:
+    batch_dict = tokenizer(examples['query'],
+                           max_length=args.q_max_len,
+                           padding=PaddingStrategy.DO_NOT_PAD,
+                           truncation=True)
+
+    return batch_dict
+
+
+@torch.no_grad()
+def _worker_encode_queries(gpu_idx: int) -> Tuple:
+    # fail fast if shard does not exist
+    _get_all_shards_path()
+
+    query_id_to_text = load_queries(path=os.path.join(args.data_dir, '{}_queries.tsv'.format(args.search_split)),
+                                    task_type=args.task_type)
+    query_ids = sorted(list(query_id_to_text.keys()))
+    queries = [query_id_to_text[query_id] for query_id in query_ids]
+    dataset = Dataset.from_dict({'query_id': query_ids,
+                                 'query': queries})
+    dataset = dataset.shard(num_shards=torch.cuda.device_count(),
+                            index=gpu_idx,
+                            contiguous=True)
+
+    # only keep data for current shard
+    query_ids = dataset['query_id']
+    query_id_to_text = {qid: query_id_to_text[qid] for qid in query_ids}
+
+    logger.info('GPU {} needs to process {} examples'.format(gpu_idx, len(dataset)))
+    torch.cuda.set_device(gpu_idx)
+
+    tokenizer: PreTrainedTokenizerFast = AutoTokenizer.from_pretrained(args.model_name_or_path)
+    model: BiencoderModelForInference = BiencoderModelForInference.build(args)
+    model.eval()
+    model.cuda()
+
+    dataset.set_transform(partial(_query_transform_func, tokenizer))
+
+    data_collator = DataCollatorWithPadding(tokenizer, pad_to_multiple_of=8)
+    data_loader = DataLoader(
+        dataset,
+        batch_size=512,
+        shuffle=False,
+        drop_last=False,
+        num_workers=args.dataloader_num_workers,
+        collate_fn=data_collator,
+        pin_memory=True)
+
+    encoded_embeds = []
+    for batch_dict in tqdm.tqdm(data_loader, desc='query encoding', mininterval=5):
+        batch_dict = move_to_cuda(batch_dict)
+
+        with torch.cuda.amp.autocast() if args.fp16 else nullcontext():
+            outputs: BiencoderOutput = model(query=batch_dict, passage=None)
+        encoded_embeds.append(outputs.q_reps)
+
+    query_embeds = torch.cat(encoded_embeds, dim=0)
+    logger.info('Done query encoding for worker {}'.format(gpu_idx))
+
+    return query_embeds, query_ids, query_id_to_text
+
+
+@torch.no_grad()
+def _worker_batch_search(gpu_idx: int):
+    embeds_path_list = _get_all_shards_path()
+
+    query_embeds, query_ids, query_id_to_text = _worker_encode_queries(gpu_idx)
+    assert query_embeds.shape[0] == len(query_ids), '{} != {}'.format(query_embeds.shape[0], len(query_ids))
+
+    query_id_to_topk = defaultdict(list)
+    psg_idx_offset = 0
+    for shard_idx, shard_path in enumerate(embeds_path_list):
+        shard_psg_embed = torch.load(shard_path, map_location=lambda storage, loc: storage).to(query_embeds.device)
+        logger.info('Load {} passage embeddings from {}'.format(shard_psg_embed.shape[0], shard_path))
+
+        for start in tqdm.tqdm(range(0, len(query_ids), args.search_batch_size),
+                               desc="search shard {}".format(shard_idx),
+                               mininterval=5):
+            batch_query_embed = query_embeds[start:(start + args.search_batch_size)]
+            batch_query_ids = query_ids[start:(start + args.search_batch_size)]
+            batch_score = torch.mm(batch_query_embed, shard_psg_embed.t())
+            batch_sorted_score, batch_sorted_indices = torch.topk(batch_score, k=args.search_topk, dim=-1, largest=True)
+            for batch_idx, query_id in enumerate(batch_query_ids):
+                cur_scores = batch_sorted_score[batch_idx].cpu().tolist()
+                cur_indices = [idx + psg_idx_offset for idx in batch_sorted_indices[batch_idx].cpu().tolist()]
+                query_id_to_topk[query_id] += list(zip(cur_scores, cur_indices))
+                query_id_to_topk[query_id] = sorted(query_id_to_topk[query_id], key=lambda t: (-t[0], t[1]))
+                query_id_to_topk[query_id] = query_id_to_topk[query_id][:args.search_topk]
+
+        psg_idx_offset += shard_psg_embed.shape[0]
+
+    out_path = _get_topk_result_save_path(worker_idx=gpu_idx)
+    with open(out_path, 'w', encoding='utf-8') as writer:
+        for query_id in query_id_to_text:
+            for rank, (score, doc_id) in enumerate(query_id_to_topk[query_id]):
+                writer.write('{}\t{}\t{}\t{}\n'.format(query_id, doc_id, rank + 1, round(score, 4)))
+
+    logger.info('Write scores to {} done'.format(out_path))
+
+
+def _compute_and_save_metrics(worker_cnt: int):
+    preds: Dict[str, List[ScoredDoc]] = {}
+    for worker_idx in range(worker_cnt):
+        path = _get_topk_result_save_path(worker_idx)
+        preds.update(load_msmarco_predictions(path))
+    out_path = os.path.join(args.search_out_dir, '{}.msmarco.txt'.format(args.search_split))
+    save_preds_to_msmarco_format(preds, out_path)
+    logger.info('Merge done: save {} predictions to {}'.format(len(preds), out_path))
+
+    path_qrels = os.path.join(args.data_dir, '{}_qrels.txt'.format(args.search_split))
+    if os.path.exists(path_qrels):
+        qrels = load_qrels(path=path_qrels)
+        all_metrics = trec_eval(qrels=qrels, predictions=preds)
+        all_metrics['mrr'] = compute_mrr(qrels=qrels, predictions=preds)
+
+        logger.info('{} trec metrics = {}'.format(args.search_split, json.dumps(all_metrics, ensure_ascii=False, indent=4)))
+        save_json_to_file(all_metrics, os.path.join(args.search_out_dir, 'metrics_{}.json'.format(args.search_split)))
+    else:
+        logger.warning('No qrels found for {}'.format(args.search_split))
+
+    # do some cleanup
+    for worker_idx in range(worker_cnt):
+        path = _get_topk_result_save_path(worker_idx)
+        os.remove(path)
+
+
+def _batch_search_queries():
+    logger.info('Args={}'.format(str(args)))
+    gpu_count = torch.cuda.device_count()
+    if gpu_count == 0:
+        logger.error('No gpu available')
+        return
+
+    logger.info('Use {} gpus'.format(gpu_count))
+    torch.multiprocessing.spawn(_worker_batch_search, args=(), nprocs=gpu_count)
+    logger.info('Done batch search queries')
+
+    _compute_and_save_metrics(gpu_count)
+
+
+if __name__ == '__main__':
+    _batch_search_queries()