Update app.py

app.py

@@ -1,21 +1,24 @@
+from transformers import AutoTokenizer, AutoModel
+import torch
 import gradio as gr
-from sentence_transformers import SentenceTransformer
 
-# ✅ Load a model that outputs 1024-dim vectors
-model = SentenceTransformer('intfloat/e5-large')
+# Load Bio_ClinicalBERT
+tokenizer = AutoTokenizer.from_pretrained("emilyalsentzer/Bio_ClinicalBERT")
+model = AutoModel.from_pretrained("emilyalsentzer/Bio_ClinicalBERT")
 
 def embed_text(text):
-    # Optionally prepend "passage: " if using e5 models
-    processed_text = "passage: " + text.strip()
-    embedding = model.encode(processed_text).tolist()
+    inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True)
+    with torch.no_grad():
+        outputs = model(**inputs)
+    # Mean pooling
+    embedding = outputs.last_hidden_state.mean(dim=1).squeeze().tolist()
     return embedding
 
-# Gradio interface
 iface = gr.Interface(
     fn=embed_text,
     inputs=gr.Textbox(lines=5, label="Enter patient text"),
     outputs="json",
-    title="Clinical Text Embedding API (1024-dim)"
+    title="Clinical Text Embedding API (Bio_ClinicalBERT)"
 )
 
 iface.launch()