main.py

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
import argparse
import pickle
import time
import numpy as np
import os
import json
import random
from transformers import AutoModel, AutoTokenizer
from utils import Recorder
from data_utils import to_cuda, collate_mp, ReRankingDataset
from torch.utils.data import DataLoader
import torch.distributed as dist
import torch.multiprocessing as mp
from torch.nn.parallel import DistributedDataParallel as DDP
from functools import partial
from model import RankingLoss, ReRanker, ContextualReRanker, ContextualTop1ReRanker
import tqdm
import math
import logging
import wandb
logging.getLogger("transformers.tokenization_utils").setLevel(logging.ERROR)
logging.getLogger("transformers.tokenization_utils_base").setLevel(logging.ERROR)
logging.getLogger("transformers.tokenization_utils_fast").setLevel(logging.ERROR)


def base_setting(args):
    args.batch_size = getattr(args, 'batch_size', 1)
    args.epoch = getattr(args, 'epoch', 5)
    args.report_freq = getattr(args, "report_freq", 100)
    args.accumulate_step = getattr(args, "accumulate_step", 12)
    args.margin = getattr(args, "margin", 0.01)
    args.gold_margin = getattr(args, "gold_margin", 0)
    args.model_type = getattr(args, "model_type", 'roberta-base')
    args.warmup_steps = getattr(args, "warmup_steps", 10000)
    args.grad_norm = getattr(args, "grad_norm", 0)
    args.seed = getattr(args, "seed", 970903)
    args.no_gold = getattr(args, "no_gold", False)
    args.pretrained = getattr(args, "pretrained", None)
    args.max_lr = getattr(args, "max_lr", 2e-3)
    args.scale = getattr(args, "scale", 1)
    args.train_file = getattr(args, "train_file", "test_set.json")
    args.valid_file = getattr(args, "valid_file", "test_set.json")
    args.max_len = getattr(args, "max_len", 64)
    args.max_num = getattr(args, "max_num", 16)
    args.cand_weight = getattr(args, "cand_weight", 1)
    args.gold_weight = getattr(args, "gold_weight", 1)


