# Lower case the first and last names
df["FIRST_NAME"] = df["FIRST_NAME"].apply(lambda x: str(x).lower())
df["LAST_NAME"] = df["LAST_NAME"].apply(lambda x: str(x).lower())
first_names = set(df["FIRST_NAME"].values)
last_names = set(df["LAST_NAME"].values)
sample_with_names = preprocess_parallel(has_name,
all_sentences,
chunksize=1000,
names=first_names | last_names)
print(len(sample_with_names), len(all_sentences))
subject_id_to_patient_info = get_subject_id_to_patient_info("medcat")
name_to_subject_id = {}
for subject_id, patient_info in subject_id_to_patient_info.items():
name_to_subject_id.setdefault(patient_info.FIRST_NAME.lower(),
[]).append(subject_id)
name_to_subject_id.setdefault(patient_info.LAST_NAME.lower(),
[]).append(subject_id)
found_sentences = 0
for sample_sentence, sample_names in tqdm(sample_with_names):
entities = [x for x in get_entities(sample_sentence)]
subject_ids = list(
set([
subject_id for name in sample_names
if name in name_to_subject_id
def main(model: gensim.models.KeyedVectors, condition_type: str,
metrics_output_path: str):
subject_id_to_patient_info = get_subject_id_to_patient_info(
condition_type=condition_type)
condition_code_to_count = get_condition_code_to_count(
condition_type=condition_type)
condition_code_to_description = get_condition_code_to_descriptions(
condition_type=condition_type)
set_to_use = filter_condition_code_by_count(condition_code_to_count,
min_count=50,
max_count=500000)
condition_code_to_index: Dict[str, int] = dict(
zip(set_to_use, range(len(set_to_use))))
mean_condition_embeddings = []
max_condition_embeddings = []
all_condition_embeddings = []
for condition in set_to_use:
desc = condition_code_to_description[condition]
condition_embeddings = get_embedding(model, desc)
mean_embedding = normalize(
np.mean(condition_embeddings, axis=0, keepdims=True))
mean_condition_embeddings.append(mean_embedding)
max_embedding = normalize(
np.max(condition_embeddings, axis=0, keepdims=True))
max_condition_embeddings.append(max_embedding)
all_condition_embeddings.append(normalize(condition_embeddings))
mean_condition_embeddings = np.concatenate(
mean_condition_embeddings,
axis=0) # Shape = (Num_Conditions, Embedding Size)
max_condition_embeddings = np.concatenate(
max_condition_embeddings,
axis=0) # Shape = (Num_Conditions, Embedding Size)
mean_differential_sim, max_differential_sim, all_pair_differential_sim = [], [], []
for subject_id, patient_info in tqdm(subject_id_to_patient_info.items()):
name = patient_info.FIRST_NAME + " " + patient_info.LAST_NAME
name_embeddings = get_embedding(model, name)
mean_name_embedding = normalize(
name_embeddings.mean(0)) # Shape = (Embedding Size,)
max_name_embedding = normalize(
name_embeddings.max(0)) # Shape = (Embedding Size,)
mean_similarities = mean_condition_embeddings @ mean_name_embedding
max_similarities = max_condition_embeddings @ max_name_embedding
name_embeddings = normalize(name_embeddings)
all_pair_similarities = []
for condition_embeddings in all_condition_embeddings:
similarity_matrix = condition_embeddings @ name_embeddings.T
all_pair_similarities.append(np.max(similarity_matrix))
condition_labels = get_condition_labels_as_vector(
patient_info.CONDITIONS, condition_code_to_index)
if sum(condition_labels) == 0: continue
mean_differential_sim.append(
differential_score(condition_labels, mean_similarities))
max_differential_sim.append(
differential_score(condition_labels, max_similarities))
all_pair_differential_sim.append(
differential_score(condition_labels, all_pair_similarities))
print(f"Mean Mean Pos-Neg {np.average(mean_differential_sim)}")
print(f"SD Mean Pos-Neg {np.std(mean_differential_sim)}")
print(f"Mean Max Pos-Neg {np.average(max_differential_sim)}")
print(f"SD Max Pos-Neg {np.std(max_differential_sim)}")
print(f"Mean All Pair Pos-Neg {np.average(all_pair_differential_sim)}")
print(f"SD All Pair Pos-Neg {np.std(all_pair_differential_sim)}")
from experiments.MLM.common import mean_std_as_string
with open(f"{metrics_output_path}/results.txt", "w") as f:
f.write(mean_std_as_string("Mean Sim", mean_differential_sim))
f.write(mean_std_as_string("Max Sim", max_differential_sim))
f.write(mean_std_as_string("All Pair Sim", all_pair_differential_sim))
def train_and_evaluate(
model: BertModel,
tokenizer: BertTokenizer,
condition_type: str,
sampling_bin: int,
n: int,
metrics_output_path: str,
):
"""Train and evaluate the model on N conditions.
@param n is the number of conditions to sample the bin from.
"""
### Get Relevant Data
subject_id_to_patient_info = get_subject_id_to_patient_info(condition_type=condition_type)
condition_code_to_count = get_condition_code_to_count(condition_type=condition_type)
condition_code_to_description = get_condition_code_to_descriptions(condition_type=condition_type)
set_to_use = filter_condition_code_by_count(condition_code_to_count, min_count=0, max_count=500000)
binned_conditions = get_frequency_bins(condition_code_to_count, condition_type)
subject_ids = sorted(list(subject_id_to_patient_info.keys()))
train_subject_ids, test_subject_ids = train_test_split(
subject_ids, train_size=0.5, random_state=2021, shuffle=True
)
### Filter condition in each bin so we have atleast one positive training example
### And One positive test example
### Otherwise, we can't train a LR model or calculate roc_auc_score
train_set_conditions = get_non_zero_count_conditions(
set_to_use, train_subject_ids, subject_id_to_patient_info
)
test_set_conditions = get_non_zero_count_conditions(
set_to_use, test_subject_ids, subject_id_to_patient_info
)
binned_conditions = [set(bin_) & train_set_conditions & test_set_conditions for bin_ in binned_conditions]
binned_conditions = [sorted(list(bin_)) for bin_ in binned_conditions]
### Sample condition in selected bin
condition_bin = binned_conditions[sampling_bin]
np.random.seed(2021)
sampled_conditions = np.random.choice(condition_bin, size=n, replace=False)
## Train a Classifier for Each Condition
auc_score_list, precision_at_10_list = [], []
for condition in tqdm(sampled_conditions):
desc = condition_code_to_description[condition]
## Get all train templates and labels for all train patients, for this condition
train_templates = []
train_labels = []
for subject_id in train_subject_ids:
patient_info = subject_id_to_patient_info[subject_id]
template = generate_name_condition_template(
patient_info.FIRST_NAME, patient_info.LAST_NAME, patient_info.GENDER, desc
)
label = condition in patient_info.CONDITIONS
train_templates.append(template)
train_labels.append(label)
## Resample to Upsample positive examples
negative_indices = [i for i, x in enumerate(train_labels) if x == 0]
positive_indices = [i for i, x in enumerate(train_labels) if x == 1]
# We set replace = to False in another file; does this matter?
positive_indices = resample(
positive_indices, replace=True, n_samples=len(negative_indices), random_state=2021
)
total_indices = negative_indices + positive_indices
train_templates = [train_templates[i] for i in total_indices]
train_labels = [train_labels[i] for i in total_indices]
## Train the LR model
train_embeddings = get_cls_embeddings(model, tokenizer, train_templates, disable_tqdm=True)
clf = MLPClassifier(hidden_layer_sizes=(128,), random_state=2021).fit(train_embeddings, train_labels)
## Get all test templates and labels for all test patients, for this condition
test_templates = []
test_labels = []
for subject_id in test_subject_ids:
patient_info = subject_id_to_patient_info[subject_id]
template = generate_name_condition_template(
patient_info.FIRST_NAME, patient_info.LAST_NAME, patient_info.GENDER, desc
)
label = condition in patient_info.CONDITIONS
test_templates.append(template)
test_labels.append(label)
## Get Embeddings for all test patients, and make prediction with LR model
test_embeddings = get_cls_embeddings(model, tokenizer, test_templates, disable_tqdm=True)
test_predictions = clf.predict_proba(test_embeddings)[:, 1]
auc_score = roc_auc_score(test_labels, test_predictions)
precision_at_10 = precision_at_k(test_labels, test_predictions, k=10)
auc_score_list.append(auc_score)
precision_at_10_list.append(precision_at_10)
from experiments.MLM.common import mean_std_as_string
with open(f"{metrics_output_path}/results.txt", "w") as f:
f.write(mean_std_as_string("Model AUC", auc_score_list))
f.write(mean_std_as_string("Model P@K", precision_at_10_list))
def main(model, tokenizer: BertTokenizerFast, condition_type: str,
metrics_output_path: str):
"""Compute the BERT representations + cosine similarities."""
## Get Relevant data
subject_id_to_patient_info = get_subject_id_to_patient_info(
condition_type=condition_type)
condition_code_to_count = get_condition_code_to_count(
condition_type=condition_type)
condition_code_to_description = get_condition_code_to_descriptions(
condition_type=condition_type)
set_to_use = filter_condition_code_by_count(condition_code_to_count,
min_count=0,
max_count=500000)
condition_code_to_index: Dict[str, int] = dict(
zip(set_to_use, range(len(set_to_use))))
mean_differential_sim, max_differential_sim, all_pair_differential_sim = [], [], []
## For each patient and condition, get a template, pass through BERT and return similarities
all_subject_ids = sorted(list(subject_id_to_patient_info.keys()))
all_subject_ids = sorted(
resample(all_subject_ids,
replace=False,
n_samples=10000,
random_state=2021))
for subject_id in tqdm(all_subject_ids):
patient_info = subject_id_to_patient_info[subject_id]
templates = []
for condition in set_to_use:
desc = condition_code_to_description[condition]
templates.append(
generate_template(patient_info.FIRST_NAME,
patient_info.LAST_NAME, patient_info.GENDER,
desc))
name = patient_info.FIRST_NAME + " " + patient_info.LAST_NAME
name_length = len(tokenizer.tokenize(name))
## Following info may change if we change the template structure.
## Following are on basis of structure [CLS] {title} {name} is a yo patient with {condition} [SEP]
example_template = tokenizer.tokenize(templates[0])
name_start_index = 2 # Name Starts after [CLS] {title}
name_end_index = name_start_index + name_length
condition_start_index = example_template.index("patient") + 2
condition_end_index = -1
assert (tokenizer.convert_tokens_to_string(
example_template[name_start_index:name_end_index]) == " ".join(
name.lower().split())), breakpoint()
assert (tokenizer.convert_tokens_to_string(
example_template[condition_start_index:condition_end_index]) ==
" ".join(condition_code_to_description[
set_to_use[0]].lower().split())), breakpoint()
## Pass all templates to BERT and return similarities
mean_similarities, max_similarities, all_pair_similarities = get_name_condition_similarities(
model,
tokenizer,
templates,
name_start_index,
name_end_index,
condition_start_index,
condition_end_index,
)
condition_labels = get_condition_labels_as_vector(
patient_info.CONDITIONS, condition_code_to_index)
mean_differential_sim.append(
differential_score(condition_labels, mean_similarities))
max_differential_sim.append(
differential_score(condition_labels, max_similarities))
all_pair_differential_sim.append(
differential_score(condition_labels, all_pair_similarities))
print(f"Mean Mean Pos-Neg {np.average(mean_differential_sim)}")
print(f"SD Mean Pos-Neg {np.std(mean_differential_sim)}")
print(f"Mean Max Pos-Neg {np.average(max_differential_sim)}")
print(f"SD Max Pos-Neg {np.std(max_differential_sim)}")
print(f"Mean All Pair Pos-Neg {np.average(all_pair_differential_sim)}")
print(f"SD All Pair Pos-Neg {np.std(all_pair_differential_sim)}")
from experiments.MLM.common import mean_std_as_string
with open(f"{metrics_output_path}/results.txt", "w") as f:
f.write(mean_std_as_string("Mean Sim", mean_differential_sim))
f.write(mean_std_as_string("Max Sim", max_differential_sim))
f.write(mean_std_as_string("All Pair Sim", all_pair_differential_sim))
def run_probe(
model: BertModel,
tokenizer: BertTokenizerFast,
condition_type: str,
template_mode: str,
prober: str,
metrics_output_path: str,
):
"""Train and evaluate the model trained on the data.
Args:
condition_type: icd9 or Stanza
template_mode: Choices in [name_and_condition, condition_only].
Specify if name should be included in template
prober: LR or MLP
"""
### Get Relevant Data
subject_id_to_patient_info = get_subject_id_to_patient_info(
condition_type=condition_type)
condition_code_to_count = get_condition_code_to_count(
condition_type=condition_type)
condition_code_to_description = get_condition_code_to_descriptions(
condition_type=condition_type)
set_to_use = filter_condition_code_by_count(condition_code_to_count,
min_count=0,
max_count=500000)
condition_code_to_index: Dict[str, int] = dict(
zip(set_to_use, range(len(set_to_use))))
## Divide patients into train and test group
all_subject_ids = sorted(list(subject_id_to_patient_info.keys()))
## Sample 10K subjects because 27K takes timeeeeeee
n_samples = min(len(all_subject_ids), 10000)
all_subject_ids = sorted(
resample(all_subject_ids,
replace=False,
n_samples=n_samples,
random_state=2021))
train_subject_ids, test_subject_ids = train_test_split(all_subject_ids,
random_state=2021,
test_size=0.5,
shuffle=True)
print(f"Train Subject Ids : {len(train_subject_ids)}")
print(f"Test Subject Ids : {len(test_subject_ids)}")
import pdb
pdb.set_trace()
## Get training example by generating template for all train patients and all conditions
subject_condition_templates = []
subject_condition_labels = []
for subject_id in train_subject_ids:
patient_info = subject_id_to_patient_info[subject_id]
for condition in set_to_use:
desc = condition_code_to_description[condition]
if template_mode == "name_and_condition":
template = generate_name_condition_template(
patient_info.FIRST_NAME, patient_info.LAST_NAME,
patient_info.GENDER, desc)
elif template_mode == "condition_only":
template = generate_condition_only_template(desc)
else:
raise NotImplementedError(f"{template_mode} is not available")
subject_condition_templates.append(template)
condition_labels = get_condition_labels_as_vector(
patient_info.CONDITIONS, condition_code_to_index)
subject_condition_labels += list(condition_labels)
## Downsample negative labels since most patients only have few positive conditions
negative_indices = [
i for i, x in enumerate(subject_condition_labels) if x == 0
]
positive_indices = [
i for i, x in enumerate(subject_condition_labels) if x == 1
]
negative_indices = resample(negative_indices,
replace=False,
n_samples=len(positive_indices),
random_state=2021)
total_indices = negative_indices + positive_indices
train_templates = [subject_condition_templates[i] for i in total_indices]
train_labels = [subject_condition_labels[i] for i in total_indices]
print(len(train_templates))
## Get [CLS] token embedding for each template and train a LR classifier
train_cls_embeddings = get_cls_embeddings(model, tokenizer,
train_templates)
print(f"Training {prober} Model")
if prober == "LR":
classifier = LogisticRegression(random_state=2021, max_iter=10000).fit(
train_cls_embeddings, train_labels)
elif prober == "MLP":
classifier = MLPClassifier(hidden_layer_sizes=(128, ),
random_state=2021).fit(
train_cls_embeddings, train_labels)
else:
raise NotImplementedError(f"{prober} not implemented")
print(f"{prober} Model Trained")
## Get templates and labels for test set patients
auc_scores, paks = [], []
for subject_id in tqdm(test_subject_ids):
test_templates = []
patient_info = subject_id_to_patient_info[subject_id]
for condition in set_to_use:
desc = condition_code_to_description[condition]
if template_mode == "name_and_condition":
template = generate_name_condition_template(
patient_info.FIRST_NAME, patient_info.LAST_NAME,
patient_info.GENDER, desc)
elif template_mode == "condition_only":
template = generate_condition_only_template(desc)
else:
raise NotImplementedError(f"{template_mode} is not available")
test_templates.append(template)
condition_labels = get_condition_labels_as_vector(
patient_info.CONDITIONS, condition_code_to_index)
test_labels = list(condition_labels)
test_cls_embeddings = get_cls_embeddings(model,
tokenizer,
test_templates,
disable_tqdm=True)
test_predictions = classifier.predict_proba(test_cls_embeddings)[:, 1]
try:
model_auc = roc_auc_score(test_labels, test_predictions)
model_precision_at_k = precision_at_k(test_labels,
test_predictions,
k=10)
auc_scores.append(model_auc)
paks.append(model_precision_at_k)
except:
continue
from experiments.MLM.common import mean_std_as_string
with open(f"{metrics_output_path}/results.txt", "w") as f:
f.write(mean_std_as_string("Model AUC", auc_scores))
f.write(mean_std_as_string("Model P@K", paks))
def evaluate(
model: BertForMaskedLM,
tokenizer: BertTokenizer,
condition_type: str,
template_idx: int,
max_count: int,
metrics_output_path: str,
):
"""
Evaluate the performance of the model in terms of being able to predict
conditions associated with certain names (via masked language modelling task on templates).
For each patient,
For each condition,
create template with name filled and condition masked.\\
get MLM probability distribution for [MASK] tokens.\\
get average probability of condition's wordpieces.\\
Compute ROC / P@10 against true condition labels.
Report Average over patients.
Ideally we want higher average probability for conditions that patient have, than for conditions they
don't have.
We also include a condition only baseline (where template doesn't contain the patient name). The
algorithm above remains same.
### Args:
condition_type: Which conditions to load for patients. Currently take value in [icd9, medcat]
template: Which template to use for probing the model.
"""
### Load relevant data
subject_id_to_patient_info = get_subject_id_to_patient_info(
condition_type=condition_type)
condition_code_to_count = get_condition_code_to_count(
condition_type=condition_type)
condition_code_to_description = get_condition_code_to_descriptions(
condition_type=condition_type)
set_to_use = filter_condition_code_by_count(condition_code_to_count,
min_count=0,
max_count=max_count)
print(len(set_to_use))
condition_code_to_index: Dict[str, int] = dict(
zip(set_to_use, range(len(set_to_use))))
### Get list of unique condition lengths (in wordpieces) to generate templates
condition_code_to_wordpiece_ids: Dict[str, List[str]] = {
condition: tokenizer.convert_tokens_to_ids(
tokenizer.tokenize(condition_code_to_description[condition]))
for condition in set_to_use
} ## Time Saving Measure since we only need condition wordpiece ids moving forward
condition_wordpiece_lengths = [
len(condition_code_to_wordpiece_ids[condition])
for condition in set_to_use
]
from collections import Counter
print(sorted(list(Counter(condition_wordpiece_lengths).items())))
set_to_use_lengths = np.array(condition_wordpiece_lengths)
condition_wordpiece_lengths: List[int] = sorted(
list(set(condition_wordpiece_lengths))) ## Keep Unique Lengths only
### Get Condition Frequency counts
condition_baseline_counts = np.array(
get_condition_counts_as_vector(condition_code_to_count,
condition_code_to_index))
condition_baseline_freq = condition_baseline_counts / np.sum(
condition_baseline_counts)
### Get Condition only template logits
## Generate Template for each unique condition wordpiece length
condition_only_templates = [
condition_only_template(length)
for length in condition_wordpiece_lengths
]
logits = get_logits_from_templates(model,
tokenizer,
condition_only_templates,
normalize=True)
logits = {
length: logit
for length, logit in zip(condition_wordpiece_lengths, logits)
}
# Isn't the start index always 1 here? Yes. This is to keep code consistent.
start_indices = [
tokenizer.tokenize(template).index("[MASK]")
for template in condition_only_templates
]
start_indices = {
length: start_index
for length, start_index in zip(condition_wordpiece_lengths,
start_indices)
}
condition_only_logits: List[float] = []
for condition in set_to_use:
condition_wp_ids = condition_code_to_wordpiece_ids[condition]
condition_length = len(condition_wp_ids)
condition_only_logits.append(
get_average_predicted_score(logits[condition_length],
condition_wp_ids,
start_indices[condition_length]))
condition_only_logits = normalize_logits(condition_only_logits,
condition_baseline_freq,
condition_wordpiece_lengths,
set_to_use_lengths)
### Get Subject Specific Condition logits
k = 10
rocs, precisions_at_k, spearmans = {}, {}, {}
for length in condition_wordpiece_lengths + ["all"]:
rocs[length] = {"baseline": [], "condition_only": [], "model": []}
precisions_at_k[length] = {
"baseline": [],
"condition_only": [],
"model": []
}
spearmans[length] = {"baseline": [], "condition_only": [], "model": []}
if length == "all":
rocs[length]["bin_prob"] = []
precisions_at_k[length]["bin_prob"] = []
spearmans[length]["bin_prob"] = []
patient_to_run = [
sum(
get_condition_labels_as_vector(patient_info.CONDITIONS,
condition_code_to_index)) > 0
for subject_id, patient_info in subject_id_to_patient_info.items()
]
print(sum(patient_to_run))
condition_bin_prob_baselines = np.zeros_like(condition_baseline_freq)
for length in condition_wordpiece_lengths:
mask = set_to_use_lengths == length
condition_bin_prob_baselines[mask] = condition_baseline_freq[
mask].mean()
for subject_id, patient_info in tqdm(subject_id_to_patient_info.items()):
condition_labels = get_condition_labels_as_vector(
patient_info.CONDITIONS, condition_code_to_index)
condition_labels = np.array(condition_labels)
if condition_labels.sum() == 0:
continue ## Skip if patient is negative for all conditions
## Generate Template for each unique condition wordpiece length
templates = []
for length in condition_wordpiece_lengths:
template = name_with_condition_template(patient_info.FIRST_NAME,
patient_info.LAST_NAME,
patient_info.GENDER,
length, template_idx)
templates.append(template)
## Get logits for all templates
logits = get_logits_from_templates(model,
tokenizer,
templates,
normalize=True)
logits = {
length: logit
for length, logit in zip(condition_wordpiece_lengths, logits)
}
## Get Start index for (masked) condition in each template
## Not sure we need to do this for all templates? -> Once we have the index of the mask of one, it should be the same for all?
start_indices = [
tokenizer.tokenize(template).index("[MASK]")
for template in templates
]
start_indices = {
length: start_index
for length, start_index in zip(condition_wordpiece_lengths,
start_indices)
}
## For each condition, get corresponding logit array and then compute average score
condition_subject_logits = []
for condition in set_to_use:
condition_wp_ids = condition_code_to_wordpiece_ids[condition]
condition_length = len(condition_wp_ids)
condition_subject_logits.append(
get_average_predicted_score(logits[condition_length],
condition_wp_ids,
start_indices[condition_length]))
condition_subject_logits = normalize_logits(
condition_subject_logits, condition_baseline_freq,
condition_wordpiece_lengths, set_to_use_lengths)
for length in condition_wordpiece_lengths + ["all"]:
mask = (set_to_use_lengths == length if length != "all" else
np.full_like(set_to_use_lengths, True, dtype=bool))
length_condition_baseline_counts = condition_baseline_counts[mask]
length_condition_only_logits = condition_only_logits[mask]
length_condition_subject_logits = condition_subject_logits[mask]
length_condition_labels = condition_labels[mask]
if length_condition_labels.sum() == 0 or mask.sum() < 2:
## If patient doesn't have any positive condition or only one condition in bin
continue
### Calculate and store metrics for this patient
_baseline_roc = roc_auc_score(length_condition_labels,
length_condition_baseline_counts)
_condition_only_roc = roc_auc_score(length_condition_labels,
length_condition_only_logits)
_model_roc = roc_auc_score(length_condition_labels,
length_condition_subject_logits)
rocs[length]["baseline"].append(_baseline_roc)
rocs[length]["condition_only"].append(_condition_only_roc)
rocs[length]["model"].append(_model_roc)
_baseline_spearman = spearmanr(
length_condition_baseline_counts,
length_condition_baseline_counts).correlation
_condition_only_spearman = spearmanr(
length_condition_baseline_counts, length_condition_only_logits)
_model_spearman = spearmanr(length_condition_baseline_counts,
length_condition_subject_logits)
spearmans[length]["baseline"].append(_baseline_spearman)
spearmans[length]["condition_only"].append(
_condition_only_spearman)
spearmans[length]["model"].append(_model_spearman)
_model_precision_at_k = precision_at_k(
length_condition_labels, length_condition_subject_logits, k)
_condition_only_precision_at_k = precision_at_k(
length_condition_labels, length_condition_only_logits, k)
_baseline_precision_at_k = precision_at_k(
length_condition_labels, length_condition_baseline_counts, k)
precisions_at_k[length]["baseline"].append(
_baseline_precision_at_k)
precisions_at_k[length]["condition_only"].append(
_condition_only_precision_at_k)
precisions_at_k[length]["model"].append(_model_precision_at_k)
if length == "all":
rocs[length]["bin_prob"].append(
roc_auc_score(length_condition_labels,
condition_bin_prob_baselines))
spearmans[length]["bin_prob"].append(
spearmanr(length_condition_baseline_counts,
condition_bin_prob_baselines).correlation)
precisions_at_k[length]["bin_prob"].append(
precision_at_k(length_condition_labels,
condition_bin_prob_baselines, k))
### Computing and print metrics (averaged over patients)
with open(f"{metrics_output_path}/results.txt", "w") as f:
for length in ["all"] + condition_wordpiece_lengths:
bin_length = (set_to_use_lengths == length if length != "all" else
np.full_like(set_to_use_lengths, True, dtype=bool))
f.write(
f"Length {length} # Num conditions in Length bin {bin_length.sum()}\n"
)
if len(rocs[length]["model"]) == 0:
continue
for method, values in rocs[length].items():
f.write(mean_std_as_string(f"{method} AUC", values))
for method, values in precisions_at_k[length].items():
f.write(mean_std_as_string(f"{method} P@K", values))
for method, values in spearmans[length].items():
f.write(mean_std_as_string(f"{method} Spearman", values))
f.write("\n")
def train_and_evaluate(
model: BertForSequenceClassification,
tokenizer: BertTokenizer,
condition_type: str,
sampling_bin: int,
n: int,
metrics_output_path: str,
):
"""Train and evaluate the model on N conditions.
@param model is the model to encode CLS tokens with.
@param tokenizer is a BERT tokenizer.
@param condition_type are we using the icd/medcat extracted conditions?
@param b is which frequency to sample from.
@param n is the number of conditions to sample the bin from.
@return all AUCs and precision @ K scores.
"""
### Get Relevant Data
subject_id_to_patient_info = get_subject_id_to_patient_info(
condition_type=condition_type)
condition_code_to_count = get_condition_code_to_count(
condition_type=condition_type)
condition_code_to_description = get_condition_code_to_descriptions(
condition_type=condition_type)
set_to_use = filter_condition_code_by_count(condition_code_to_count,
min_count=0,
max_count=500000)
binned_conditions = get_frequency_bins(condition_code_to_count,
condition_type)
subject_ids = sorted(list(subject_id_to_patient_info.keys()))
train_subject_ids, test_subject_ids = train_test_split(subject_ids,
train_size=0.5,
random_state=2021,
shuffle=True)
### Filter condition in each bin so we have atleast one positive training examples
### And One positive test example
### Otherwise, we can't train a LR model or calculate roc_auc_score
train_set_conditions = get_non_zero_count_conditions(
set_to_use, train_subject_ids, subject_id_to_patient_info)
test_set_conditions = get_non_zero_count_conditions(
set_to_use, test_subject_ids, subject_id_to_patient_info)
binned_conditions = [
set(bin_) & train_set_conditions & test_set_conditions
for bin_ in binned_conditions
]
binned_conditions = [sorted(list(bin_)) for bin_ in binned_conditions]
### Sample condition in selected bin
condition_bin = binned_conditions[sampling_bin]
np.random.seed(2021)
sampled_conditions = np.random.choice(condition_bin, size=n, replace=False)
## Train a Classifier for Each Condition
auc_score_list, precision_at_10_list = [], []
for condition in tqdm(sampled_conditions):
desc = condition_code_to_description[condition]
train_templates = []
train_labels = []
for subject_id in train_subject_ids:
patient_info = subject_id_to_patient_info[subject_id]
template = generate_name_condition_template(
patient_info.FIRST_NAME, patient_info.LAST_NAME,
patient_info.GENDER, desc)
label = condition in patient_info.CONDITIONS
train_templates.append(template)
train_labels.append(label)
## Resample to Upsample positive examples
negative_indices = [i for i, x in enumerate(train_labels) if x == 0]
positive_indices = [i for i, x in enumerate(train_labels) if x == 1]
positive_indices = resample(positive_indices,
replace=True,
n_samples=len(negative_indices),
random_state=2021)
total_indices = negative_indices + positive_indices
### Divide Train Set into Train and Validation Set
training_indices, validation_indices = train_test_split(
total_indices, train_size=0.85, random_state=2021, shuffle=True)
# Not too sure we can ensure the validation templates have a positive label in it...
# Or if there is only 1, that it doesn't end up in the validation set.
validation_templates = [train_templates[i] for i in validation_indices]
validation_labels = [train_labels[i] for i in validation_indices]
np.random.seed(2021)
np.random.shuffle(training_indices)
train_templates = [train_templates[i] for i in training_indices]
train_labels = [train_labels[i] for i in training_indices]
### Train the BERT Model
train_dataset = get_as_dataset(tokenizer, train_templates,
train_labels)
validation_dataset = get_as_dataset(tokenizer, validation_templates,
validation_labels)
clf = train_model(model, train_dataset, validation_dataset)
### Get Test Templates
test_templates = []
test_labels = []
for subject_id in test_subject_ids:
patient_info = subject_id_to_patient_info[subject_id]
template = generate_name_condition_template(
patient_info.FIRST_NAME, patient_info.LAST_NAME,
patient_info.GENDER, desc)
label = condition in patient_info.CONDITIONS
test_templates.append(template)
test_labels.append(label)
### Get Test Predictions
test_dataset = get_as_dataset(tokenizer, test_templates, test_labels)
test_predictions = clf.predict(test_dataset)
test_predictions = test_predictions.predictions[:, 1]
### Calculate Metrics
auc_score = roc_auc_score(test_labels, test_predictions)
precision_at_10 = precision_at_k(test_labels, test_predictions, k=10)
auc_score_list.append(auc_score)
precision_at_10_list.append(precision_at_10)
from experiments.MLM.common import mean_std_as_string
with open(f"{metrics_output_path}/results.txt", "w") as f:
f.write(mean_std_as_string("Model AUC", auc_score_list))
f.write(mean_std_as_string("Model P@K", precision_at_10_list))
from argparse import ArgumentParser
parser = ArgumentParser()
parser.add_argument("--condition-type", required=True)
parser.add_argument("--output-file", required=True)
args = parser.parse_args()
def generate_name_condition_template(first_name: str, last_name: str,
gender: str, condition_description: str):
title = "Mr" if gender == "M" else "Mrs" # I guess just assume married w/e idk ?
return f"{title} {first_name} {last_name} is a yo patient with {condition_description}"
data = get_subject_id_to_patient_info(args.condition_type)
desc = get_condition_code_to_descriptions(args.condition_type)
templates = [
"\n".join([
generate_name_condition_template(d.FIRST_NAME, d.LAST_NAME, d.GENDER,
desc[c]) for c in d.CONDITIONS
]) for d in data.values()
]
subject_ids = [d for d in data.keys()]
import pandas as pd
df = pd.DataFrame({"SUBJECT_ID": subject_ids, "TEXT": templates})
df.to_csv(args.output_file, index=False)