add comet on all other operations

src/streamlit_app.py

@@ -1,5 +1,5 @@
 # ================================
-#   ✅ Cache-Safe Multimodal App
+#   ✅ Cache-Safe Multimodal App with Full Opik Tracking
 # ================================
 
 import os
@@ -27,15 +27,16 @@ 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
-
-
+from opik import track, log_event
+from openai import OpenAI
 
 # ========== 🔑 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():
@@ -57,9 +58,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]
@@ -71,67 +69,83 @@ def load_medical_data():
     return dataset
 
 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())
-
-# 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])
+TEXT_COLUMN = "complaints" if "complaints" in data.features else list(data.features.keys())[0]
 
 # ========== 🧠 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 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
+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",
         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("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.")
+    with st.spinner("Processing your query..."):
+        session_data = process_query(query)
+        st.markdown(f"### 🤖 Explanation by GPT:\n{session_data['gpt_explanation']}")
 
-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.")