def evaluation(args):
    # load data
    base_setting(args)
    tok = AutoTokenizer.from_pretrained(args.model_type)
    collate_fn = partial(collate_mp, pad_token_id=tok.pad_token_id, is_test=True)
    test_set = ReRankingDataset(args.valid_file, args.model_type, is_test=False, maxlen=args.max_len, is_sorted=False, maxnum=args.max_num, task_type=args.task_type, org_query=(args.contextual or args.contextual_wtop1))
    dataloader = DataLoader(test_set, batch_size=args.batch_size, shuffle=False, num_workers=0, collate_fn=collate_fn)
    # build models
    model_path = args.pretrained if args.pretrained is not None else args.model_type
    if args.contextual_wtop1:
        scorer = ContextualTop1ReRanker(model_path, tok.pad_token_id, args.feature_extractor_q, args.feature_extractor_c, post_processing=args.post_processing)
    elif args.contextual:
        scorer = ContextualReRanker(model_path, tok.pad_token_id, args.feature_extractor, post_processing=args.post_processing)
    else:
        scorer = ReRanker(model_path, tok.pad_token_id)
    if args.cuda:
        scorer = scorer.cuda()
    # scorer.load_state_dict(torch.load(os.path.join(args.output_dir, args.model_pt), map_location=f'cuda:{args.gpuid[0]}'))
    state_dict = torch.load(args.model_pt, map_location=f'cuda:{args.gpuid[0]}')
    # if args.contextual:
    #     new_state_dict = {}
    #     for key in state_dict:
    #         if not key.startswith("feature_extractor"):
    #             new_state_dict[key] = state_dict[key]
    #     state_dict = new_state_dict
    scorer.load_state_dict(state_dict)
    scorer.eval()
    model_name = args.model_pt.split("/")[0]

    # loss = 0
    cnt = 0
    acc = 0
    best_rank = 0
    mean_rank = 0
    pred_rank = 0
    top1s = []
    top1s_gt = []
    pred_rank_list = []
    with torch.no_grad():
        for (i, batch) in tqdm.tqdm(enumerate(dataloader), total=len(dataloader), desc="Testing"):
            if args.cuda:
                to_cuda(batch, args.gpuid[0])
            ranks = batch["ranks"]
            random_indices = torch.randperm(ranks.size(0)).to(ranks.device)
            zero_lead = torch.zeros(1, dtype=random_indices.dtype).to(ranks.device)
            shuffler = torch.cat([zero_lead, random_indices+1], dim=0).to(ranks.device)
            unshuffler = torch.argsort(random_indices).to(ranks.device)
            ranks = ranks[random_indices]
            batch["input_ids"] = batch["input_ids"][shuffler]
            if args.fp16:
                with torch.cuda.amp.autocast():
                    output = scorer(batch["input_ids"]) if not args.contextual_wtop1 else scorer(batch["input_ids"], batch["ctx_ids"])
            else:
                output = scorer(batch["input_ids"]) if not args.contextual_wtop1 else scorer(batch["input_ids"], batch["ctx_ids"])
            if 'chuck_sizes' not in batch or len(batch['chuck_sizes']) == 1:
                all_ranks, all_output = [ranks], [output]
            else:
                pivot = 0
                all_ranks, all_output = [], []
                for chuck_size in batch["chuck_sizes"]:
                    all_ranks.append(ranks[pivot:pivot+chuck_size])
                    all_output.append(output[pivot:pivot+chuck_size])
                    pivot += chuck_size

            for ranks, output in zip(all_ranks, all_output):
                top1 = torch.argmin(output, dim=0)
                top1_gt = int(torch.argmin(ranks, dim=0))
                min_rank = torch.min(ranks, dim=0)[0]
                top1_ranks = (min_rank == ranks).nonzero().squeeze(1)
                best_rank += int(ranks[int(top1_gt)])
                mean_rank += torch.mean(ranks, dim=0)
                pred_rank += int(ranks[int(top1)])
                pred_rank_list.append(int(ranks[int(top1)]))
                top1s.append(int(top1))
                top1s_gt.append(int(top1_gt))
                # print(f"top1: {int(top1)} with rank={ranks[int(top1)]}, gt: {top1_gt} with rank={ranks[top1_gt]}")
                if top1 in top1_ranks:
                    acc += 1
                # loss += RankingLoss(output, batch["ranks"], args.margin, args.loss_type)
                cnt += 1
    print(f"accuracy: {acc / cnt:.6f} avg best rank: {best_rank / cnt:.6f} avg mean rank: {mean_rank / cnt:.6f} avg pred rank: {pred_rank / cnt:.6f}")
    pred_rank_list = np.array(pred_rank_list)
    topk = {}
    for k in [1, 5, 10, 20, 50, 100]:
        topk[k] = float(np.mean(pred_rank_list <= k))
        print("Top-%d:"%k, topk[k])
    return top1s, top1s_gt


