ClinicalBERT

To perform K-fold cross-validation, the code uses sklearn.model_selection.KFold on the non-readmitted and readmitted ID pools.

Note on indexing: The code treats the first fold as a special case (generated manually) and uses i=0 to represent Fold 2, i=1 for Fold 3, etc., when iterating through the KFold results.

For each fold i:

1. train_id_folds[i] contains the training IDs.
2. val_id_folds[i] contains the validation IDs.
3. test_id_folds[i] contains the test IDs.

These folds can be used to generate complete datasets for Discharge Summaries or Early Notes tasks by saving them to ./good_datasets/fold{i+2}/....

from sklearn.model_selection import KFold

# Setup KFold
kf = KFold(n_splits = 4, shuffle = True, random_state = 1)

# Loop through folds to populate val_id_folds and test_id_folds
# (Logic involves sampling and dropping duplicates to ensure clean splits)

To generate pretraining data for Clinical BERT, use the create_pretraining_data.py script. It is recommended to perform pretraining in two stages: first on short sequences (128), then continuing on longer sequences (512) using the first model as the checkpoint.

Stage 1: 128 Max Sequence Length

python create_pretraining_data.py \
  --input_file=PRETRAIN_DATA_PATH/clinical_sentences_pretrain.txt \
  --output_file=PRETRAIN_DATA_PATH/tf_examples_128.tfrecord \
  --vocab_file=INITIAL_MODEL_PATH/vocab.txt \
  --do_lower_case=True \
  --max_seq_length=128 \
  --max_predictions_per_seq=20 \
  --masked_lm_prob=0.15 \
  --random_seed=12345 \
  --dupe_factor=3

Stage 2: 512 Max Sequence Length

python create_pretraining_data.py \
  --input_file=PRETRAIN_DATA_PATH/clinical_sentences_pretrain.txt \
  --output_file=PRETRAIN_DATA_PATH/tf_examples_512.tfrecord \
  --vocab_file=INITIAL_MODEL_PATH/vocab.txt \
  --do_lower_case=True \
  --max_seq_length=512 \
  --max_predictions_per_seq=76 \
  --masked_lm_prob=0.15 \
  --random_seed=12345 \
  --dupe_factor=3

Clinical text requires specific cleaning to remove de-identification markers and standardize medical abbreviations.

Cleaning Steps:

• Remove de-identified brackets [ ].
• Remove segmenter markers like 1.2..
• Standardize abbreviations: dr. $\rightarrow$ doctor, m.d. $\rightarrow$ md.
• Remove metadata like admission date: and discharge date:.
• Remove special characters like --, __, or ==.
• Remove all digits and normalize whitespace.
• Convert all text to lowercase.

Sentence Segmentation: Use spacy with a sentencizer to split notes into sentences. For very short segments (length < 20), it is recommended to append them to the previous sentence to handle abbreviations that are segmented as single lines.

import re
import string
import spacy
from spacy.lang.en import English

def preprocess1(x):
    y=re.sub('\\[(.*?)\\]','',x) # remove de-identified brackets
    y=re.sub('[0-9]+\.','',y) # remove 1.2. markers
    y=re.sub('dr\.', 'doctor', y)
    y=re.sub('m\.d\.', 'md', y)
    y=re.sub('admission date:', '', y)
    y=re.sub('discharge date:', '', y)
    y=re.sub('--|__|==', '', y)
    y = y.translate(str.maketrans('', '', string.digits))
    y = " ".join(y.split())
    return y

def preprocessing(df_notes):
    df_notes['TEXT']=df_notes['TEXT'].fillna(' ')
    df_notes['TEXT']=df_notes['TEXT'].str.replace('\n',' ')
    df_notes['TEXT']=df_notes['TEXT'].str.replace('\r',' ')
    df_notes['TEXT']=df_notes['TEXT'].apply(str.strip)
    df_notes['TEXT']=df_notes['TEXT'].str.lower()
    df_notes['TEXT']=df_notes['TEXT'].apply(lambda x: preprocess1(x))
    return df_notes

# Sentence segmentation logic
nlp = English()
nlp.add_pipe(nlp.create_pipe('sentencizer'))

def toSentence(x):
    doc = nlp(x)
    text=[]
    try:
        for sent in doc.sents:
            st=str(sent).strip()
            if len(st)<20:
                if len(text)!=0:
                    text[-1]=' '.join((text[-1],st))
                else:
                    text=[st]
            else:
                text.append((st))
    except:
        print(doc)
    return text

Pretraining for Clinical BERT is performed using run_pretraining.py. It is recommended to use a two-stage approach:

1. Stage 1: Pretrain for 100,000 steps using a max_seq_length of 128.
2. Stage 2: Pretrain for another 100,000 steps using a max_seq_length of 512, initializing from the checkpoint produced in Stage 1.

Stage 1 Command:

python run_pretraining.py \
  --input_file=PRETRAIN_DATA_PATH/tf_examples_128.tfrecord \
  --output_dir=PRETRAINED_MODEL_PATH/pretraining_output \
  --do_train=True \
  --do_eval=True \
  --bert_config_file=INITIAL_DATA_PATH/bert_config.json \
  --init_checkpoint=INITIAL_DATA_PATH/bert_model.ckpt \
  --train_batch_size=64 \
  --max_seq_length=128 \
  --max_predictions_per_seq=20 \
  --num_train_steps=100000 \
  --num_warmup_steps=10 \
  --learning_rate=2e-5

Stage 2 Command:

