Update src/streamlit_app.py

src/streamlit_app.py

@@ -1,156 +1,108 @@
-# streamlit_app.py
+import streamlit as st
+from transformers import AutoTokenizer, AutoModelForCausalLM
+import torch
 import os
-import json
-import time
-
-# -----------------------------
-# IMPORTANT: set cache dirs BEFORE importing transformers/huggingface_hub
-# -----------------------------
-os.environ.setdefault("HF_HOME", os.environ.get("HF_HOME", "/tmp/huggingface"))
-os.environ.setdefault("TRANSFORMERS_CACHE", os.environ.get("TRANSFORMERS_CACHE", "/tmp/huggingface/transformers"))
-os.environ.setdefault("HF_DATASETS_CACHE", os.environ.get("HF_DATASETS_CACHE", "/tmp/huggingface/datasets"))
-os.environ.setdefault("HUGGINGFACE_HUB_CACHE", os.environ.get("HUGGINGFACE_HUB_CACHE", "/tmp/huggingface/hub"))
-os.environ.setdefault("XDG_CACHE_HOME", os.environ.get("XDG_CACHE_HOME", "/tmp/huggingface"))
-os.environ.setdefault("HOME", os.environ.get("HOME", "/tmp"))
-
-# create cache dirs (best-effort)
-for d in [os.environ["HF_HOME"], os.environ["TRANSFORMERS_CACHE"], os.environ["HF_DATASETS_CACHE"], os.environ["HUGGINGFACE_HUB_CACHE"]]:
+import shutil
+
+# Define the custom cache directory for Hugging Face models
+cache_dir = "/tmp/biogpt_app_cache"
+
+# --- PROACTIVE CACHE CLEARING ---
+# Set environment variables to point Hugging Face and Streamlit to our custom cache directory
+# This is done to prevent PermissionErrors in read-only environments.
+os.environ["STREAMLIT_CACHE_DIR"] = "/tmp/streamlit_cache"
+os.environ["HF_HOME"] = cache_dir
+os.environ["TRANSFORMERS_CACHE"] = cache_dir
+os.environ["XDG_CACHE_HOME"] = cache_dir
+os.environ["STREAMLIT_BROWSER_GATHER_USAGE_STATS"] = "false"
+
+# Clear the cache directory before attempting to download the model.
+if os.path.exists(cache_dir):
     try:
-        os.makedirs(d, exist_ok=True)
-        os.chmod(d, 0o777)
-    except Exception:
-        pass
-
-import streamlit as st
-import requests
-
-# Optional heavy imports will be inside local-model branch
-LOCAL_MODE = os.environ.get("USE_LOCAL_MODEL", "0") == "1"
-
-# default model id the user provided; keep as-is
-DEFAULT_MODEL_ID = "kirubel1738/biogpt-pubmedqa-finetuned"
-
-st.set_page_config(page_title="BioGPT (PubMedQA) demo", layout="centered")
-
-st.title("BioGPT — PubMedQA demo")
-st.caption("Defaults to the Hugging Face Inference API (recommended for Spaces / CPU).")
-
-st.markdown(
-    """
-**How it works**
-- By default the app will call Hugging Face's Inference API for the model you specify (fast and avoids memory issues).
-- If you set `USE_LOCAL_MODEL=1` in your environment, the app will attempt to load the model locally using `transformers` (only for GPUs/large memory machines).
-"""
-)
-
-col1, col2 = st.columns([3,1])
-
-with col1:
-    model_id = st.text_input("Model repo id", value=DEFAULT_MODEL_ID, help="Hugging Face repo id (e.g. username/modelname).")
-    prompt = st.text_area("Question / prompt", height=180, placeholder="Enter a PubMed-style question or prompt...")
-with col2:
-    max_new_tokens = st.slider("Max new tokens", 16, 1024, 128)
-    temperature = st.slider("Temperature", 0.0, 1.5, 0.0, step=0.05)
-    method = st.radio("Run method", ("Inference API (recommended)", "Local model (heavy)"), index=0)
+        st.info("Clearing old cache to ensure a fresh download...")
+        shutil.rmtree(cache_dir)
+    except Exception as e:
+        st.error(f"Failed to clear old cache. Please check directory permissions. Error: {e}")
+        st.stop()
 
-# override radio if user set USE_LOCAL_MODEL env var
-if LOCAL_MODE:
-    method = "Local model (heavy)"
+# Ensure the new cache directory exists before model loading
+try:
+    os.makedirs(cache_dir, exist_ok=True)
+except Exception as e:
+    st.error(f"Failed to create cache directory at {cache_dir}. Error: {e}")
+    st.stop()
 
-hf_token = os.environ.get("HUGGINGFACE_HUB_TOKEN") or os.environ.get("HF_TOKEN") or os.environ.get("HUGGINGFACE_API_TOKEN")
+st.set_page_config(page_title="BioGPT-PubMedQA Chatbot", layout="centered")
+st.title("🧬 BioGPT-PubMedQA Chatbot")
+st.write("A fine-tuned BioGPT model for biomedical Q&A.")
 
-def call_inference_api(model_id: str, prompt: str, max_new_tokens: int, temperature: float):
+# Load model once using Streamlit's resource caching
+@st.cache_resource
+def load_model(cache_directory):
     """
-    Simple POST to Hugging Face Inference API.
-    If you want to use the InferenceClient from huggingface_hub you can swap this.
+    Loads the tokenizer and model from Hugging Face Hub,
+    explicitly using the specified cache directory.
     """