def test(dataloader, scorer, args, gpuid, target):
    scorer.eval()
    if args.cuda:
        scorer.cuda(gpuid)
    loss = 0
    cnt = 0
    acc = 0
    best_rank = 0
    mean_rank = 0
    pred_rank = 0
    pred_rank_list = []
    best_rank_list = []
    with torch.no_grad():
        for (i, batch) in tqdm.tqdm(enumerate(dataloader), total=len(dataloader), desc="Testing"):
            if args.cuda:
                to_cuda(batch, args.gpuid[0])
            ranks = batch["ranks"]
            if args.fp16:
                with torch.cuda.amp.autocast():
                    output = scorer(batch["input_ids"]) if not args.contextual_wtop1 else scorer(batch["input_ids"], batch["ctx_ids"])
            else:
                output = scorer(batch["input_ids"]) if not args.contextual_wtop1 else scorer(batch["input_ids"], batch["ctx_ids"])
            if 'chuck_sizes' not in batch or len(batch['chuck_sizes']) == 1:
                all_ranks, all_output = [ranks], [output]
            else:
                pivot = 0
                all_ranks, all_output = [], []
                for chuck_size in batch["chuck_sizes"]:
                    all_ranks.append(ranks[pivot:pivot+chuck_size])
                    all_output.append(output[pivot:pivot+chuck_size])
                    pivot += chuck_size

            for ranks, output in zip(all_ranks, all_output):
                top1 = torch.argmin(output, dim=0)
                min_rank = torch.min(ranks, dim=0)[0]
                top1_ranks = (min_rank == ranks).nonzero().squeeze(1)
                best_rank += int(min_rank)
                mean_rank += float(torch.mean(ranks, dim=0))
                pred_rank += int(ranks[int(top1)])
                pred_rank_list.append(int(ranks[int(top1)]))
                best_rank_list.append(int(min_rank))
                if top1 in top1_ranks:
                    acc += 1
                loss += float(RankingLoss(output, ranks, args.margin, args.loss_type))
                cnt += 1
    print(f"\naccuracy: {acc / cnt:.6f} loss: {loss / cnt:.6f} avg best rank: {best_rank / cnt:.6f} avg mean rank: {mean_rank / cnt:.6f} avg pred rank: {pred_rank / cnt:.6f}")
    pred_rank_list = np.array(pred_rank_list)
    best_rank_list = np.array(best_rank_list)
    topk = {}
    for k in [1, 5, 10, 20, 50, 100]:
        topk[f"eval-{target}/top-%d-acc"%k] = float(np.mean(pred_rank_list <= k))
        print("Target %s Top-%d:"%(target, k), topk[f"eval-{target}/top-%d-acc"%k])
    for k in [1, 5, 10, 20, 50, 100]:
        print("Target %s Upper bound - Top-%d:"%(target, k), float(np.mean(best_rank_list <= k)))
    return float(loss / cnt), float(acc / cnt), float(best_rank / cnt), float(mean_rank / cnt), float(pred_rank / cnt), topk

def work(x):
    cache, train_set, x = x
    cache.append(train_set[x])


