Update app.py

app.py

@@ -14,37 +14,83 @@ supabase = create_client(SUPA_URL, SUPA_KEY)
 embedder = SentenceTransformer("paraphrase-MiniLM-L3-v2")
 def fetch_mems(query, k=5):
     vec = embedder.encode(query).tolist()
-    # call your RPC
-    data = supabase.rpc("match_memories", {"query_embedding": vec, "match_count": k}).execute().data
+    data = supabase.rpc(
+        "match_memories",
+        {"query_embedding": vec, "match_count": k}
+    ).execute().data
     return 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 LoRA model from HF hub —
+# — Load tokenizer & adapter from HF hub —
 REPO = "sourize/phi2-memory-lora"
-tokenizer = AutoTokenizer.from_pretrained(REPO, trust_remote_code=True, padding_side="left")
-model_base = AutoModelForCausalLM.from_pretrained(REPO, trust_remote_code=True)
-model_base.resize_token_embeddings(len(tokenizer))
-model = PeftModel.from_pretrained(model_base, REPO)
-pipe  = pipeline("text-generation", model=model, tokenizer=tokenizer, device=0,
-                 do_sample=True, top_p=0.9, temperature=0.8)
 
+# 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) 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):
+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)
-    prompt = f"Memory:\n{block}\n\nUser: {u}\nAssistant:"
-    out    = pipe(prompt, max_length=200)[0]["generated_text"].split("Assistant:")[-1].strip()
-    add_mem("assistant", out)
-    return out
+
+    # 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
+    add_mem("assistant", reply)
+    return reply
 
 user = st.text_input("You:")
 if user:
@@ -52,6 +98,8 @@ if user:
     st.session_state.history.append(("You", user))
     st.session_state.history.append(("Bot", resp))
 
-for s, t in st.session_state.history:
-    style = "### You:" if s=="You" else "**Bot:**"
-    st.markdown(f"{style} {t}")
+for speaker, text in st.session_state.history:
+    if speaker == "You":
+        st.markdown(f"**You:** {text}")
+    else:
+        st.markdown(f"**Assistant:** {text}")