Transformer embeddings for clinical NLP

import pandas as pd
import numpy as np
import polars as pl

from os import listdir

import matplotlib.pyplot as plt
import seaborn as sns

from sentence_transformers import SentenceTransformer, models
from datasets import load_dataset
from sentence_transformers import InputExample
from torch.utils.data import DataLoader
from sentence_transformers import losses

sns.set()
%matplotlib inline

source_folder = '/Users/alexlee/Desktop/Data/clinical/clinical_abbreviations/'

filenames = listdir(source_folder)

df1 = (
    pd
    .read_csv(f'{source_folder}/{filenames[0]}', sep='=', header=None, names=['abbreviation', 'sense'])
    .assign(abbreviation=lambda df_: df_.abbreviation.str.strip(), 
            sense=lambda df_: df_.sense.str.strip())
)

df2 = (
    pd
    .read_csv(f'{source_folder}/{filenames[1]}', sep='\t', header=None, names=['abbreviation', 'sense', 'similarity'])
    .assign(abbreviation=lambda df_: df_.abbreviation.str.strip(), 
            sense=lambda df_: df_.sense.str.strip())
)

df3 = (
    pd
    .read_csv(f'{source_folder}/{filenames[2]}', sep=',', names=['abbreviation', 'sense'])
    .assign(abbreviation=lambda df_: df_.abbreviation.str.strip(), 
            sense=lambda df_: df_.sense.str.strip())
)

df4 = (
    pd
    .read_csv(f'{source_folder}/vanderbilt_clinic_notes.txt', sep='\t')
    .assign(abbreviation=lambda df_: df_.abbreviation.str.strip(), 
            sense=lambda df_: df_.sense.str.strip())
) 

df5 = (
    pd
    .read_csv(f'{source_folder}/vanderbilt_discharge_sums.txt', sep='\t')
    .assign(abbreviation=lambda df_: df_.abbreviation.str.strip(), 
            sense=lambda df_: df_.sense.str.strip())
) 

df_all = pd.concat([df1, df2, df3, df4, df4])

df_all = (
    df_all
    .loc[:, ['abbreviation', 'sense']]
    .drop_duplicates()
    .sort_values(by=['abbreviation'])
    .query('abbreviation.isnull() == False')
    .query('sense.isnull() == False')
    .reset_index()
    .iloc[:, 1:]
)

# training data
train_data = df_all.values

## Step 1: use an existing language model
word_embedding_model = models.Transformer('distilroberta-base')

## Step 2: use a pool function over the token embeddings
pooling_model = models.Pooling(word_embedding_model.get_word_embedding_dimension())

## Join steps 1 and 2 using the modules argument
model = SentenceTransformer(modules=[word_embedding_model, pooling_model])

train_examples = []

n_examples = len(train_data)

# convert each of the examples 
for i in range(n_examples):
  example = train_data[i]
  train_examples.append(InputExample(texts=[example[0], example[1]]))

train_dataloader = DataLoader(train_examples, shuffle=True, batch_size=16)

train_loss = losses.MultipleNegativesRankingLoss(model=model)

model.fit(train_objectives=[(train_dataloader, train_loss)], epochs=10) 

model.save('clinical_embeddings_100524')

# initially inspecting some of the matches
texts = ['sob', 
         'shortness of breath', 
         'hbg', 
         'plt', 
         'bilirubin', 
         'haemoglobin', 
         'platelets', 
         'alp', 
         'alkaline phosphatase', 
         'hb', 'hb.', 'plt.', 'plat', 's.o.b', 'sob on arrival']

texts_emb = model.encode(texts)

s1 = []
s2 = []
scores = []

for n in range(len(texts)):
    for m in range(len(texts)):
        s1.append(texts[n])
        s2.append(texts[m])
        score = texts_emb[n].dot(texts_emb[m].T)
        scores.append(score)

df = pd.DataFrame(data={'text1': s1, 
                        'text2': s2, 
                        'similarity': scores})

query = 'sob on arrival'

df.query(f'text1 == "{query}"').sort_values(by='similarity', ascending=False)