import json
import os
from typing import Dict, List
from pathlib import Path
from datetime import datetime as dt
import logging
from tenacity import (
retry,
stop_after_attempt,
wait_random_exponential,
)
from openai import OpenAI
import google.generativeai as genai
from langchain.callbacks.manager import CallbackManager
from langchain.callbacks.streaming_stdout import StreamingStdOutCallbackHandler
from langchain.llms import Ollama
import pandas as pd
from config.config import *
from prompts.prompt import *
logging.basicConfig(filename=f'logs/{dt.now().strftime("%Y%m%d")}.log', level=logging.INFO, format='%(asctime)s\n%(message)s')
@retry(wait=wait_random_exponential(min=1, max=60), stop=stop_after_attempt(6))
def query_llm(
model: str,
llm,
systemPrompt: str,
userPrompt: str,
):
if model in ['gpt-4-1106-preview', 'gpt-3.5-turbo', 'gpt-3.5-turbo-16k', 'gpt-3.5-turbo-1106']:
try:
result = llm.chat.completions.create(
model=model,
messages=[
{'role': 'system', 'content': systemPrompt},
{'role': 'user', 'content': userPrompt},
],
)
except Exception as e:
logging.info(f'{e}')
raise e
return result.choices[0].message.content, result.usage.prompt_tokens, result.usage.completion_tokens
elif model in ['gemini-pro']:
try:
response = llm.generate_content(systemPrompt + userPrompt)
except Exception as e:
logging.info(f'{e}')
raise e
return response.text, 0, 0
elif model in ['llama2:70b']:
try:
response = llm(systemPrompt + userPrompt)
except Exception as e:
logging.info(f'{e}')
raise e
return response, 0, 0
def format_input(
patient: List,
dataset: str,
form: str,
features: List[str],
):
feature_values = {}
numerical_features = ['Diastolic blood pressure', 'Fraction inspired oxygen', 'Glucose', 'Heart Rate', 'Height', 'Mean blood pressure', 'Oxygen saturation', 'Respiratory rate', 'Systolic blood pressure', 'Temperature', 'Weight', 'pH']
categorical_features = ['Capillary refill rate', 'Glascow coma scale eye opening', 'Glascow coma scale motor response', 'Glascow coma scale total', 'Glascow coma scale verbal response']
if dataset == 'mimic-iv':
for i, feature in enumerate(features):
if feature in numerical_features:
feature_values[feature] = [str(visit[i]) for visit in patient]
for categorical_feature in categorical_features:
indexes = [i for i, f in enumerate(features) if f.startswith(categorical_feature)]
feature_values[categorical_feature] = []
for visit in patient:
values = [visit[i] for i in indexes]
if 1 not in values:
feature_values[categorical_feature].append('unknown')
else:
for i in indexes:
if visit[i] == 1:
feature_values[categorical_feature].append(features[i].split('->')[-1])
break
features = categorical_features + numerical_features
elif dataset == 'tjh':
for i, feature in enumerate(features):
feature_values[feature] = [str(visit[i]) for visit in patient]
detail = ''
if form == 'string':
for feature in features:
detail += f'- {feature}: \"{", ".join(feature_values[feature])}\"\n'
elif form == 'list':
for feature in features:
detail += f'- {feature}: [{", ".join(feature_values[feature])}]\n'
elif form == 'batches':
for i, visit in enumerate(patient):
detail += f'Visit {i + 1}:\n'
for feature in features:
value = feature_values[feature][i] if i < len(feature_values[feature]) else 'unknown'
detail += f'- {feature}: {value}\n'
detail += '\n'
return detail
def run(
config: Dict,
output_logits: bool=True,
output_prompts: bool=False,
logits_root: str='logits',
prompts_root: str='logs',
):
logging.info(f'Running config: {config}\n\n')
prompt_tokens = 0
completion_tokens = 0
dataset = config['dataset']
assert dataset in ['tjh', 'mimic-iv'], f'Unknown dataset: {dataset}'
task = config['task']
assert task in ['outcome', 'los', 'readmission', 'multitask'], f'Unknown task: {task}'
time = config['time']
if time == 0:
time_des = 'upon-discharge'
elif time == 1:
time_des = '1month'
elif time == 2:
time_des = '6months'
else:
raise ValueError(f'Unknown time: {time}')
if config['unit'] is True or config['reference_range'] is True:
unit_range = ''
unit_values = dict(json.load(open(UNIT[dataset])))
range_values = dict(json.load(open(REFERENCE_RANGE[dataset])))
for feature in unit_values.keys():
unit_range += f'- {feature}: '
if config['unit'] is True:
unit_range = unit_range + unit_values[feature] + ' '
if config['reference_range'] is True:
unit_range = unit_range + range_values[feature]
unit_range += '\n'
else:
unit_range = ''
form = config['form']
assert form in ['string', 'batches', 'list'], f'Unknown form: {form}'
nshot = config['n_shot']
if nshot == 0:
example = ''
elif nshot == 1:
example = f'Here is an example of input information:\n'
example += 'Example #1:'
example += EXAMPLE[dataset][task][0] + '\n'
else:
example = f'Here are {nshot} examples of input information:\n'
for i in range(nshot):
example += f'Example #{i + 1}:'
example += EXAMPLE[dataset][task][i] + '\n'
if config.get('prompt_engineering') is True:
example = COT[dataset]
response_format = RESPONSE_FORMAT['cot']
else:
response_format = RESPONSE_FORMAT[task]
if task == 'outcome':
task_description = TASK_DESCRIPTION_AND_RESPONSE_FORMAT[task][time_des]
else:
task_description = TASK_DESCRIPTION_AND_RESPONSE_FORMAT[task]
model = config['model']
if model in ['gpt-4-1106-preview', 'gpt-3.5-turbo', 'gpt-3.5-turbo-16k', 'gpt-3.5-turbo-1106']:
llm = OpenAI(api_key=OPENAI_API_KEY)
elif model in ['gemini-pro']:
genai.configure(api_key=GOOGLE_API_KEY, transport='rest')
llm = genai.GenerativeModel(model)
elif model in ['llama2:70b']:
llm = Ollama(model=model)
else:
raise ValueError(f'Unknown model: {model}')
dataset_path = f'datasets/{dataset}/processed/fold_llm'
impute = config.get('impute', 1)
if impute in [1, 2]:
xs = pd.read_pickle(os.path.join(dataset_path, 'test_x.pkl'))
else:
xs = pd.read_pickle(os.path.join(dataset_path, 'test_x_no_impute.pkl'))
ys = pd.read_pickle(os.path.join(dataset_path, 'test_y.pkl'))
pids = pd.read_pickle(os.path.join(dataset_path, 'test_pid.pkl'))
features = pd.read_pickle(os.path.join(dataset_path, 'all_features.pkl'))[2:]
record_times = pd.read_pickle(os.path.join(dataset_path, 'test_x_record_times.pkl'))
labels = []
preds = []
if output_logits:
logits_path = os.path.join(logits_root, dataset, task, model)
Path(logits_path).mkdir(parents=True, exist_ok=True)
sub_dst_name = f'{form}_{str(nshot)}shot_{time_des}'
if config['unit'] is True:
sub_dst_name += '_unit'
if config['reference_range'] is True:
sub_dst_name += '_range'
if config.get('prompt_engineering') is True:
sub_dst_name += '_cot'
if impute == 0:
sub_dst_name += '_no_impute'
elif impute == 1:
sub_dst_name += '_impute'
elif impute == 2:
sub_dst_name += '_impute_info'
sub_logits_path = os.path.join(logits_path, sub_dst_name)
Path(sub_logits_path).mkdir(parents=True, exist_ok=True)
if output_prompts:
prompts_path = os.path.join(prompts_root, dataset, task, model)
Path(prompts_path).mkdir(parents=True, exist_ok=True)
sub_dst_name = f'{form}_{str(nshot)}shot_{time_des}'
if config['unit'] is True:
sub_dst_name += '_unit'
if config['reference_range'] is True:
sub_dst_name += '_range'
if config.get('prompt_engineering') is True:
sub_dst_name += '_cot'
if impute == 0:
sub_dst_name += '_no_impute'
elif impute == 1:
sub_dst_name += '_impute'
elif impute == 2:
sub_dst_name += '_impute_info'
sub_prompts_path = os.path.join(prompts_path, sub_dst_name)
Path(sub_prompts_path).mkdir(parents=True, exist_ok=True)
for x, y, pid, record_time in zip(xs, ys, pids, record_times):
if isinstance(pid, float):
pid = str(round(pid))
length = len(x)
sex = 'male' if x[0][0] == 1 else 'female'
age = x[0][1]
x = [visit[2:] for visit in x]
detail = format_input(
patient=x,
dataset=dataset,
form=form,
features=features,
)
input_format_description = INPUT_FORMAT_DESCRIPTION[form]
if impute == 0:
input_format_description += MISSING_VALUE_DESCRIPTION
elif impute == 2:
input_format_description += INSTRUCTING_MISSING_VALUE
userPrompt = USERPROMPT.format(
INPUT_FORMAT_DESCRIPTION=input_format_description,
TASK_DESCRIPTION_AND_RESPONSE_FORMAT=task_description,
UNIT_RANGE_CONTEXT=unit_range,
EXAMPLE=example,
SEX=sex,
AGE=age,
LENGTH=length,
RECORD_TIME_LIST=', '.join(list(map(str, record_time))),
DETAIL=detail,
RESPONSE_FORMAT=response_format,
)
if output_prompts:
with open(os.path.join(sub_prompts_path, f'{pid}.txt'), 'w') as f:
f.write(userPrompt)
if output_logits:
try:
result, prompt_token, completion_token = query_llm(
model=model,
llm=llm,
systemPrompt=SYSTEMPROMPT[dataset],
userPrompt=userPrompt
)
except Exception as e:
# logging.info(f'PatientID: {patient.iloc[0]["PatientID"]}:\n')
logging.info(f'Query LLM Exception: {e}')
continue
prompt_tokens += prompt_token
completion_tokens += completion_token
if task == 'outcome':
label = y[0][0]
elif task == 'readmission':
label = y[0][2]
elif task == 'los':
label = [yi[1] for yi in y]
elif task == 'multitask':
label = [y[0][0], y[0][2]]
else:
raise ValueError(f'Unknown task: {task}')
try:
if config.get('prompt_engineering') is True:
pred = result
elif task in ['los', 'multitask']:
pred = [float(p) for p in result.split(',')]
else:
pred = float(result)
except:
if task == 'los':
pred = [0] * len(label)
elif task == 'multitask':
pred = [0.501, 0.501]
else:
pred = 0.501
logging.info(f'PatientID: {pid}:\nResponse: {result}\n')
pd.to_pickle({
'prompt': userPrompt,
'pred': pred,
'label': label,
}, os.path.join(sub_logits_path, f'{pid}.pkl'))
labels.append(label)
preds.append(pred)
if output_logits:
logging.info(f'Prompts: {prompt_tokens}, Completions: {completion_tokens}, Total: {prompt_tokens + completion_tokens}\n\n')
pd.to_pickle({
'config': config,
'preds': preds,
'labels': labels,
}, os.path.join(logits_path, sub_dst_name + '.pkl'))
if __name__ == '__main__':
for config in params:
run(config, output_logits=True, output_prompts=False)
