Update app.py

app.py

@@ -1,22 +1,16 @@
 import os
 import streamlit as st
-from transformers import (
-    pipeline,
-    AutoTokenizer,
-    AutoModelForCausalLM,
-    TextIteratorStreamer
-)
+from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM
 from peft import PeftModel
 from supabase import create_client
 from sentence_transformers import SentenceTransformer
-import threading
 
-# ── 1) Supabase setup ───────────────────────────────────────────────────────
+# ── Supabase setup ─────────────────────────────────────────────────────────
 SUPA_URL = os.getenv("SUPABASE_URL")
 SUPA_KEY = os.getenv("SUPABASE_SERVICE_ROLE_KEY")
 supabase = create_client(SUPA_URL, SUPA_KEY)
 
-# ── 2) Embedder & memory RPC ────────────────────────────────────────────────
+# ── Embedder & memory RPC ──────────────────────────────────────────────────
 @st.cache_resource(show_spinner=False)
 def get_embedder():
     return SentenceTransformer("paraphrase-MiniLM-L3-v2")
@@ -26,7 +20,9 @@ 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
+    return supabase.rpc("match_memories", 
+                        {"query_embedding": vec, "match_count": k}
+                      ).execute().data
 
 def add_mem(speaker, text):
     vec = embedder.encode(text).tolist()
@@ -34,80 +30,72 @@ def add_mem(speaker, text):
         "speaker": speaker, "text": text, "embedding": vec
     }).execute()
 
-# ── 3) Model + tokenizer (cached) ───────────────────────────────────────────
+# ── Model + tokenizer ──────────────────────────────────────────────────────
 @st.cache_resource(show_spinner=False)
-def load_model():
+def load_generator():
     REPO = "sourize/phi2-memory-lora"
-    # tokenizer
-    tok = AutoTokenizer.from_pretrained(REPO, trust_remote_code=True, padding_side="left")
-    if tok.pad_token_id is None:
-        tok.add_special_tokens({"pad_token": "[PAD]"})
-    # base + resize
-    base = AutoModelForCausalLM.from_pretrained("microsoft/phi-2", trust_remote_code=True)
-    base.resize_token_embeddings(len(tok))
-    # adapter overlay
-    model = PeftModel.from_pretrained(base, REPO, device_map="auto", torch_dtype="auto")
+    # 1) Tokenizer
+    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
+    base = AutoModelForCausalLM.from_pretrained(
+        "microsoft/phi-2", trust_remote_code=True, torch_dtype="auto"
+    )
+    base.resize_token_embeddings(len(tokenizer))
+    # 3) Overlay LoRA adapter
+    model = PeftModel.from_pretrained(
+        base, REPO, device_map="auto", torch_dtype="auto"
+    )
     model.eval()
-    # prepare a streaming pipeline
-    return tok, pipeline(
+    # 4) Pipeline (greedy, small output for speed)
+    gen = pipeline(
         "text-generation",
         model=model,
-        tokenizer=tok,
+        tokenizer=tokenizer,
         device_map="auto",
         max_new_tokens=64,
-        do_sample=False,
+        do_sample=False,         # greedy decoding
         use_cache=True,
         return_full_text=False,
-        streamer=TextIteratorStreamer  # enable streaming
     )
+    return tokenizer, gen
 
-tokenizer, generator = load_model()
+tokenizer, generator = load_generator()
 
-# ── 4) Streamlit UI setup ───────────────────────────────────────────────────
+# ── Streamlit UI ──────────────────────────────────────────────────────────
 st.set_page_config(layout="wide")
 st.title("🧠 Memory-Aware Phi-2 Chat")
 
 if "history" not in st.session_state:
     st.session_state.history = []  # list of (role, message)
 
-# ── 5) Chat function ────────────────────────────────────────────────────────
-def chat(user_input: str):
-    add_mem("user", user_input)
-    # retrieve top-3 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:"
-
-    # stream generation
-    streamer = generator.tokenizer.streamer if hasattr(generator.tokenizer, "streamer") else None
-    if streamer:
-        # If using TextIteratorStreamer, kick off async thread
-        thread = threading.Thread(target=generator, kwargs={"prompt": prompt})
-        thread.start()
-        output = ""
-        for token in streamer:
-            output += token
-            # update the last message in session_state so UI refreshes
-            st.session_state.history[-1] = ("Bot", output)
-            st.experimental_rerun()
-        thread.join()
-    else:
-        output = generator(prompt)[0]["generated_text"]
-    reply = output.strip()
-    add_mem("assistant", reply)
-    return reply
-
-# ── 6) Render chat bubbles & input ──────────────────────────────────────────
+# Render all previous messages as chat bubbles
 for role, msg in st.session_state.history:
     if role == "You":
         st.chat_message("user").write(msg)
     else:
         st.chat_message("assistant").write(msg)
 
+# Input box at the bottom
 user_input = st.chat_input("Type your message...")
+
 if user_input:
-    # append placeholder so streamer can fill it
+    # 1) show user bubble
     st.session_state.history.append(("You", user_input))
-    st.session_state.history.append(("Bot", ""))  
-    # run chat (which will update the last bubble)
-    chat(user_input)
+
+    # 2) store user turn
+    add_mem("user", user_input)
+
+    # 3) retrieve memories and build prompt
+    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
+    with st.spinner("Thinking..."):
+        out = generator(prompt)[0]["generated_text"].strip()
+
+    # 5) show bot bubble and record
+    st.session_state.history.append(("Bot", out))
+    add_mem("assistant", out)