app.py

import os
import torch
from datasets import load_dataset, DatasetDict
from transformers import AutoTokenizer, AutoModel
import chromadb
import gradio as gr
import numpy as np
from sklearn.metrics import precision_score, recall_score, f1_score

# Mean Pooling - Take attention mask into account for correct averaging
def meanpooling(output, mask):
    embeddings = output[0]  # First element of model_output contains all token embeddings
    mask = mask.unsqueeze(-1).expand(embeddings.size()).float()
    return torch.sum(embeddings * mask, 1) / torch.clamp(mask.sum(1), min=1e-9)

# Load the dataset
dataset = load_dataset("thankrandomness/mimic-iii")

# Split the dataset into train and validation sets
split_dataset = dataset['train'].train_test_split(test_size=0.2, seed=42)
dataset = DatasetDict({
    'train': split_dataset['train'],
    'validation': split_dataset['test']
})

# Load the model and tokenizer
tokenizer = AutoTokenizer.from_pretrained("neuml/pubmedbert-base-embeddings-matryoshka")
model = AutoModel.from_pretrained("neuml/pubmedbert-base-embeddings-matryoshka")

# Function to normalize embeddings to unit vectors
def normalize_embedding(embedding):
    norm = np.linalg.norm(embedding)
    return (embedding / norm).tolist() if norm > 0 else embedding

# Function to embed and normalize text
def embed_text(text):
    inputs = tokenizer(text, padding=True, truncation=True, max_length=512, return_tensors='pt')
    with torch.no_grad():
        output = model(**inputs)
    embeddings = meanpooling(output, inputs['attention_mask'])
    normalized_embeddings = normalize_embedding(embeddings.numpy())
    return normalized_embeddings

# Initialize ChromaDB client
client = chromadb.Client()
collection = client.create_collection(name="pubmedbert_matryoshka_embeddings")

# Function to upsert data into ChromaDB
def upsert_data(dataset_split):
    for i, row in enumerate(dataset_split):
        for note in row['notes']:
            text = note.get('text', '')
            annotations_list = []
            
            for annotation in note.get('annotations', []):
                try:
                    code = annotation['code']
                    code_system = annotation['code_system']
                    description = annotation['description']
                    annotations_list.append({"code": code, "code_system": code_system, "description": description})
                except KeyError as e:
                    print(f"Skipping annotation due to missing key: {e}")

            if text and annotations_list:
                embeddings = embed_text([text])[0]

                # Upsert data, embeddings, and annotations into ChromaDB
                for j, annotation in enumerate(annotations_list):
                    collection.upsert(
                        ids=[f"note_{note['note_id']}_{j}"],
                        embeddings=[embeddings],
                        metadatas=[annotation]
                    )
            else:
                print(f"Skipping note {note['note_id']} due to missing 'text' or 'annotations'")

# Upsert training data
upsert_data(dataset['train'])

# Define retrieval function with similarity threshold
def retrieve_relevant_text(input_text):
    input_embedding = embed_text([input_text])[0]
    results = collection.query(
        query_embeddings=[input_embedding],
        n_results=5,
        include=["metadatas", "documents", "distances"]
    )
    
    output = []
    #print("Retrieved items and their similarity scores:")
    for metadata, distance in zip(results['metadatas'][0], results['distances'][0]):
        #print(f"Code: {metadata['code']}, Similarity Score: {distance}")
        #if distance <= similarity_threshold:
        output.append({
            "similarity_score": distance,
            "code": metadata['code'],
            "code_system": metadata['code_system'],
            "description": metadata['description']
        })
    
    # if not output:
    #     print("No results met the similarity threshold.")
    return output

# Evaluate retrieval efficiency on the validation/test set
def evaluate_efficiency(dataset_split):
    y_true = []
    y_pred = []
    total_similarity = 0
    total_items = 0

    for i, row in enumerate(dataset_split):
        for note in row['notes']:
            text = note.get('text', '')
            annotations_list = [annotation['code'] for annotation in note.get('annotations', []) if 'code' in annotation]
            
            if text and annotations_list:
                retrieved_results = retrieve_relevant_text(text)
                retrieved_codes = [result['code'] for result in retrieved_results]
                
                # Sum up similarity scores for average calculation
                for result in retrieved_results:
                    total_similarity += result['similarity_score']
                    total_items += 1

                # Ground truth
                y_true.extend(annotations_list)
                # Predictions (limit to length of true annotations to avoid mismatch)
                y_pred.extend(retrieved_codes[:len(annotations_list)])
                
                # for result in retrieved_results:
                #     print(f"  Code: {result['code']}, Similarity Score: {result['similarity_score']:.2f}")

    # Debugging output to check for mismatches and understand results
    # print("Sample y_true:", y_true[:10])
    # print("Sample y_pred:", y_pred[:10])

    if total_items > 0:
        avg_similarity = total_similarity / total_items
    else:
        avg_similarity = 0

    if len(y_true) != len(y_pred):
        min_length = min(len(y_true), len(y_pred))
        y_true = y_true[:min_length]
        y_pred = y_pred[:min_length]
    
    # Calculate metrics
    precision = precision_score(y_true, y_pred, average='macro', zero_division=0)
    recall = recall_score(y_true, y_pred, average='macro', zero_division=0)
    f1 = f1_score(y_true, y_pred, average='macro', zero_division=0)
    
    return precision, recall, f1, avg_similarity

# Calculate retrieval efficiency metrics
precision, recall, f1, avg_similarity = evaluate_efficiency(dataset['validation'])

# Gradio interface
def gradio_interface(input_text):
    results = retrieve_relevant_text(input_text)
    formatted_results = [
        f"Result {i + 1}:\n"
        f"Similarity Score: {result['similarity_score']:.2f}\n"
        f"Code: {result['code']}\n"
        f"Code System: {result['code_system']}\n"
        f"Description: {result['description']}\n"
        "-------------------"
        for i, result in enumerate(results)
    ]
    return "\n".join(formatted_results)

# Display retrieval efficiency metrics
# metrics = f"Precision: {precision:.2f}, Recall: {recall:.2f}, F1 Score: {f1:.2f}"
metrics = f"Accuracy: {avg_similarity:.2f}"

with gr.Blocks() as interface:
    gr.Markdown("# Automated Medical Coding POC")
    # gr.Markdown(metrics)
    with gr.Row():
        with gr.Column():
            text_input = gr.Textbox(label="Input Text")
            submit_button = gr.Button("Submit")
        with gr.Column():
            text_output = gr.Textbox(label="Retrieved Results", lines=10)
    submit_button.click(fn=gradio_interface, inputs=text_input, outputs=text_output)

interface.launch()