-    api_url = f"https://api-inference.huggingface.co/models/{model_id}"
-    headers = {"Authorization": f"Bearer {hf_token}"} if hf_token else {}
-    payload = {
-        "inputs": prompt,
-        "parameters": {"max_new_tokens": max_new_tokens, "temperature": temperature},
-        "options": {"wait_for_model": True}
-    }
+    model_name = "kirubel1738/biogpt-pubmedqa-finetuned"
+    
     try:
-        r = requests.post(api_url, headers=headers, json=payload, timeout=120)
+        tokenizer = AutoTokenizer.from_pretrained(model_name, cache_dir=cache_directory)
+        model = AutoModelForCausalLM.from_pretrained(
+            model_name,
+            torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
+            device_map="auto",
+            cache_dir=cache_directory
+        )
+        return tokenizer, model
     except Exception as e:
-        return False, f"Request failed: {e}"
-    if r.status_code != 200:
-        try:
-            error = r.json()
-        except Exception:
-            error = r.text
-        return False, f"API error ({r.status_code}): {error}"
-    try:
-        resp = r.json()
-        # handle several possible response schemas
-        if isinstance(resp, dict) and "error" in resp:
-            return False, resp["error"]
-        # often it's a list of dicts with 'generated_text'
-        if isinstance(resp, list):
-            out_texts = []
-            for item in resp:
-                if isinstance(item, dict):
-                    # common key: 'generated_text'
-                    for k in ("generated_text", "text", "content"):
-                        if k in item:
-                            out_texts.append(item[k])
-                            break
-                    else:
-                        out_texts.append(json.dumps(item))
-                else:
-                    out_texts.append(str(item))
-            return True, "\n\n".join(out_texts)
-        # fallback
-        return True, str(resp)
-    except Exception as e:
-        return False, f"Could not parse response: {e}"
-
-# Local model loader (only if method chosen)
-generator = None
-if method.startswith("Local"):
-    st.warning("Local model mode selected — this requires transformers + torch and lots of RAM/GPU. Only use if you know the model fits your hardware.")
-    try:
-        from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
-        import torch
-        device = 0 if torch.cuda.is_available() else -1
-        st.info(f"torch.cuda.is_available={torch.cuda.is_available()} -- device set to {device}")
-        with st.spinner("Loading tokenizer & model (this can take a while)..."):
-            tokenizer = AutoTokenizer.from_pretrained(model_id, cache_dir=os.environ.get("TRANSFORMERS_CACHE"))
-            model = AutoModelForCausalLM.from_pretrained(model_id, cache_dir=os.environ.get("TRANSFORMERS_CACHE"), low_cpu_mem_usage=True)
-            generator = pipeline("text-generation", model=model, tokenizer=tokenizer, device=device)
-    except Exception as e:
-        st.error(f"Local model load failed: {e}")
+        st.error(f"Failed to load model. Please ensure the model name is correct and it is publicly accessible.")
+        st.exception(e)
         st.stop()
 
-if st.button("Generate"):
-    if not prompt or prompt.strip() == "":
-        st.error("Please enter a prompt.")
-        st.stop()
-
-    if method.startswith("Inference"):
-        if ("kirubel1738/biogpt-pubmedqa-finetuned" in model_id) and not hf_token:
-            st.info("If the model is private or rate-limited, set HUGGINGFACE_HUB_TOKEN as a secret in Spaces or as an env var locally.")
-        with st.spinner("Querying Hugging Face Inference API..."):
-            ok, out = call_inference_api(model_id, prompt, max_new_tokens, float(temperature))
-        if not ok:
-            st.error(out)
+# Load the model, passing the cache directory
+try:
+    tokenizer, model = load_model(cache_dir)
+except Exception as e:
+    st.error(f"An unexpected error occurred during model loading: {e}")
+    st.stop()
+
+# Maintain chat history
+if "messages" not in st.session_state:
+    st.session_state["messages"] = []
+
+# Display chat history
+for msg in st.session_state["messages"]:
+    with st.chat_message(msg["role"]):
+        st.markdown(msg["content"])
+
+# Input box for user
+if prompt := st.chat_input("Ask me a biomedical question..."):
+    st.session_state["messages"].append({"role": "user", "content": prompt})
+    
+    with st.chat_message("user"):
+        st.markdown(prompt)
+    
+    formatted_prompt = f"""### Question:{prompt}### Answer:"""
+    inputs = tokenizer(formatted_prompt, return_tensors="pt").to(model.device)
+    
+    with st.spinner("Thinking..."):
+        with torch.no_grad():
+            outputs = model.generate(
+                **inputs,
+                max_new_tokens=200,
+                do_sample=True,
+                temperature=0.7,
+                top_p=0.9,
+                eos_token_id=tokenizer.eos_token_id,
+            )
+        decoded = tokenizer.decode(outputs[0], skip_special_tokens=True)
+        
+        if "### Answer:" in decoded:
+            answer = decoded.split("### Answer:")[-1].strip()
         else:
-            st.success("Done")
-            st.text_area("Model output", value=out, height=320)
-    else:
-        # local model generation
-        try:
-            with st.spinner("Running local generation..."):
-                results = generator(prompt, max_new_tokens=max_new_tokens, do_sample=True, temperature=temperature)
-                if isinstance(results, list) and len(results) > 0 and "generated_text" in results[0]:
-                    out = results[0]["generated_text"]
-                else:
-                    out = str(results)
-            st.success("Done")
-            st.text_area("Model output", value=out, height=320)
-        except Exception as e:
-            st.error(f"Local generation failed: {e}")
-
-st.markdown("---")
-st.caption("If you run into permissions errors in Spaces, ensure the HF cache env vars above point to a writable directory (we already set them to /tmp/huggingface in this container).")
+            answer = decoded.strip()
+    
+    st.session_state["messages"].append({"role": "assistant", "content": answer})
+    
+    with st.chat_message("assistant"):
+        st.markdown(answer)