Update app.py

app.py

@@ -1,24 +1,35 @@
-from transformers import AutoTokenizer, AutoModel
 import torch
+from transformers import AutoTokenizer, AutoModel
 import gradio as gr
 
-# Load Bio_ClinicalBERT
+# Load Bio_ClinicalBERT model and tokenizer
 tokenizer = AutoTokenizer.from_pretrained("emilyalsentzer/Bio_ClinicalBERT")
 model = AutoModel.from_pretrained("emilyalsentzer/Bio_ClinicalBERT")
+model.eval()  # Disable dropout for inference
+
+# Attention-masked mean pooling function
+def mean_pooling(model_output, attention_mask):
+    token_embeddings = model_output.last_hidden_state  # (batch_size, seq_len, hidden_dim)
+    input_mask_expanded = attention_mask.unsqueeze(-1).expand(token_embeddings.size()).float()
+    pooled = torch.sum(token_embeddings * input_mask_expanded, dim=1) / \
+             torch.clamp(input_mask_expanded.sum(dim=1), min=1e-9)
+    return pooled
 
+# Embedding function
 def embed_text(text):
-    inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True)
+    inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True, max_length=512)
     with torch.no_grad():
         outputs = model(**inputs)
-    # Mean pooling
-    embedding = outputs.last_hidden_state.mean(dim=1).squeeze().tolist()
+    embedding = mean_pooling(outputs, inputs['attention_mask']).squeeze().tolist()
     return embedding
 
+# Gradio interface
 iface = gr.Interface(
     fn=embed_text,
-    inputs=gr.Textbox(lines=5, label="Enter patient text"),
+    inputs=gr.Textbox(lines=5, label="Enter clinical text"),
     outputs="json",
-    title="Clinical Text Embedding API (Bio_ClinicalBERT)"
+    title="Clinical Text Embedding API",
+    description="Generate dense vector embeddings using Bio_ClinicalBERT with attention-masked mean pooling"
 )
 
 iface.launch()