import dataclasses
import json
import random
from enum import Enum, auto
from pathlib import Path
from typing import List, Any
import numpy as np
import pandas as pd
import torch
from torch.utils.data import DataLoader, Dataset
from tqdm import tqdm
from local_config import PATH_TO_MIMIC_NLE
class SeparatorStyle(Enum):
"""Different separator style."""
SINGLE = auto()
TWO = auto()
@dataclasses.dataclass
class Conversation:
"""A class that keeps all conversation history."""
system: str
roles: List[str]
messages: List[List[str]]
offset: int
sep_style: SeparatorStyle = SeparatorStyle.SINGLE
sep: str = "###"
sep2: str = None
# Used for gradio server
skip_next: bool = False
conv_id: Any = None
def get_prompt(self):
if self.sep_style == SeparatorStyle.SINGLE:
ret = self.system
for role, message in self.messages:
if message:
ret += self.sep + " " + role + ": " + message
else:
ret += self.sep + " " + role + ":"
return ret
elif self.sep_style == SeparatorStyle.TWO:
seps = [self.sep, self.sep2]
ret = self.system + seps[0]
for i, (role, message) in enumerate(self.messages):
if message:
ret += role + ": " + message + seps[i % 2]
else:
ret += role + ":"
return ret
else:
raise ValueError(f"Invalid style: {self.sep_style}")
def append_message(self, role, message):
self.messages.append([role, message])
def dict(self):
return {
"system": self.system,
"roles": self.roles,
"messages": self.messages,
"offset": self.offset,
"sep": self.sep,
"sep2": self.sep2,
"conv_id": self.conv_id,
}
def create_conv():
conv = Conversation(
system="A chat between a curious user and an artificial intelligence assistant acting as an experienced radiologist. "
"The assistant gives professional, detailed, and polite answers to the user's questions.",
roles=["USER", "ASSISTANT"],
messages=[],
offset=0,
sep_style=SeparatorStyle.TWO,
sep=" ",
sep2="</s>",
)
return conv
def create_direct_task_data(lang_model, tokenizer, val_dataset, task_name):
prompts = pd.read_csv(f"data/instruct_prompts/{task_name}_prompts.csv")["instruction"].tolist()
data_loader = DataLoader(val_dataset, batch_size=12, shuffle=False, num_workers=0)
report_jsons = []
print("Dataloader len: ", len(data_loader))
for _, batch in tqdm(enumerate(data_loader)):
# Create prompts for every report
# sample batchsize questions from EL_prompts
batch_prompts = random.choices(prompts, k=len(batch["text_input"]))
batch_instructions = []
for text_target, prompt in zip(batch["text_target"], batch_prompts):
conv = create_conv()
conv.append_message(conv.roles[0], "Report: " + text_target + "\n" + prompt)
conv.append_message(conv.roles[1], None)
batch_instructions.append(conv.get_prompt())
inputs = tokenizer.batch_encode_plus(batch_instructions, return_tensors="pt", padding=True)
input_ids = inputs["input_ids"].to(torch.device("cuda"))
# generate answers with no-lora vicuna
generation_output = lang_model.generate(
input_ids=input_ids,
dicom=None,
return_dict_in_generate=True,
output_scores=True,
max_new_tokens=256
)
preds = tokenizer.batch_decode(generation_output.sequences, skip_special_tokens=True)
preds = [p.split("ASSISTANT:")[1] for idx, p in enumerate(preds)]
# iterate over batch elements
for i in range(len(batch["text_input"])):
text_target = batch["text_target"][i] # GT report
task_prompt = batch_prompts[i]
task_instruction = batch_instructions[i]
answer = preds[i]
dicom = batch["dicom"][i]
# sample random prompt for every report
reports_json = {
"gt_report": text_target,
"task": task_prompt,
"instruction": task_instruction,
"input": "",
"output": answer,
"dicom": dicom,
"task_type": task_name
}
report_jsons.append(reports_json)
# save
with open(f"data/large_instruct_data/instruct_large_{task_name}.json", "w") as f:
json.dump(report_jsons, f, ensure_ascii=False, indent=4)
def create_cp_task_data(val_dataset, task_name):
prompts = pd.read_csv(f"data/instruct_prompts/{task_name}_prompts.csv")["instruction"].tolist()
data_loader = DataLoader(val_dataset, batch_size=200, shuffle=False, num_workers=200)
report_jsons = []
for _, batch in tqdm(enumerate(data_loader)):
# Create prompts for every report
# sample batchsize questions from EL_prompts
batch_prompts = random.choices(prompts, k=len(batch["text_input"]))
# iterate over batch elements
for i in range(len(batch["text_input"])):
text_target = batch["text_target"][i] # GT report
task_prompt = batch_prompts[i]
cp_indices = np.where(batch["chexpert_labels"][i] == 1.)
cp_findings = [val_dataset.dataset.dataset.chexpert_cols[i] for i in cp_indices[0]]
if task_name == "CPbQA": # binary QA
if "No Finding" in cp_findings:
cp_findings.remove("No Finding")
# 50% sample finding from cp_findings, 50% sample finding from val_dataset.dataset.dataset.chexpert_cols - cp_findings
if random.random() < 0.6 and len(cp_findings) > 0:
finding = random.choice(cp_findings) # answer: yes
answer = 'yes'
else:
finding = random.choice(list(set(val_dataset.dataset.dataset.chexpert_cols[1:]) - set(cp_findings))) # answer: no
answer = 'no'
task_prompt = task_prompt.replace("<X>", finding)
elif task_name == "CPaQA": # give all findings
answer = ', '.join(cp_findings)
dicom = batch["dicom"][i]
# sample random prompt for every report
reports_json = {
"gt_report": text_target,
"task": task_prompt,
"input": "",
"output": answer,
"dicom": dicom,
"task_type": task_name
}
report_jsons.append(reports_json)
# save
with open(f"data/large_instruct_data/instruct_large_{task_name}.json", "w") as f:
json.dump(report_jsons, f, ensure_ascii=False, indent=4)
class CorrectionDataset(Dataset):
def __init__(self, data):
self.data = data
def __len__(self):
return len(self.data)
def __getitem__(self, idx):
sample = self.data[idx]
fp = sample["fp"]
fn = sample["fn"]
fp_str = ', '.join(fp)
fp_str = fp_str.rsplit(', ', 1)
fp_str = ' and '.join(fp_str)
fn_str = ', '.join(fn)
fn_str = fn_str.rsplit(', ', 1)
fn_str = ' and '.join(fn_str)
gt_report = sample["gt_report"]
pred_report = sample["pred_report"]
dicom = sample["dicom"]
return {'gt_report': gt_report, 'pred_report': pred_report, 'fp': fp_str, 'fn': fn_str, 'dicom': dicom}
def create_correction_task_data(lang_model, tokenizer):
# load correction json
with open("data/instruct_prompts/instruct_task_correction_preds.json") as f:
correction_preds = json.load(f)
# create pytorch dataset from json
correction_dataset = CorrectionDataset(correction_preds)
data_loader = DataLoader(correction_dataset, batch_size=12, shuffle=False, num_workers=12)
prompts_both = pd.read_csv(f"data/instruct_prompts/CO_both_prompts.csv")["instruction"].tolist()
prompts_add = pd.read_csv(f"data/instruct_prompts/CO_add_prompts.csv")["instruction"].tolist()
prompts_rem = pd.read_csv(f"data/instruct_prompts/CO_rem_prompts.csv")["instruction"].tolist()
report_jsons = []
for _, batch in tqdm(enumerate(data_loader)):
# use very clear, fixed prompt for data generation -> in training use random prompts
fixed_batch_prompts = []
for fp, fn in zip(batch["fp"], batch["fn"]):
fixed_corr_prompt = "Please provide an adapted report. "
if fp != "":
fixed_corr_prompt += f"Do not mention {fp}. "
if fn != "":
fixed_corr_prompt += f"Mention {fn}. "
if fp == "" and fn == "":
fixed_corr_prompt = "NOCHANGE"
fixed_batch_prompts.append(fixed_corr_prompt.strip())
batch_prompts = []
for fp, fn in zip(batch["fp"], batch["fn"]):
if fp == "" and fn == "":
batch_prompts.append("NOCHANGE")
elif fp == "":
batch_prompts.append(random.choice(prompts_add).replace("<add>", fn))
elif fn == "":
batch_prompts.append(random.choice(prompts_rem).replace("<rem>", fp))
else:
batch_prompts.append(random.choice(prompts_both).replace("<add>", fn).replace("<rem>", fp))
batch_instructions = []
for pred_report, prompt in zip(batch["pred_report"], fixed_batch_prompts):
conv = create_conv()
conv.append_message(conv.roles[0], "Please write a radiology report for the given x-ray.")
conv.append_message(conv.roles[1], pred_report)
conv.append_message(conv.roles[0], prompt)
conv.append_message(conv.roles[1], None)
batch_instructions.append(conv.get_prompt())
inputs = tokenizer.batch_encode_plus(batch_instructions, return_tensors="pt", padding=True)
input_ids = inputs["input_ids"].to(torch.device("cuda"))
# generate answers with no-lora vicuna
generation_output = lang_model.generate(
input_ids=input_ids,
dicom=None,
return_dict_in_generate=True,
output_scores=True,
max_new_tokens=256
)
preds = tokenizer.batch_decode(generation_output.sequences, skip_special_tokens=True)
preds = [p.split("ASSISTANT:")[-1].strip() for idx, p in enumerate(preds)]
# iterate over batch elements
for i in range(len(batch["pred_report"])):
gt_report = batch["gt_report"][i] # GT report
incorrect_report = batch["pred_report"][i] # predicted report that will be corrected
task_prompt = batch_prompts[i]
task_instruction = batch_instructions[i]
answer = preds[i]
dicom = batch["dicom"][i]
if task_prompt == "NOCHANGE":
continue # we don't want to train for correction on already correct reports
# sample random prompt for every report
reports_json = {
"gt_report": gt_report,
"incorrect_report": incorrect_report,
"task": task_prompt,
"instruction": task_instruction,
"input": "",
"output": answer,
"dicom": dicom,
"task_type": 'CO'
}
report_jsons.append(reports_json)
# save
with open(f"data/large_instruct_data/instruct_large_CO.json", "w") as f:
json.dump(report_jsons, f, ensure_ascii=False, indent=4)
def create_nle_task_data():
MIMIC_DIAGNOSISLIST = ['Atelectasis', 'Consolidation', 'Edema', 'Enlarged Cardiomediastinum', 'Lung Lesion', 'Lung Opacity', 'Pleural Effusion',
'Pleural Other', 'Pneumonia', 'Pneumothorax']
# load mimic_nle json
mimic_nle = []
with open(f'{PATH_TO_MIMIC_NLE}/mimic-nle/mimic-nle-train.json', 'r') as f:
for line in f:
obj = json.loads(line)
mimic_nle.append(obj)
prompts = pd.read_csv(f"data/instruct_prompts/RE_prompts.csv")["instruction"].tolist()
report_jsons = []
reports = pd.read_csv('mimic-cxr/reports_processed/mimic_cxr_sectioned.csv')
reports = reports.dropna(subset=['findings'])
reports['findings'] = reports['findings'].apply(lambda x: x.replace('\n', ''))
for sample in tqdm(mimic_nle):
report_id = sample["report_ID"]
gt_report = reports[reports["Note_file"] == f"{report_id}.txt"]["findings"].tolist()
if len(gt_report) == 0: # report did have no findings section
continue
gt_report = gt_report[0]
nle = sample['nle']
if nle not in gt_report: # sort out samples that reference the impression instead of the findings section
continue
dicom = reports[reports["Note_file"] == f"{report_id}.txt"]["dicom_id"].tolist()[0]
task_prompt = random.choice(prompts)
diagnoses = [d for idx, d in enumerate(MIMIC_DIAGNOSISLIST) if sample["diagnosis_label"][idx] == 1]
diagnoses_string = ", ".join(diagnoses)
diagnoses_string = diagnoses_string.rsplit(', ', 1)
diagnoses_string = ' and '.join(diagnoses_string)
task_prompt = task_prompt.replace("<X>", diagnoses_string)
# sample random prompt for every report
reports_json = {
"gt_report": gt_report,
"task": task_prompt,
"input": "",
"output": sample['nle'],
"dicom": dicom,
"task_type": 'RE'
}
report_jsons.append(reports_json)
# save
print(len(report_jsons))
with open(f"data/large_instruct_data/instruct_large_RE.json", "w") as f:
json.dump(report_jsons, f, ensure_ascii=False, indent=4)
