Update src/streamlit_app.py

src/streamlit_app.py

@@ -1,111 +1,157 @@
-import streamlit as st
-from transformers import AutoTokenizer, AutoModelForCausalLM
-import torch
+
+# streamlit_app.py
 import os
-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):
+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"]]:
     try:
-        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()
+        os.makedirs(d, exist_ok=True)
+        os.chmod(d, 0o777)
+    except Exception:
+        pass
 
-# 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()
+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 Chatbot", layout="centered")
-st.title("🧬 BioGPT-PubMedQA Chatbot")
-st.write("A fine-tuned BioGPT model for biomedical Q&A.")
+st.set_page_config(page_title="BioGPT (PubMedQA) demo", layout="centered")
 
-# Detect device (CPU or GPU)
-device = "cuda" if torch.cuda.is_available() else "cpu"
-dtype = torch.float16 if torch.cuda.is_available() else torch.float32
+st.title("BioGPT — PubMedQA demo")
+st.caption("Defaults to the Hugging Face Inference API (recommended for Spaces / CPU).")
 
-# Load model once using Streamlit's resource caching
-@st.cache_resource
-def load_model(cache_directory):
+st.markdown(
     """
-    Loads the tokenizer and model from Hugging Face Hub,
-    explicitly using the specified cache directory.
+**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)
+
+# override radio if user set USE_LOCAL_MODEL env var
+if LOCAL_MODE:
+    method = "Local model (heavy)"
+
+hf_token = os.environ.get("HUGGINGFACE_HUB_TOKEN") or os.environ.get("HF_TOKEN") or os.environ.get("HUGGINGFACE_API_TOKEN")
+
+def call_inference_api(model_id: str, prompt: str, max_new_tokens: int, temperature: float):
+    """
+    Simple POST to Hugging Face Inference API.
+    If you want to use the InferenceClient from huggingface_hub you can swap this.
     """
-    model_name = "kirubel1738/biogpt-pubmedqa-finetuned"
-    
+    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}
+    }
     try:
-        tokenizer = AutoTokenizer.from_pretrained(model_name, cache_dir=cache_directory)
-        model = AutoModelForCausalLM.from_pretrained(
-            model_name,
-            torch_dtype=dtype,
-            cache_dir=cache_directory
-        ).to(device)
-        return tokenizer, model
+        r = requests.post(api_url, headers=headers, json=payload, timeout=120)
     except Exception as e:
-        st.error("Failed to load model. Please ensure the model name is correct and it is publicly accessible.")
-        st.exception(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.stop()
 
-# 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(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()
+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)
         else:
-            answer = decoded.strip()
-    
-    st.session_state["messages"].append({"role": "assistant", "content": answer})
-    
-    with st.chat_message("assistant"):
-        st.markdown(answer)
+            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).")