Update src/streamlit_app.py

src/streamlit_app.py

@@ -1,8 +1,10 @@
 # streamlit_app.py
 import os
 import streamlit as st
-from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
+import torch
+from transformers import AutoTokenizer, pipeline
 from peft import PeftModel
+from unsloth import FastLanguageModel
 
 # -----------------------------
 # Ensure cache dirs are writable in Spaces
@@ -13,17 +15,45 @@ os.environ.setdefault("HF_DATASETS_CACHE", "/tmp/huggingface/datasets")
 os.environ.setdefault("HUGGINGFACE_HUB_CACHE", "/tmp/huggingface/hub")
 os.environ.setdefault("XDG_CACHE_HOME", "/tmp/huggingface")
 
-# Base and adapter model IDs
-BASE_MODEL = "microsoft/BioGPT-Large-PubMedQA"
-ADAPTER_MODEL = "kirubel1738/biogpt-pubmedqa-finetuned"
+# -----------------------------
+# Model IDs
+# -----------------------------
+BASE_MODEL = "unsloth/llama-3-8b-bnb-4bit"
+ADAPTER_MODEL = "kirubel1738/llama3-biology-qa"
 
+# -----------------------------
+# Load model once
+# -----------------------------
 @st.cache_resource
 def load_model():
-    """Load BioGPT with PubMedQA adapter on CPU."""
+    """Load LLaMA-3 8B with PEFT adapter entirely on CPU."""
+    st.info("Loading LLaMA-3 model on CPU... This may take a while.")
+
+    # Load tokenizer
     tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL)
-    base_model = AutoModelForCausalLM.from_pretrained(BASE_MODEL, device_map=None)
-    model = PeftModel.from_pretrained(base_model, ADAPTER_MODEL)  # apply adapter
-    generator = pipeline("text-generation", model=model, tokenizer=tokenizer, device=-1)
+
+    # Load base model in 4-bit on CPU
+    base_model, _ = FastLanguageModel.from_pretrained(
+        model_name=BASE_MODEL,
+        max_seq_length=2048,
+        dtype=None,
+        load_in_4bit=True,
+        device_map={"": "cpu"}  # force CPU
+    )
+
+    # Apply adapter
+    model = PeftModel.from_pretrained(base_model, ADAPTER_MODEL)
+
+    # Text-generation pipeline on CPU
+    generator = pipeline(
+        "text-generation",
+        model=model,
+        tokenizer=tokenizer,
+        device=-1,  # CPU
+        max_new_tokens=256,
+        do_sample=True,
+        temperature=0.7
+    )
     return generator
 
 # Load once
@@ -32,24 +62,17 @@ generator = load_model()
 # -----------------------------
 # Streamlit UI
 # -----------------------------
-st.set_page_config(page_title="BioGPT — PubMedQA demo", layout="centered")
-st.title("🧬 BioGPT — PubMedQA Demo")
-
-st.write("Ask a biomedical question and get an answer ")
-st.write(" generated by BioGPT-Large-PubMedQA fine-tuned on MMLU + SciQ dataset.")
+st.set_page_config(page_title="LLaMA-3 Biology QA", layout="centered")
+st.title("🧬 LLaMA-3 — Biology QA Demo")
+st.write("Ask a biology question and get an answer generated by LLaMA-3 fine-tuned on the Biology QA dataset.")
 
-user_input = st.text_area("Enter your biomedical question:", height=150)
+user_input = st.text_area("Enter your biology question:", height=150)
 
 if st.button("Get Answer"):
     if user_input.strip():
         with st.spinner("Generating answer..."):
             try:
-                result = generator(
-                    user_input,
-                    max_new_tokens=128,
-                    do_sample=True,
-                    temperature=0.7
-                )
+                result = generator(user_input)
                 output_text = result[0]["generated_text"]
                 st.success("Answer:")
                 st.write(output_text)
@@ -59,8 +82,4 @@ if st.button("Get Answer"):
         st.warning("Please enter a question.")
 
 st.markdown("---")
-st.caption("Model: microsoft/biogpt + adapter kirubel1738/biogpt-pubmedqa-finetuned | Runs on CPU")
-
-
-
-
+st.caption(f"Model: {BASE_MODEL} + adapter {ADAPTER_MODEL} | Runs on CPU")