revert

src/streamlit_app.py

@@ -1,15 +1,9 @@
 # ================================
-#   ✅ Cache-Safe Multimodal App with Full Opik Tracking
+#   ✅ Cache-Safe Multimodal App
 # ================================
 
 import os
 
-# ---- Disable Comet auto-patching (MUST be set BEFORE importing openai/comet_llm/comet_ml) ----
-# Disable all Comet auto-logging / monkey-patching
-os.environ["COMET_DISABLE_AUTO_LOGGING"] = "1"
-# Optionally: only disable LLM auto-logging
-os.environ["COMET_DISABLE_AUTO_LOGGING_LLM"] = "1"
-
 # ====== Force all cache dirs to /tmp (writable in most environments) ======
 CACHE_BASE = "/tmp/cache"
 os.environ["HF_HOME"] = f"{CACHE_BASE}/hf_home"
@@ -22,29 +16,26 @@ os.environ["STREAMLIT_STATIC_DIR"] = f"{CACHE_BASE}/streamlit_static"
 
 # Create the directories before imports
 for path in os.environ.values():
-    if isinstance(path, str) and path.startswith(CACHE_BASE):
+    if path.startswith(CACHE_BASE):
         os.makedirs(path, exist_ok=True)
 
-# ====== Now safe to import libraries ======
+# ====== Imports ======
 import streamlit as st
 import torch
 from sentence_transformers import SentenceTransformer, util
 from transformers import CLIPProcessor, CLIPModel
 from datasets import load_dataset, get_dataset_split_names
 from PIL import Image
-
 import openai
-from openai import OpenAI   # OK to import after openai is present
-from opik import track, log_event
+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")
-
-client = OpenAI(api_key=openai.api_key)
-
 # ========== 📥 Load Models ==========
 @st.cache_resource(show_spinner=False)
 def load_models():
@@ -77,83 +68,67 @@ def load_medical_data():
     return dataset
 
 data = load_medical_data()
-TEXT_COLUMN = "complaints" if "complaints" in data.features else list(data.features.keys())[0]
+
+from openai import OpenAI
+client = OpenAI(api_key=openai.api_key)
+# Temporary debug display
+#st.write("Dataset columns:", data.features.keys())
+
+# After seeing the real column name, let's say it's "text" instead of "description":
+text_field = "text" if "text" in data.features else list(data.features.keys())[0]
+
+# Then use dynamic access:
+#text_embeddings = embed_texts(data[text_field])
 
 # ========== 🧠 Embedding Function ==========
+@st.cache_data(show_spinner=False)
+def embed_texts(_texts):
+    return text_model.encode(_texts, convert_to_tensor=True)
+
+# Pick which text column to use
+TEXT_COLUMN = "complaints"  # or "general_complaint", depending on your needs
+
+# ========== 🧑‍⚕️ App UI ==========
+st.title("🩺 Multimodal Medical Chatbot")
+
+query = st.text_input("Enter your medical question or symptom description:")
+
 @track
-def embed_texts_tracked(texts, model_name="all-MiniLM-L6-v2"):
-    embeddings = text_model.encode(texts, convert_to_tensor=True)
-    log_event("embedding_generated", {
-        "model": model_name,
-        "num_texts": len(texts),
-        "embedding_shape": list(embeddings.shape)
-    })
-    return embeddings
-
-# ========== 🔍 Case Selection ==========
-@track
-def select_top_case(query_embedding, text_embeddings, k=1):
-    cos_scores = util.pytorch_cos_sim(query_embedding, text_embeddings)[0]
-    top_result = torch.topk(cos_scores, k=k)
-    idx = top_result.indices[0].item()
-    score = float(top_result.values[0].item())
-    log_event("case_selected", {
-        "case_index": idx,
-        "similarity_score": score
-    })
-    return idx, score
-
-# ========== 🖼️ Display Case ==========
-@track
-def display_case(case):
-    st.image(case['image'], caption="Most relevant medical image", use_container_width=True)
-    st.markdown(f"**Case Description:** {case[TEXT_COLUMN]}")
-    log_event("case_displayed", {
-        "case_id": case.get("id", None),
-        "description_preview": case[TEXT_COLUMN][:100] + "..."
-    })
-    return case
-
-# ========== 🤖 GPT Completion ==========
-@track
-def get_chat_completion_openai(client, prompt: str, case_id=None):
-    response = client.chat.completions.create(
-        model="gpt-4o",
+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
     )
-    answer = response.choices[0].message.content
-    log_event("gpt_response", {
-        "case_id": case_id,
-        "prompt_length": len(prompt),
-        "response_length": len(answer)
-    })
-    return answer
-
-# ========== 🔄 Full Query Processing ==========
-@track
-def process_query(query):
-    text_embeddings = embed_texts_tracked(data[TEXT_COLUMN])
-    query_embedding = embed_texts_tracked([query])[0]
-    idx, score = select_top_case(query_embedding, text_embeddings)
-    case = display_case(data[idx])
-    explanation = get_chat_completion_openai(client, f"Explain this case in plain English: {case[TEXT_COLUMN]}", case_id=idx)
-    return {
-        "query": query,
-        "case_id": idx,
-        "similarity_score": score,
-        "gpt_explanation": explanation
-    }
-
-# ========== 🖥️ Streamlit UI ==========
-st.title("🩺 Multimodal Medical Chatbot")
 
-query = st.text_input("Enter your medical question or symptom description:")
 
 if query:
-    with st.spinner("Processing your query..."):
-        session_data = process_query(query)
-        st.markdown(f"### 🤖 Explanation by GPT:\n{session_data['gpt_explanation']}")
+    with st.spinner("Searching medical cases..."):
+        text_embeddings = embed_texts(data[TEXT_COLUMN])
+        query_embedding = embed_texts([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]
+
+        # Show Image
+        st.image(selected['image'], caption="Most relevant medical image", use_container_width=True)
+
+        # Show Text
+        st.markdown(f"**Case Description:** {selected[TEXT_COLUMN]}")
+
+        # GPT Explanation
+        if openai.api_key:
+            prompt = f"Explain this case in plain English: {selected[TEXT_COLUMN]}"
+
+            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.")
+st.caption("This chatbot is for educational purposes only and does not provide medical advice.")