def run(rank, args):
    base_setting(args)
    torch.manual_seed(args.seed)
    torch.cuda.manual_seed_all(args.seed)
    np.random.seed(args.seed)
    random.seed(args.seed)
    if args.fp16:
        scaler = torch.cuda.amp.GradScaler()

    config = vars(args)
    if args.wandb:
        wandb_run = wandb.init(
            project=args.wandb_project,
            name=args.wandb_name,
            config=config
        )
    gpuid = args.gpuid[rank]
    is_master = rank == 0
    is_mp = len(args.gpuid) > 1
    world_size = len(args.gpuid)
    if is_master:
        id = len(os.listdir(args.output_dir))
        recorder = Recorder(id, args.log, name=args.wandb_name)
    tok = AutoTokenizer.from_pretrained(args.model_type)
    collate_fn = partial(collate_mp, pad_token_id=tok.pad_token_id, is_test=False)
    collate_fn_val = partial(collate_mp, pad_token_id=tok.pad_token_id, is_test=True)
    train_set = ReRankingDataset(args.train_file, args.model_type, is_test=True, maxlen=args.max_len, maxnum=args.max_num, null_rank=args.null_rank, task_type=args.task_type, org_query=(args.contextual or args.contextual_wtop1))
    if False and args.task_type == "title":
        import multiprocessing as mp
        n_workers = 60
        pool = mp.Pool(n_workers)
        print("Multi Processing")
        cache = []
        for _ in tqdm.tqdm(pool.imap_unordered(work, [(cache, train_set, x) for x in range(len(train_set))]), total=len(train_set)):
            pass
        import pickle
        fw = open('train-title.pkl', 'wb')
        pickle.dump(cache, fw)
        #self.cache = pool.map(self.__getitem__, list(range(self.num)))
        #for i in tqdm.trange(self.num, desc="Saving the cache"):
        #    self.cache[i] = self.__getitem__(i)
        #self.cache_done = True
        #del self.data
        import pdb; pdb.set_trace()


    if args.valid_files is None:
        target, valid_file = args.valid_file.split(':')
        val_sets = [(target, ReRankingDataset(valid_file, args.model_type, is_test=True, maxlen=512, is_sorted=False, maxnum=args.max_num, task_type=args.task_type, org_query=(args.contextual or args.contextual_wtop1)))]
    else:
        val_sets = []
        for valid_file in args.valid_files.split(','):
            target, valid_file = valid_file.split(':')
            val_sets.append((target, ReRankingDataset(valid_file, args.model_type, is_test=True, maxlen=512, is_sorted=False, maxnum=args.max_num, task_type=args.task_type, org_query=(args.contextual or args.contextual_wtop1))))
    if is_mp:
        train_sampler = torch.utils.data.distributed.DistributedSampler(
    	 train_set, num_replicas=world_size, rank=rank, shuffle=True)
        dataloader = DataLoader(train_set, batch_size=args.batch_size, shuffle=False, num_workers=8, collate_fn=collate_fn, sampler=train_sampler)
        val_samplers = [torch.utils.data.distributed.DistributedSampler(
    	 val_set, num_replicas=world_size, rank=rank) for val_set in val_sets]
        val_dataloaders = [(target, DataLoader(val_set, batch_size=8, shuffle=False, num_workers=8, collate_fn=collate_fn_val, sampler=val_samplers[i])) for i, (target, val_set) in enumerate(val_sets)]
    else:
        dataloader = DataLoader(train_set, batch_size=args.batch_size, shuffle=True, num_workers=8, collate_fn=collate_fn)
        val_dataloaders = [(target, DataLoader(val_set, batch_size=args.eval_batch_size, shuffle=False, num_workers=8, collate_fn=collate_fn_val)) for target, val_set in val_sets]
    # build models
    model_path = args.pretrained if args.pretrained is not None else args.model_type
    if args.contextual_wtop1:
        scorer = ContextualTop1ReRanker(model_path, tok.pad_token_id, args.feature_extractor_q, args.feature_extractor_c, post_processing=args.post_processing)
    elif args.contextual:
        scorer = ContextualReRanker(model_path, tok.pad_token_id, args.feature_extractor, post_processing=args.post_processing)
    else:
        scorer = ReRanker(model_path, tok.pad_token_id)
    # scorer = ReRanker(model_path, tok.pad_token_id) if not args.contextual else ContextualReRanker(model_path, tok.pad_token_id, args.feature_extractor)
    if len(args.model_pt) > 0:
        scorer.load_state_dict(torch.load(os.path.join(args.output_dir, args.model_pt), map_location=f'cuda:{gpuid}'))
    if args.cuda:
        if len(args.gpuid) == 1:
            scorer = scorer.cuda()
        else:
            dist.init_process_group("nccl", rank=rank, world_size=world_size)
            scorer = nn.parallel.DistributedDataParallel(scorer.to(gpuid), [gpuid], find_unused_parameters=True)
    scorer.train()
    init_lr = args.max_lr / args.warmup_steps
    s_optimizer = optim.Adam(scorer.parameters(), lr=init_lr)
    if is_master:
        recorder.write_config(args, [scorer], __file__)
    min_avg_ranks = {target:10000 for target, _ in val_sets}
    best_scores = {target:None for target, _ in val_sets}
    best_selected_scores = {target:None for target, _ in val_sets if not target.startswith('dev')}
    all_step_cnt = 0
    # start training
    for epoch in range(args.epoch):
        s_optimizer.zero_grad()
        step_cnt = 0
        sim_step = 0
        avg_loss = 0
        print(f"Epoch {epoch}")
        progress_bar = tqdm.tqdm(enumerate(dataloader), total=len(dataloader), desc="Training epoch {}".format(epoch))
        eval_steps = len(dataloader)//args.eval_per_epoch
        did_first_eval = False
        for (i, batch) in progress_bar:
            if args.cuda:
                to_cuda(batch, gpuid)
            step_cnt += 1
            try:
                if args.fp16:
                    with torch.cuda.amp.autocast():
                        output = scorer(batch["input_ids"]) if not args.contextual_wtop1 else scorer(batch["input_ids"], batch["ctx_ids"])
                        if args.batch_size == 1:
                            loss = args.scale * RankingLoss(output, batch["ranks"], args.margin, args.loss_type)
                        else:
                            pivot = 0
                            loss = 0.0
                            for chuck_size in batch["chuck_sizes"]:
                                loss += args.scale * RankingLoss(output[pivot:pivot+chuck_size], batch["ranks"][pivot:pivot+chuck_size], args.margin, args.loss_type)
                                pivot += chuck_size
                            loss = loss / len(batch["chuck_sizes"])
                else:
                    output = scorer(batch["input_ids"]) if not args.contextual_wtop1 else scorer(batch["input_ids"], batch["ctx_ids"])
                    if args.batch_size == 1:
                        loss = args.scale * RankingLoss(output, batch["ranks"], args.margin, args.loss_type)
                    else:
                        pivot = 0
                        loss = 0.0
                        for chuck_size in batch["chuck_sizes"]:
                            loss += args.scale * RankingLoss(output[pivot:pivot+chuck_size], batch["ranks"][pivot:pivot+chuck_size], args.margin, args.loss_type)
                            pivot += chuck_size
                        loss = loss / len(batch["chuck_sizes"])
                    loss = loss / args.accumulate_step
                avg_loss += float(loss.item())
                if args.fp16:
                    scaler.scale(loss).backward()
                else:
                    loss.backward()
            except RuntimeError as e:
                if "out of memory" in str(e):
                    print(
                        "attempting to recover from OOM in forward/backward pass"
                    )
                    s_optimizer.zero_grad()
                    if args.cuda:
                        torch.cuda.empty_cache()
                    print("skip batch (%d)"%i)
                    continue
                else:
                    raise e
            progress_bar.set_postfix({"loss": loss.item()})
            if args.wandb:
                wandb.log({"train/loss": loss.item()})
            if step_cnt == args.accumulate_step:
                # optimize step      
                if args.grad_norm > 0:
                    nn.utils.clip_grad_norm_(scorer.parameters(), args.grad_norm)
                step_cnt = 0
                sim_step += 1
                all_step_cnt += 1
                lr = args.max_lr * min(all_step_cnt ** (-0.5), all_step_cnt * (args.warmup_steps ** (-1.5)))
                for param_group in s_optimizer.param_groups:
                    param_group['lr'] = lr
                if args.fp16:
                    scaler.step(s_optimizer)
                else:
                    s_optimizer.step()
                s_optimizer.zero_grad()
                scaler.update()
            if sim_step % args.report_freq == 0 and step_cnt == 0 and is_master:
                print("id: %d"%id)
                recorder.print("epoch: %d, batch: %d, avg loss: %.6f"%(epoch+1, sim_step, 
                 avg_loss / args.report_freq))
                recorder.print(f"learning rate: {lr:.6f}")
                recorder.plot("loss", {"loss": avg_loss / args.report_freq}, all_step_cnt)
                recorder.print()
                avg_loss = 0

            # if all_step_cnt % 1000 == 0: # and step_cnt == 0 all_step_cnt != 0 and 
            if ((i % eval_steps == 0 and i>0) or i == len(dataloader)-1) or (not did_first_eval and not args.skip_eval): # and step_cnt == 0 all_step_cnt != 0 and 
                curr_scores = {}
                is_dev_best = []
                for target, val_dataloader in val_dataloaders:
                    loss, acc, best_rank, mean_rank, pred_rank, topk = test(val_dataloader, scorer, args, gpuid, target)
                    did_first_eval = True
                    curr_scores[target] = topk
                    if args.wandb:
                        d = {f"eval-{target}/loss": float(loss),
                             f"eval-{target}/acc": float(acc),
                             f"eval-{target}/best_rank": float(best_rank),
                             f"eval-{target}/mean_rank": float(mean_rank),
                             f"eval-{target}/pred_rank": float(pred_rank)}
                        for k in topk:
                            d[k] = topk[k]
                        wandb.log(d) 
                    if pred_rank < min_avg_ranks[target] and is_master:
                        min_avg_ranks[target] = pred_rank
                        best_scores[target] = topk
                        if target.startswith('dev'):
                            is_dev_best.append(target)
                        if is_mp:
                            recorder.save(scorer.module, f"scorer-{target}-best.bin")
                        else:
                            recorder.save(scorer, f"scorer-{target}-best.bin")
                        recorder.print("eval target: %s"%target)
                        recorder.print("best - epoch: %d, batch: %d"%(epoch, i / args.accumulate_step))
                        recorder.print("best - loss: %.6f, acc: %.6f, best rank: %.6f, mean rank: %.6f, pred rank: %.6f"%(loss, acc, best_rank, mean_rank, pred_rank))
                        print(target + ', ' + ', '.join([str(best_scores[target][f"eval-{target}/top-%d-acc"%k]) for k in [1,5,10,20,50,100]]), flush=True)
                for dev_target in is_dev_best:
                    target = dev_target[4:]
                    best_selected_scores[target] = curr_scores[target]
                print("current best scores - epoch: %d, batch: %d"%(epoch, i / args.accumulate_step), flush=True)
                for target in best_scores:
                    print(target + ', ' + ', '.join([str(best_scores[target][f"eval-{target}/top-%d-acc"%k]) for k in [1,5,10,20,50,100]]), flush=True)
                print("current best selected scores - epoch: %d, batch: %d"%(epoch, i / args.accumulate_step), flush=True)
                for target in best_selected_scores:
                    print(target + ', ' + ', '.join([str(best_selected_scores[target][f"eval-{target}/top-%d-acc"%k]) for k in [1,5,10,20,50,100]]), flush=True)
                recorder.save(scorer, "scorer-last.bin")
                recorder.save(s_optimizer, "optimizer.bin")
                # if is_master:
                #     recorder.print("val rouge: %.6f"%(1 - loss))
    recorder.save(scorer, "scorer-last.bin")
    torch.save(scorer.state_dict(), os.path.join(args.output_dir, args.wandb_name+"-scorer-last.bin"))
    print("final best scores - epoch: %d, batch: %d"%(epoch, i / args.accumulate_step), flush=True)
    for target in best_scores:
        print(target + ', ' + ', '.join([str(best_scores[target][f"eval-{target}/top-%d-acc"%k]) for k in [1,5,10,20,50,100]]), flush=True)
    print("final best selected scores - epoch: %d, batch: %d"%(epoch, i / args.accumulate_step), flush=True)
    for target in best_selected_scores:
        print(target + ', ' + ', '.join([str(best_selected_scores[target][f"eval-{target}/top-%d-acc"%k]) for k in [1,5,10,20,50,100]]), flush=True)

