Update app.py

app.py

@@ -1,12 +1,19 @@
 import os
 import streamlit as st
 from transformers import (
-    pipeline, AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
+    pipeline, AutoTokenizer, AutoModelForCausalLM
 )
 from peft import PeftModel
 from supabase import create_client
 from sentence_transformers import SentenceTransformer
 
+# Try to import bitsandbytes for 4-bit; fall back if missing
+try:
+    from transformers import BitsAndBytesConfig
+    BNB_AVAILABLE = True
+except ImportError:
+    BNB_AVAILABLE = False
+
 # ── Supabase setup ─────────────────────────────────────────────────────────
 SUPA_URL = os.getenv("SUPABASE_URL")
 SUPA_KEY = os.getenv("SUPABASE_SERVICE_ROLE_KEY")
@@ -43,25 +50,37 @@ def load_generator():
     )
     if tokenizer.pad_token_id is None:
         tokenizer.add_special_tokens({"pad_token": "[PAD]"})
-    # 2) Quantization config for 4-bit
-    bnb_config = BitsAndBytesConfig(
-        load_in_4bit=True,
-        bnb_4bit_quant_type="nf4",
-        bnb_4bit_compute_dtype="float16",
-        low_cpu_mem_usage=True,
-    )
-    # 3) Load base model in 4-bit + resize embeddings
-    base = AutoModelForCausalLM.from_pretrained(
-        "microsoft/phi-2",
-        trust_remote_code=True,
-        quantization_config=bnb_config,
-        device_map="auto"
-    )
+
+    # 2) Base model load (with or without 4-bit)
+    if BNB_AVAILABLE:
+        bnb_config = BitsAndBytesConfig(
+            load_in_4bit=True,
+            bnb_4bit_quant_type="nf4",
+            bnb_4bit_compute_dtype="float16",
+            low_cpu_mem_usage=True,
+        )
+        base = AutoModelForCausalLM.from_pretrained(
+            "microsoft/phi-2",
+            trust_remote_code=True,
+            quantization_config=bnb_config,
+            device_map="auto"
+        )
+    else:
+        base = AutoModelForCausalLM.from_pretrained(
+            "microsoft/phi-2",
+            trust_remote_code=True,
+            torch_dtype="auto",
+            device_map="auto"
+        )
+
+    # 3) Resize embeddings & overlay LoRA
     base.resize_token_embeddings(len(tokenizer))
-    # 4) Overlay LoRA adapter
-    model = PeftModel.from_pretrained(base, REPO, device_map="auto", torch_dtype="auto")
+    model = PeftModel.from_pretrained(
+        base, REPO, device_map="auto", torch_dtype="auto"
+    )
     model.eval()
-    # 5) Pipeline with greedy sampling + constraints
+
+    # 4) Pipeline (greedy-ish sampling)
     gen = pipeline(
         "text-generation",
         model=model,
@@ -80,9 +99,9 @@ tokenizer, generator = load_generator()
 
 # ── System prompt to reduce hallucinations ──────────────────────────────────
 SYSTEM = (
-    "You are a helpful assistant.\\n"
-    "Answer **only** using the information in the memory below.\\n"
-    "If the answer is not in memory, reply: \"I don't know.\"\\n"
+    "You are a helpful assistant.\n"
+    "Answer **only** using the information in the memory below.\n"
+    "If the answer is not in memory, reply: \"I don't know.\"\n"
 )
 
 # ── Streamlit UI ──────────────────────────────────────────────────────────
@@ -94,10 +113,7 @@ if "history" not in st.session_state:
 
 # Render existing chat history
 for role, msg in st.session_state.history:
-    if role == "You":
-        st.chat_message("user").write(msg)
-    else:
-        st.chat_message("assistant").write(msg)
+    st.chat_message("user" if role=="You" else "assistant").write(msg)
 
 # Input box at the bottom
 user_input = st.chat_input("Type your message...")
@@ -120,7 +136,7 @@ Memory:
 User: {user_input}
 Assistant:"""
 
-    # Generate reply synchronously with spinner
+    # Generate reply with spinner
     with st.spinner("Thinking..."):
         try:
             out = generator(prompt)[0]["generated_text"].strip()
@@ -130,4 +146,4 @@ Assistant:"""
 
     # Append assistant reply
     st.session_state.history.append(("Bot", out))
-    add_mem("assistant", out)
+    add_mem("assistant", out)