|
import os |
|
import torch |
|
from datasets import load_dataset, DatasetDict |
|
from transformers import AutoTokenizer, AutoModel |
|
import chromadb |
|
import gradio as gr |
|
from sklearn.metrics import precision_score, recall_score, f1_score |
|
|
|
|
|
def meanpooling(output, mask): |
|
embeddings = output[0] |
|
mask = mask.unsqueeze(-1).expand(embeddings.size()).float() |
|
return torch.sum(embeddings * mask, 1) / torch.clamp(mask.sum(1), min=1e-9) |
|
|
|
|
|
dataset = load_dataset("thankrandomness/mimic-iii-sample") |
|
|
|
|
|
split_dataset = dataset['train'].train_test_split(test_size=0.2, seed=42) |
|
dataset = DatasetDict({ |
|
'train': split_dataset['train'], |
|
'validation': split_dataset['test'] |
|
}) |
|
|
|
|
|
tokenizer = AutoTokenizer.from_pretrained("neuml/pubmedbert-base-embeddings-matryoshka") |
|
model = AutoModel.from_pretrained("neuml/pubmedbert-base-embeddings-matryoshka") |
|
|
|
|
|
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']) |
|
return embeddings.numpy().tolist() |
|
|
|
|
|
client = chromadb.Client() |
|
collection = client.create_collection(name="pubmedbert_matryoshka_embeddings") |
|
|
|
|
|
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] |
|
|
|
|
|
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_data(dataset['train']) |
|
|
|
|
|
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 = [] |
|
for metadata, distance in zip(results['metadatas'][0], results['distances'][0]): |
|
output.append({ |
|
"similarity_score": distance, |
|
"code": metadata['code'], |
|
"code_system": metadata['code_system'], |
|
"description": metadata['description'] |
|
}) |
|
return output |
|
|
|
|
|
def evaluate_efficiency(dataset_split): |
|
y_true = [] |
|
y_pred = [] |
|
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] |
|
|
|
|
|
y_true.extend(annotations_list) |
|
|
|
y_pred.extend(retrieved_codes[:len(annotations_list)]) |
|
|
|
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] |
|
|
|
precision = precision_score(y_true, y_pred, average='macro') |
|
recall = recall_score(y_true, y_pred, average='macro') |
|
f1 = f1_score(y_true, y_pred, average='macro') |
|
|
|
return precision, recall, f1 |
|
|
|
|
|
precision, recall, f1 = evaluate_efficiency(dataset['validation']) |
|
|
|
|
|
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) |
|
|
|
|
|
metrics = f"Precision: {precision:.2f}, Recall: {recall:.2f}, F1 Score: {f1:.2f}" |
|
|
|
with gr.Blocks() as interface: |
|
gr.Markdown("# Text Retrieval with Efficiency Metrics") |
|
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() |
|
|