import pickle, json, sys

# Usage: python to_json.py input.pkl output.json

with open(sys.argv[1], "rb") as f:
    data = pickle.load(f)

ks = ["test_span_table", "predict_span_table", "predict_sentence_table"]
for d in data:
    for k in ks:
        if k in d:
            d[k] = d[k].tolist()

with open(sys.argv[2], "w") as f:
    json.dump(data, f, ensure_ascii=False)