Spaces:
Running
Running
File size: 1,721 Bytes
50e583f |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 |
__author__ = "qiao"
"""
Running the TrialGPT matching for three cohorts (sigir, TREC 2021, TREC 2022).
"""
import json
from nltk.tokenize import sent_tokenize
import os
import sys
from tqdm import tqdm
from TrialGPT import trialgpt_matching
if __name__ == "__main__":
corpus = sys.argv[1]
model = sys.argv[2]
dataset = json.load(open(f"dataset/{corpus}/retrieved_trials.json"))
output_path = f"results/matching_results_{corpus}_{model}.json"
# Dict{Str(patient_id): Dict{Str(label): Dict{Str(trial_id): Str(output)}}}
if os.path.exists(output_path):
output = json.load(open(output_path))
else:
output = {}
for instance in tqdm(dataset):
# Dict{'patient': Str(patient), '0': Str(NCTID), ...}
patient_id = instance["patient_id"]
patient = instance["patient"]
sents = sent_tokenize(patient)
sents.append("The patient will provide informed consent, and will comply with the trial protocol without any practical issues.")
sents = [f"{idx}. {sent}" for idx, sent in enumerate(sents)]
patient = "\n".join(sents)
# initialize the patient id in the output
if patient_id not in output:
output[patient_id] = {"0": {}, "1": {}, "2": {}}
for label in ["2", "1", "0"]:
if label not in instance: continue
for trial in instance[label]:
trial_id = trial["NCTID"]
# already calculated and cached
if trial_id in output[patient_id][label]:
continue
# in case anything goes wrong (e.g., API calling errors)
try:
results = trialgpt_matching(trial, patient, model)
output[patient_id][label][trial_id] = results
with open(output_path, "w") as f:
json.dump(output, f, indent=4)
except Exception as e:
print(e)
continue
|