def main(args):
    # set env
    if len(args.gpuid) > 1:
        os.environ['MASTER_ADDR'] = 'localhost'
        os.environ['MASTER_PORT'] = f'{args.port}'
        mp.spawn(run, args=(args,), nprocs=len(args.gpuid), join=True)
    else:
        run(0, args)

if __name__ ==  "__main__":
    parser = argparse.ArgumentParser(description='Training Parameter')
    parser.add_argument("--cuda", action="store_true")
    parser.add_argument("--wtop1", action="store_true")
    parser.add_argument("--fp16", action="store_true")
    parser.add_argument("--skip_eval", action="store_true")
    parser.add_argument("--contextual", action="store_true")
    parser.add_argument("--contextual_wtop1", action="store_true")
    parser.add_argument("--wandb", action="store_true")
    parser.add_argument("--wandb_project", default='ugar', type=str)
    parser.add_argument("--wandb_name", default='roberta-base', type=str)
    parser.add_argument("--model_type", default='roberta-base', type=str)
    parser.add_argument("--feature_extractor", default='facebook/dpr-question_encoder-single-nq-base', type=str)
    parser.add_argument("--feature_extractor_q", default='facebook/dpr-question_encoder-single-nq-base', type=str)
    parser.add_argument("--feature_extractor_c", default='facebook/dpr-ctx_encoder-single-nq-base', type=str)
    parser.add_argument("--gpuid", nargs='+', type=int, default=0)
    parser.add_argument("--n_workers", type=int, default=32)
    parser.add_argument("--max_len", type=int, default=64)
    parser.add_argument("--margin", type=float, default=0.01)
    parser.add_argument("--max_lr", type=float, default=2e-3)
    parser.add_argument("--epoch", type=int, default=20)
    parser.add_argument("--warmup_steps", type=int, default=10000)
    parser.add_argument("--batch_size", type=int, default=1)
    parser.add_argument("--eval_batch_size", type=int, default=1)
    parser.add_argument("--accumulate_step", type=int, default=12)
    parser.add_argument("--eval_per_epoch", type=int, default=70)
    parser.add_argument("-e", "--evaluate", action="store_true")
    parser.add_argument("-l", "--log", action="store_true")
    parser.add_argument("-p", "--port", type=int, default=12355)
    parser.add_argument("--model_pt", default="", type=str)
    parser.add_argument("--task_type", default="", type=str)
    parser.add_argument("--post_processing", default="diff", type=str)
    parser.add_argument("--encode_mode", default=None, type=str)
    parser.add_argument("--train_file", default=None, type=str)
    parser.add_argument("--valid_file", default=None, type=str)
    parser.add_argument("--valid_files", default=None, type=str)
    parser.add_argument("--pretrained", default=None, type=str)
    parser.add_argument("--bm25_dir", default="output_t0_gen_bm25", type=str)
    parser.add_argument("--output_dir", default="./cache", type=str)
    parser.add_argument("--rerank_output", default="reranked_result.json", type=str)
    parser.add_argument("--loss_type", default='weight-divide', type=str)
    parser.add_argument("--null_rank", type=int, default=101)
    parser.add_argument("--skip_eval_first", action="store_true")
    args = parser.parse_args()
    if args.cuda is False:
        if args.evaluate:
            result = evaluation(args)
        else:
            main(args)
    else:
        if args.evaluate:
            with torch.cuda.device(args.gpuid[0]):
                result, result_gt = evaluation(args)
                if args.bm25_dir is not None:
                    reranked_result = []
                    reranked_result_gt = []
                    if args.n_workers <= 1:
                        for i, r in tqdm.tqdm(enumerate(result), total=len(result), desc="Reranking"):
                            # open the bm25 file and get the r-th item
                            with open(os.path.join(args.bm25_dir, f"nq-test-{i}/results.json")) as f:
                                bm25_result = json.load(f)
                                reranked_result.append(bm25_result[r+1])
                                reranked_result_gt.append(bm25_result[result_gt[i]+1])
                    else:
                        def rerank(i, r):
                            with open(os.path.join(args.bm25_dir, f"nq-test-{i}/results.json")) as f:
                                bm25_result = json.load(f)
                                return bm25_result[r+1]
                        pool = mp.Pool(args.n_workers)
                        reranked_result = pool.starmap(rerank, enumerate(result))
                        pool.close()
                        pool.join()

                        pool = mp.Pool(args.n_workers)
                        reranked_result_gt = pool.starmap(rerank, enumerate(result_gt))
                        pool.close()
                        pool.join()
                    # write to file
                    with open(args.rerank_output, "w") as f:
                        json.dump(reranked_result, f, indent=4)
                    with open(args.rerank_output.replace(".json", "_gt.json"), "w") as f:
                        json.dump(reranked_result_gt, f, indent=4)
        elif len(args.gpuid) == 1:    
            with torch.cuda.device(args.gpuid[0]):
                main(args)
        else:
            main(args)