get_submission.py

import os
import sys
import argparse
import json
import numpy as np

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument("--test_data", default="./data/test.jsonl", help="path to the test data file", required=False)
    parser.add_argument("--preds", default="./saved/_preds.npy",
                        help="path to the prediction file generated by the run_infer.sh script")
    parser.add_argument("--output", default="./saved/results.jsonl", help="path to the output file")
    args = parser.parse_args()
    preds = np.load(args.preds)
    fout = open(args.output, "w")
    with open(args.test_data, "r") as fin:
        lines = fin.readlines()
        Cnt = 0
        for line in lines:
            data = json.loads(line)
            res = {"id": data['id']}
            tmp = []
            for mention in data['candidates']:
                tmp.append({"id": mention["id"], "type_id": int(preds[Cnt])})
                Cnt += 1
            res["predictions"] = tmp
            fout.write(json.dumps(res) + "\n")
        assert Cnt == len(preds)
    fout.close()