Update src/streamlit_app.py

src/streamlit_app.py

@@ -27,10 +27,15 @@ from transformers import CLIPProcessor, CLIPModel
 from datasets import load_dataset, get_dataset_split_names
 from PIL import Image
 import openai
+import comet_llm
+from opik import track
+
+
 
 # ========== 🔑 API Key ==========
 openai.api_key = os.getenv("OPENAI_API_KEY")
-
+os.environ["OPIK_API_KEY"] = os.getenv("OPIK_API_KEY")
+os.environ["OPIK_WORKSPACE"] = os.getenv("OPIK_WORKSPACE")
 # ========== 📥 Load Models ==========
 @st.cache_resource(show_spinner=False)
 def load_models():
@@ -52,9 +57,6 @@ clip_model, clip_processor, text_model = load_models()
 
 # ========== 📥 Load Dataset ==========
 @st.cache_resource(show_spinner=False)
-
-
-
 def load_medical_data():
     available_splits = get_dataset_split_names("univanxx/3mdbench")
     split_to_use = "train" if "train" in available_splits else available_splits[0]
@@ -67,6 +69,8 @@ def load_medical_data():
 
 data = load_medical_data()
 
+from openai import OpenAI
+client = OpenAI(api_key=openai.api_key)
 # Temporary debug display
 #st.write("Dataset columns:", data.features.keys())
 
@@ -103,19 +107,49 @@ TEXT_COLUMN = "complaints"  # or "general_complaint", depending on your needs
 st.title("🩺 Multimodal Medical Chatbot")
 
 query = st.text_input("Enter your medical question or symptom description:")
+uploaded_file = st.file_uploader("Upload an image to find similar medical cases:", type=["png", "jpg", "jpeg"])
+
+@track
+def get_chat_completion_openai(client, prompt: str):
+    return client.chat.completions.create(
+        model="gpt-4o",  # or "gpt-4" if you need the older GPT-4
+        messages=[{"role": "user", "content": prompt}],
+        temperature=0.5,
+        max_tokens=150
+    )
+
+@track
+def get_similar_prompt(query):
+    text_embeddings = embed_dataset_texts(combined_texts)  # cached
+    query_embedding = embed_query_text(query)  # recalculated each time
+
+    cos_scores = util.pytorch_cos_sim(query_embedding, text_embeddings)[0]
+    top_result = torch.topk(cos_scores, k=1)
+    idx = top_result.indices[0].item()
+    return data[idx]
+
+# Cache dataset image embeddings (takes time, so cached)
+@st.cache_data(show_spinner=True)
+def embed_dataset_images(_dataset):
+    features = []
+    for item in _dataset:
+        # Load image from URL/path or raw bytes - adapt this if needed
+        img = item["image"]
+        inputs = clip_processor(images=img, return_tensors="pt")
+        with torch.no_grad():
+            feat = clip_model.get_image_features(**inputs)
+        feat /= feat.norm(p=2, dim=-1, keepdim=True)
+        features.append(feat.cpu())
+    return torch.cat(features, dim=0)
+
+dataset_image_features = embed_dataset_images(data)
 
 if query:
     with st.spinner("Searching medical cases..."):
-        text_embeddings = embed_dataset_texts(combined_texts)  # cached
-        query_embedding = embed_query_text(query)  # recalculated each time
-    #    text_embeddings = embed_dataset_texts(data[TEXT_COLUMN])
-    #   query_embedding = embed_query_text([query])[0]
+
 
         # Compute similarity
-        cos_scores = util.pytorch_cos_sim(query_embedding, text_embeddings)[0]
-        top_result = torch.topk(cos_scores, k=1)
-        idx = top_result.indices[0].item()
-        selected = data[idx]
+        selected = get_similar_prompt(query)
 
         # Show Image
         st.image(selected['image'], caption="Most relevant medical image", use_container_width=True)
@@ -126,20 +160,36 @@ if query:
         # GPT Explanation
         if openai.api_key:
             prompt = f"Explain this case in plain English: {selected[TEXT_COLUMN]}"
-            from openai import OpenAI
-            client = OpenAI(api_key=openai.api_key)
 
-            response = client.chat.completions.create(
-            model="gpt-4o",  # or "gpt-4" if you need the older GPT-4
-            messages=[{"role": "user", "content": prompt}],
-            temperature=0.5,
-            max_tokens=150
-            )
-            explanation = response.choices[0].message.content
+            explanation = get_chat_completion_openai(client, prompt)
+            explanation = explanation.choices[0].message.content
 
             st.markdown(f"### 🤖 Explanation by GPT:\n{explanation}")
         else:
             st.warning("OpenAI API key not found. Please set OPENAI_API_KEY as a secret environment variable.")
 
-st.caption("This chatbot is for educational purposes only and does not provide medical advice.")
+if uploaded_file is not None:
+    query_image = Image.open(uploaded_file).convert("RGB")
+    st.image(query_image, caption="Your uploaded image", use_container_width=True)
+
+    # Embed uploaded image
+    inputs = clip_processor(images=query_image, return_tensors="pt")
+    with torch.no_grad():
+        query_feat = clip_model.get_image_features(**inputs)
+    query_feat /= query_feat.norm(p=2, dim=-1, keepdim=True)
+
+    # Compute cosine similarity
+    similarities = (dataset_image_features @ query_feat.T).squeeze(1)  # [num_dataset_images]
+
+    top_k = 3
+    top_results = torch.topk(similarities, k=top_k)
+
+    st.write(f"Top {top_k} similar medical cases:")
+
+    for rank, idx in enumerate(top_results.indices):
+        score = top_results.values[rank].item()
+        similar_img = data[int(idx)]['image']
+        st.image(similar_img, caption=f"Similarity: {score:.3f}", use_container_width=True)
+        st.markdown(f"**Case description:** {data[int(idx)]['complaints']}")
 
+st.caption("This chatbot is for educational purposes only and does not provide medical advice.")