Update src/streamlit_app.py

src/streamlit_app.py

@@ -1,20 +1,25 @@
+# ================================
+#   ✅ Cache-Safe Multimodal App
+# ================================
+
 import os
-# ✅ Set all relevant cache directories to a writable location
-os.environ["HF_HOME"] = "/tmp/cache"
-os.environ["TRANSFORMERS_CACHE"] = "/tmp/cache/transformers"
-os.environ["SENTENCE_TRANSFORMERS_HOME"] = "/tmp/cache/sentence_transformers"
-os.environ["HF_DATASETS_CACHE"] = "/tmp/cache/hf_datasets"
-os.environ["TORCH_HOME"] = "/tmp/cache/torch"
-
-# ✅ Create the directories if they don't exist
-for path in [
-    "/tmp/cache",
-    "/tmp/cache/transformers",
-    "/tmp/cache/sentence_transformers",
-    "/tmp/cache/hf_datasets",
-    "/tmp/cache/torch"
-]:
-    os.makedirs(path, exist_ok=True)
+
+# ====== Force all cache dirs to /tmp (writable in most environments) ======
+CACHE_BASE = "/tmp/cache"
+os.environ["HF_HOME"] = f"{CACHE_BASE}/hf_home"
+os.environ["TRANSFORMERS_CACHE"] = f"{CACHE_BASE}/transformers"
+os.environ["SENTENCE_TRANSFORMERS_HOME"] = f"{CACHE_BASE}/sentence_transformers"
+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"
+
+# Create the directories before imports
+for path in os.environ.values():
+    if path.startswith(CACHE_BASE):
+        os.makedirs(path, exist_ok=True)
+
+# ====== Imports ======
 import streamlit as st
 import torch
 from sentence_transformers import SentenceTransformer, util
@@ -22,9 +27,6 @@ from transformers import CLIPProcessor, CLIPModel
 from datasets import load_dataset
 from PIL import Image
 import openai
-from io import BytesIO
-import requests
-
 
 # ========== 🔑 API Key ==========
 openai.api_key = os.getenv("OPENAI_API_KEY")
@@ -32,9 +34,18 @@ openai.api_key = os.getenv("OPENAI_API_KEY")
 # ========== 📥 Load Models ==========
 @st.cache_resource(show_spinner=False)
 def load_models():
-    clip_model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32")
-    clip_processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32")
-    text_model = SentenceTransformer("all-MiniLM-L6-v2")
+    clip_model = CLIPModel.from_pretrained(
+        "openai/clip-vit-base-patch32",
+        cache_dir=os.environ["TRANSFORMERS_CACHE"]
+    )
+    clip_processor = CLIPProcessor.from_pretrained(
+        "openai/clip-vit-base-patch32",
+        cache_dir=os.environ["TRANSFORMERS_CACHE"]
+    )
+    text_model = SentenceTransformer(
+        "all-MiniLM-L6-v2",
+        cache_folder=os.environ["SENTENCE_TRANSFORMERS_HOME"]
+    )
     return clip_model, clip_processor, text_model
 
 clip_model, clip_processor, text_model = load_models()
@@ -42,7 +53,11 @@ clip_model, clip_processor, text_model = load_models()
 # ========== 📥 Load Dataset ==========
 @st.cache_resource(show_spinner=False)
 def load_medical_data():
-    dataset = load_dataset("univanxx/3mdbench", split="train")
+    dataset = load_dataset(
+        "univanxx/3mdbench",
+        split="train",
+        cache_dir=os.environ["HF_DATASETS_CACHE"]
+    )
     return dataset
 
 data = load_medical_data()
@@ -79,9 +94,7 @@ if query:
             prompt = f"Explain this case in plain English: {selected['description']}"
             response = openai.ChatCompletion.create(
                 model="gpt-4",
-                messages=[
-                    {"role": "user", "content": prompt}
-                ],
+                messages=[{"role": "user", "content": prompt}],
                 temperature=0.5,
                 max_tokens=150
             )
@@ -90,5 +103,5 @@ if query:
         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.")
+