Update app.py

app.py

@@ -1,6 +1,8 @@
 import os
 import streamlit as st
-from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM
+from transformers import (
+    pipeline, AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
+)
 from peft import PeftModel
 from supabase import create_client
 from sentence_transformers import SentenceTransformer
@@ -18,14 +20,15 @@ def get_embedder():
 embedder = get_embedder()
 
 @st.cache_data(show_spinner=False)
-def fetch_mems(query, k=5):
-    vec = embedder.encode(query).tolist()
-    return supabase.rpc("match_memories", 
-                        {"query_embedding": vec, "match_count": k}
-                      ).execute().data
+def fetch_mems(query, k=3):
+    vec = embedder.encode(query).astype('float32').tolist()
+    return supabase.rpc(
+        "match_memories",
+        {"query_embedding": vec, "match_count": k}
+    ).execute().data
 
 def add_mem(speaker, text):
-    vec = embedder.encode(text).tolist()
+    vec = embedder.encode(text).astype('float32').tolist()
     supabase.table("memories").insert({
         "speaker": speaker, "text": text, "embedding": vec
     }).execute()
@@ -35,34 +38,53 @@ def add_mem(speaker, text):
 def load_generator():
     REPO = "sourize/phi2-memory-lora"
     # 1) Tokenizer
-    tokenizer = AutoTokenizer.from_pretrained(REPO, trust_remote_code=True, padding_side="left")
+    tokenizer = AutoTokenizer.from_pretrained(
+        REPO, trust_remote_code=True, padding_side="left"
+    )
     if tokenizer.pad_token_id is None:
         tokenizer.add_special_tokens({"pad_token": "[PAD]"})
-    # 2) Base model & resize
+    # 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, torch_dtype="auto"
+        "microsoft/phi-2",
+        trust_remote_code=True,
+        quantization_config=bnb_config,
+        device_map="auto"
     )
     base.resize_token_embeddings(len(tokenizer))
-    # 3) Overlay LoRA adapter
-    model = PeftModel.from_pretrained(
-        base, REPO, device_map="auto", torch_dtype="auto"
-    )
+    # 4) Overlay LoRA adapter
+    model = PeftModel.from_pretrained(base, REPO, device_map="auto", torch_dtype="auto")
     model.eval()
-    # 4) Pipeline (greedy, small output for speed)
+    # 5) Pipeline with greedy sampling + constraints
     gen = pipeline(
         "text-generation",
         model=model,
         tokenizer=tokenizer,
         device_map="auto",
-        max_new_tokens=64,
-        do_sample=False,         # greedy decoding
+        max_new_tokens=32,
+        do_sample=True,
+        temperature=0.2,
+        top_p=0.8,
         use_cache=True,
-        return_full_text=False,
+        return_full_text=False
     )
     return tokenizer, gen
 
 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"
+)
+
 # ── Streamlit UI ──────────────────────────────────────────────────────────
 st.set_page_config(layout="wide")
 st.title("🧠 Memory-Aware Phi-2 Chat")
@@ -70,7 +92,7 @@ st.title("🧠 Memory-Aware Phi-2 Chat")
 if "history" not in st.session_state:
     st.session_state.history = []  # list of (role, message)
 
-# Render all previous messages as chat bubbles
+# Render existing chat history
 for role, msg in st.session_state.history:
     if role == "You":
         st.chat_message("user").write(msg)
@@ -81,21 +103,31 @@ for role, msg in st.session_state.history:
 user_input = st.chat_input("Type your message...")
 
 if user_input:
-    # 1) show user bubble
+    # Append user message
     st.session_state.history.append(("You", user_input))
-
-    # 2) store user turn
     add_mem("user", user_input)
 
-    # 3) retrieve memories and build prompt
+    # Retrieve relevant memories
     mems = fetch_mems(user_input, k=3)
     mem_block = "\n".join(f"{m['speaker']}: {m['text']}" for m in mems)
-    prompt = f"Memory:\n{mem_block}\n\nUser: {user_input}\nAssistant:"
 
-    # 4) generate reply with spinner
+    # Build prompt
+    prompt = f"""{SYSTEM}
+
+Memory:
+{mem_block}
+
+User: {user_input}
+Assistant:"""
+
+    # Generate reply synchronously with spinner
     with st.spinner("Thinking..."):
-        out = generator(prompt)[0]["generated_text"].strip()
+        try:
+            out = generator(prompt)[0]["generated_text"].strip()
+        except Exception as e:
+            out = "Sorry, I encountered an error."
+            st.error(f"Generation error: {e}")
 
-    # 5) show bot bubble and record
+    # Append assistant reply
     st.session_state.history.append(("Bot", out))
-    add_mem("assistant", out)
+    add_mem("assistant", out)