python run_pretraining.py \
  --input_file=PRETRAIN_DATA_PATH/tf_examples_512.tfrecord \
  --output_dir=PRETRAINED_MODEL_PATH/pretraining_output \
  --do_train=True \
  --do_eval=True \
  --bert_config_file=INITIAL_DATA_PATH/bert_config.json \
  --init_checkpoint=PRETRAINED_MODEL_PATH/pretraining_output_128/model.ckpt-100000 \
  --train_batch_size=16 \
  --max_seq_length=512 \
  --max_predictions_per_seq=76 \
  --num_train_steps=100000 \
  --learning_rate=2e-5

You can visualize the attention mechanism of a BertForSequenceClassification model by extracting the query and key tensors using PyTorch forward hooks. This allows you to see how much weight each token (query) places on other tokens (keys) in a given text sequence.

To implement this, you need to:

1. Define a transpose_for_scores helper to reshape the attention heads.
2. Register register_forward_hook on the query and key modules of a specific encoder layer.
3. Compute the attention probabilities using the dot-product of the query and key tensors, scaled by the square root of the head dimension.
4. Use matplotlib and seaborn to plot the resulting matrix as a heatmap.

import torch
import math
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
from pytorch_pretrained_bert import BertTokenizer

# Assuming model, bert_config, and tokenizer are already initialized

def transpose_for_scores(config, x):
    new_x_shape = x.size()[:-1] + (config.num_attention_heads, int(config.hidden_size / config.num_attention_heads))
    x = x.view(*new_x_shape)
    return x.permute(0, 2, 1, 3)

def get_attention_scores(model, i, text):
    tokenized = tokenizer.tokenize(text)
    indexed_tokens = tokenizer.convert_tokens_to_ids(tokenized)
    segment_ids = [0] * len(indexed_tokens)
    t_tensor = torch.tensor([indexed_tokens])
    s_ids = torch.tensor([segment_ids])

    outputs_query = []
    outputs_key = []

    def hook_query(module, input, output):
        outputs_query.append(output)

    def hook_key(module, input, output):
        outputs_key.append(output)

    # Register hooks on the specific layer
    model.bert.encoder.layer[i].attention.self.query.register_forward_hook(hook_query)
    model.bert.encoder.layer[i].attention.self.key.register_forward_hook(hook_key)
    
    model(t_tensor, s_ids)

    query_layer = transpose_for_scores(bert_config, outputs_query[0])
    key_layer = transpose_for_scores(bert_config, outputs_key[0])

    attention_scores = torch.matmul(query_layer, key_layer.transpose(-1, -2))
    attention_scores = attention_scores / math.sqrt(int(bert_config.hidden_size / bert_config.num_attention_heads))
    attention_probs = torch.nn.Softmax(dim=-1)(attention_scores)

    return attention_probs, tokenized

# Usage
text = 'he has experienced acute on chronic diastolic heart failure'
attention_probs, tokens = get_attention_scores(model, 0, text)

# Extracting a specific head's map (e.g., head 1)
map1 = np.asarray(attention_probs[0][1].detach().numpy())

# Plotting
f, ax = plt.subplots(figsize=(10,10))
ax.imshow(map1, interpolation='nearest', cmap='gray')
ax.set_yticks(range(len(tokens)))
ax.set_yticklabels(tokens)
ax.set_xticks(range(len(tokens)))
ax.set_xticklabels(tokens, rotation=60)
plt.show()

To prepare the dataset for readmission prediction, you must process the ADMISSIONS.csv file to calculate the time between discharge and the next admission.

Key steps include:

1. Convert ADMITTIME, DISCHTIME, and DEATHTIME to datetime objects.
2. Sort by SUBJECT_ID and ADMITTIME.
3. Calculate NEXT_ADMITTIME and NEXT_ADMISSION_TYPE using a groupby-shift operation.
4. Handle 'ELECTIVE' admissions by setting their next admission time to NaT.
5. Define the OUTPUT_LABEL as 1 if the patient is readmitted within 30 days (DAYS_NEXT_ADMIT < 30), and 0 otherwise.
6. Filter out 'NEWBORN' admission types and patients who died during the current stay (DEATHTIME is not null).

# Convert times
df_adm.ADMITTIME = pd.to_datetime(df_adm.ADMITTIME, format = '%Y-%m-%d %H:%M:%S', errors = 'coerce')
df_adm.DISCHTIME = pd.to_datetime(df_adm.DISCHTIME, format = '%Y-%m-%d %H:%M:%S', errors = 'coerce')
df_adm.DEATHTIME = pd.to_datetime(df_adm.DEATHTIME, format = '%Y-%m-%d %H:%M:%S', errors = 'coerce')

# Calculate next admission info
df_adm = df_adm.sort_values(['SUBJECT_ID','ADMITTIME'])
df_adm['NEXT_ADMITTIME'] = df_adm.groupby('SUBJECT_ID').ADMITTIME.shift(-1)
df_adm['NEXT_ADMISSION_TYPE'] = df_adm.groupby('SUBJECT_ID').ADMISSION_TYPE.shift(-1)

# Calculate days until next admit
df_adm['DAYS_NEXT_ADMIT']=  (df_adm.NEXT_ADMITTIME - df_adm.DISCHTIME).dt.total_seconds()/(24*60*60)

# Define label (1 if readmitted < 30 days)
df_adm1['OUTPUT_LABEL'] = (df_adm1.DAYS_NEXT_ADMIT < 30).astype('int')

# Filter
df_adm1 = df_adm1[df_adm1['ADMISSION_TYPE']!='NEWBORN']
df_adm1 = df_adm1[df_adm1.DEATHTIME.isnull()]