Update src/streamlit_app.py

src/streamlit_app.py

@@ -2,9 +2,8 @@
 import os
 import streamlit as st
 import torch
-from transformers import AutoTokenizer, pipeline
+from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
 from peft import PeftModel
-from unsloth import FastLanguageModel
 
 # -----------------------------
 # Ensure cache dirs are writable in Spaces
@@ -15,44 +14,25 @@ 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")
 
-# -----------------------------
-# Model IDs
-# -----------------------------
-BASE_MODEL = "unsloth/llama-3-8b-bnb-4bit"
-ADAPTER_MODEL = "kirubel1738/llama3-biology-qa"
+# Base and adapter model IDs
+BASE_MODEL = "microsoft/BioGPT-Large-PubMedQA"
+ADAPTER_MODEL = "kirubel1738/biogpt-pubmedqa-finetuned
 
-# -----------------------------
-# Load model once
-# -----------------------------
 @st.cache_resource
 def load_model():
-    """Load LLaMA-3 8B with PEFT adapter entirely on CPU."""
-    st.info("Loading LLaMA-3 model on CPU... This may take a while.")
+    """Load BioGPT with your Biology-QA adapter."""
+    # Pick device automatically
+    device = 0 if torch.cuda.is_available() else -1
 
-    # Load tokenizer
     tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL)
-
-    # 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
+    base_model = AutoModelForCausalLM.from_pretrained(BASE_MODEL)
+    model = PeftModel.from_pretrained(base_model, ADAPTER_MODEL)  # apply adapter
+    
     generator = pipeline(
         "text-generation",
         model=model,
         tokenizer=tokenizer,
-        device=-1,  # CPU
-        max_new_tokens=256,
-        do_sample=True,
-        temperature=0.7
+        device=device
     )
     return generator
 
@@ -62,9 +42,10 @@ generator = load_model()
 # -----------------------------
 # Streamlit UI
 # -----------------------------
-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.")
+st.set_page_config(page_title="BioGPT — Biology QA Demo", layout="centered")
+st.title("🧬 BioGPT — Pubmed QA Demo")
+
+st.write("Ask a biology-related question and get an answer generated by BioGPT-Large-PubMedQA fine-tuned with your Biology-QA adapter.")
 
 user_input = st.text_area("Enter your biology question:", height=150)
 
@@ -72,7 +53,13 @@ if st.button("Get Answer"):
     if user_input.strip():
         with st.spinner("Generating answer..."):
             try:
-                result = generator(user_input)
+                result = generator(
+                    user_input,
+                    max_new_tokens=128,
+                    do_sample=True,
+                    temperature=0.7,
+                    top_p=0.9
+                )
                 output_text = result[0]["generated_text"]
                 st.success("Answer:")
                 st.write(output_text)
@@ -82,4 +69,4 @@ if st.button("Get Answer"):
         st.warning("Please enter a question.")
 
 st.markdown("---")
-st.caption(f"Model: {BASE_MODEL} + adapter {ADAPTER_MODEL} | Runs on CPU")
+st.caption("Model: microsoft/BioGPT-Large-PubMedQA + adapter kirubel1738/biogpt-biology-qa | Runs on CPU/GPU automatically")