Update src/streamlit_app.py

src/streamlit_app.py

@@ -2,7 +2,7 @@
 #   ✅ Cache-Safe Multimodal App
 # ================================
 
-import os
+import shutil, os
 
 # ====== Force all cache dirs to /tmp (writable in most environments) ======
 CACHE_BASE = "/tmp/cache"
@@ -13,6 +13,10 @@ os.environ["HF_DATASETS_CACHE"] = f"{CACHE_BASE}/hf_datasets"
 os.environ["TORCH_HOME"] = f"{CACHE_BASE}/torch"
 os.environ["STREAMLIT_CACHE_DIR"] = f"{CACHE_BASE}/streamlit_cache"
 os.environ["STREAMLIT_STATIC_DIR"] = f"{CACHE_BASE}/streamlit_static"
+os.environ["STREAMLIT_CONFIG_DIR"] = "/tmp/.streamlit"
+
+# Create the directories before imports
+os.makedirs(os.environ["STREAMLIT_CONFIG_DIR"], exist_ok=True)
 
 # Create the directories before imports
 for path in os.environ.values():
@@ -26,37 +30,30 @@ 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
 import comet_llm
 from opik import track
 
-os.environ["STREAMLIT_CONFIG_DIR"] = "/tmp/.streamlit"
-os.environ["STREAMLIT_CACHE_DIR"] = f"{CACHE_BASE}/streamlit_cache"
-os.environ["STREAMLIT_STATIC_DIR"] = f"{CACHE_BASE}/streamlit_static"
-
-os.makedirs("/tmp/.streamlit", exist_ok=True)
-
-
 # ========== 🔑 API Key ==========
-openai.api_key = os.getenv("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")
 # ========== 📥 Load Models ==========
 @st.cache_resource(show_spinner=False)
 def load_models():
-    clip_model = CLIPModel.from_pretrained(
+    _clip_model = CLIPModel.from_pretrained(
         "openai/clip-vit-base-patch32",
         cache_dir=os.environ["TRANSFORMERS_CACHE"]
     )
-    clip_processor = CLIPProcessor.from_pretrained(
+    _clip_processor = CLIPProcessor.from_pretrained(
         "openai/clip-vit-base-patch32",
         cache_dir=os.environ["TRANSFORMERS_CACHE"]
     )
-    text_model = SentenceTransformer(
+    _text_model = SentenceTransformer(
         "all-MiniLM-L6-v2",
         cache_folder=os.environ["SENTENCE_TRANSFORMERS_HOME"]
     )
-    return clip_model, clip_processor, text_model
+    return _clip_model, _clip_processor, _text_model
 
 clip_model, clip_processor, text_model = load_models()
 
@@ -72,10 +69,24 @@ def load_medical_data():
     )
     return dataset
 
+# 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_img = clip_processor(images=img, return_tensors="pt")
+        with torch.no_grad():
+            feat = clip_model.get_image_features(**inputs_img)
+        feat /= feat.norm(p=2, dim=-1, keepdim=True)
+        features.append(feat.cpu())
+    return torch.cat(features, dim=0)
+
 data = load_medical_data()
+dataset_image_features = embed_dataset_images(data)
 
-from openai import OpenAI
-client = OpenAI(api_key=openai.api_key)
+client = OpenAI(api_key=OpenAI.api_key)
 # Temporary debug display
 #st.write("Dataset columns:", data.features.keys())
 
@@ -102,17 +113,37 @@ combined_texts = prepare_combined_texts(data)
 def embed_dataset_texts(_texts):
     return text_model.encode(_texts, convert_to_tensor=True)
 
-def embed_query_text(query):
-    return text_model.encode([query], convert_to_tensor=True)[0]
+def embed_query_text(_query):
+    return text_model.encode([_query], convert_to_tensor=True)[0]
+
+@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=425
+    )
+
+@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]
+
 
 # Pick which text column to use
 TEXT_COLUMN = "complaints"  # or "general_complaint", depending on your needs
 
 # ========== 🧑‍⚕️ App UI ==========
-st.title("🩺 Dr_Q_bot - Multimodal Medical Chatbot")
+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"])
+uploaded_files = st.file_uploader("Upload an image to find similar medical cases:", type=["png", "jpg", "jpeg"], accept_multiple_files=True)
 
 # Add author info in the sidebar
 with st.sidebar:
@@ -120,44 +151,8 @@ with st.sidebar:
     st.markdown("**Vasan Iyer**")
     st.markdown("**Eric J Giacomucci**")
     st.markdown("[GitHub](https://github.com/Vaiy108)")
-    st.markdown("[LinkedIn](https://linkedin.com/in/vasan-iyer)") 
-
-@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)
+    st.markdown("[LinkedIn](https://linkedin.com/in/vasan-iyer)")
 
-#if query:
 if st.button("Submit") and query:
     with st.spinner("Searching medical cases..."):
 
@@ -172,7 +167,7 @@ if st.button("Submit") and query:
         st.markdown(f"**Case Description:** {selected[TEXT_COLUMN]}")
 
         # GPT Explanation
-        if openai.api_key:
+        if OpenAI.api_key:
             prompt = f"Explain this case in plain English: {selected[TEXT_COLUMN]}"
 
             explanation = get_chat_completion_openai(client, prompt)
@@ -182,32 +177,37 @@ if st.button("Submit") and query:
         else:
             st.warning("OpenAI API key not found. Please set OPENAI_API_KEY as a secret environment variable.")
 
-if uploaded_file is not None:
-    print('uploading file')
-    print(uploaded_file)
-    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]
+if uploaded_files is not None:
+    with st.spinner("Searching medical cases..."):
+        st.write(f"Number of files: {len(uploaded_files)}")
+
+        if len(uploaded_files) > 0:
+            print(uploaded_files)
+            uploaded_file = uploaded_files[0]
+            st.write(f'uploading file {uploaded_file.name}')
+            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)
+            top_k = 3
+            top_results = torch.topk(similarities, k=top_k)
 
-    st.write(f"Top {top_k} similar medical cases:")
+            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']}")
-else:
-    print("no image")
+            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.")
+st.caption("This chatbot is for educational purposes only and does not provide medical advice.")