Update app.py

app.py

@@ -1,105 +1,113 @@
 import os
 import streamlit as st
-from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM
+from transformers import (
+    pipeline,
+    AutoTokenizer,
+    AutoModelForCausalLM,
+    TextIteratorStreamer
+)
 from peft import PeftModel
 from supabase import create_client
 from sentence_transformers import SentenceTransformer
+import threading
 
-# — Supabase creds from Secrets —
+# ── 1) Supabase setup ───────────────────────────────────────────────────────
 SUPA_URL = os.getenv("SUPABASE_URL")
 SUPA_KEY = os.getenv("SUPABASE_SERVICE_ROLE_KEY")
 supabase = create_client(SUPA_URL, SUPA_KEY)
 
-# — Embedding model & retrieval function —
-embedder = SentenceTransformer("paraphrase-MiniLM-L3-v2")
+# ── 2) Embedder & memory RPC ────────────────────────────────────────────────
+@st.cache_resource(show_spinner=False)
+def get_embedder():
+    return SentenceTransformer("paraphrase-MiniLM-L3-v2")
+
+embedder = get_embedder()
+
+@st.cache_data(show_spinner=False)
 def fetch_mems(query, k=5):
     vec = embedder.encode(query).tolist()
-    data = supabase.rpc(
-        "match_memories",
-        {"query_embedding": vec, "match_count": k}
-    ).execute().data
-    return data
+    return supabase.rpc("match_memories", {"query_embedding": vec, "match_count": k}).execute().data
 
 def add_mem(speaker, text):
     vec = embedder.encode(text).tolist()
     supabase.table("memories").insert({
-        "speaker": speaker,
-        "text": text,
-        "embedding": vec
+        "speaker": speaker, "text": text, "embedding": vec
     }).execute()
 
-# — Load tokenizer & adapter from HF hub —
-REPO = "sourize/phi2-memory-lora"
-
-# 1) Tokenizer (with your extra PAD token)
-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 Phi-2 → resize embeddings to match tokenizer
-base = AutoModelForCausalLM.from_pretrained(
-    "microsoft/phi-2", trust_remote_code=True, torch_dtype="auto"
-)
-base.resize_token_embeddings(len(tokenizer))
+# ── 3) Model + tokenizer (cached) ───────────────────────────────────────────
+@st.cache_resource(show_spinner=False)
+def load_model():
+    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")
+    model.eval()
+    # prepare a streaming pipeline
+    return tok, pipeline(
+        "text-generation",
+        model=model,
+        tokenizer=tok,
+        device_map="auto",
+        max_new_tokens=64,
+        do_sample=False,
+        use_cache=True,
+        return_full_text=False,
+        streamer=TextIteratorStreamer  # enable streaming
+    )
+
+tokenizer, generator = load_model()
+
+# ── 4) Streamlit UI setup ───────────────────────────────────────────────────
+st.set_page_config(layout="wide")
+st.title("🧠 Memory-Aware Phi-2 Chat")
 
-# 3) Overlay your LoRA adapter
-model = PeftModel.from_pretrained(
-    base,
-    REPO,
-    torch_dtype="auto",
-    device_map="auto"        # let accelerate pick CPU/GPU
-)
-model.eval()
-
-# 4) Build the generation pipeline
-pipe = pipeline(
-    "text-generation",
-    model=model,
-    tokenizer=tokenizer,
-    device=0,                # or device_map="auto"
-    do_sample=True,
-    top_p=0.9,
-    temperature=0.8,
-)
-
-# — Streamlit UI —  
-st.title("🧠 Memory-Aware Phi-2 Bot")
 if "history" not in st.session_state:
-    st.session_state.history = []
-
-def chat(u: str) -> str:
-    # store user turn
-    add_mem("user", u)
-
-    # fetch & format memories
-    mems = fetch_mems(u, 3)
-    block = "\n".join(f"{m['speaker']}: {m['text']}" for m in mems)
-
-    # build prompt
-    prompt = f"""Memory:
-{block}
-
-User: {u}
-Assistant:"""
-
-    # generate reply
-    out = pipe(prompt, max_length=200)[0]["generated_text"]
-    reply = out.split("Assistant:")[-1].strip()
-
-    # store assistant turn
+    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
 
-user = st.text_input("You:")
-if user:
-    resp = chat(user)
-    st.session_state.history.append(("You", user))
-    st.session_state.history.append(("Bot", resp))
-
-for speaker, text in st.session_state.history:
-    if speaker == "You":
-        st.markdown(f"**You:** {text}")
+# ── 6) Render chat bubbles & input ──────────────────────────────────────────
+for role, msg in st.session_state.history:
+    if role == "You":
+        st.chat_message("user").write(msg)
     else:
-        st.markdown(f"**Assistant:** {text}")
+        st.chat_message("assistant").write(msg)
+
+user_input = st.chat_input("Type your message...")
+if user_input:
+    # append placeholder so streamer can fill it
+    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)