rerank.py

import argparse
import train
import data
import torch

import torch_xla.core.xla_model as xm
device = xm.xla_device()
# device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')

print('device in modeling.py:', device)


def main_cli():
    parser = argparse.ArgumentParser('CEDR model re-ranking')
    parser.add_argument('--model', choices=train.MODEL_MAP.keys(), default='vanilla_bert')
    parser.add_argument('--datafiles', type=argparse.FileType('rt'), nargs='+')
    parser.add_argument('--run', type=argparse.FileType('rt'))
    parser.add_argument('--model_weights', type=argparse.FileType('rb'))
    parser.add_argument('--out_path', type=argparse.FileType('wt'))
    args = parser.parse_args()
    # model = train.MODEL_MAP[args.model]().to(device)
    model = train.MODEL_MAP[args.model]()
    dataset = data.read_datafiles(args.datafiles)
    run = data.read_run_dict(args.run)
    if args.model_weights is not None:
        # model.load(args.model_weights.name, map_location=device)
        model.load(args.model_weights.name, map_location='cpu')
    model = model.to(device)
    train.run_model(model, dataset, run, args.out_path.name, desc='rerank')


if __name__ == '__main__':
    main_cli()