Update app.py

app.py

@@ -2,75 +2,70 @@ import os
 import streamlit as st
 from huggingface_hub import InferenceClient
 
-# ── Configuration ──────────────────────────────────────────────────────────
-HF_TOKEN       = os.getenv("HF_TOKEN")  # set this in your Space Secrets
-MODEL_ID       = "sourize/phi2-memory-lora"
-CONTEXT_TURNS  = 6
+# ── Config ─────────────────────────────────────────────────────────
+MODEL_ID = "sourize/phi2-memory-lora"
+HF_TOKEN = os.getenv("HF_TOKEN")  # Set this in your Space secrets
+CONTEXT_TURNS = 6
 MAX_NEW_TOKENS = 128
+TEMPERATURE = 0.7
+TOP_P = 0.9
 
 SYSTEM_PROMPT = (
-    "You are a helpful assistant for DeepTalks, based on Phi-2 fine-tuned by Sourish.\n"
+    "You are a helpful assistant for DeepTalks with base Phi-2\n"
+    "fine-tuned by Sourish for domain support.\n"
     "Answer **only** using the conversation context below.\n"
-    "Do NOT output lines beginning with 'User:' or 'Assistant:'.\n"
+    "Do NOT output any lines beginning with 'User:' or 'Assistant:'.\n"
     "If you don't know, say \"I don't know.\"\n"
 )
 
-# ── Initialize Hugging Face Inference Client ─────────────────────────────────
-client = InferenceClient(token=HF_TOKEN)
+client = InferenceClient(model=MODEL_ID, token=HF_TOKEN)
 
-def query_hf(prompt: str) -> str:
-    # model_id and prompt are positional args, params passed via params=
-    response = client.text_generation(
-        MODEL_ID,
-        prompt,
-        params={
-            "max_new_tokens": MAX_NEW_TOKENS,
-            "do_sample": True,
-            "temperature": 0.7,
-            "top_p": 0.9,
-            "return_full_text": False,
-        },
-    )
-    text = response.generated_text.strip()
-    # Remove stray markers
-    for marker in ("User:", "Assistant:"):
-        if marker in text:
-            text = text.split(marker)[0].strip()
-    return text or "I don't know."
-
-# ── Streamlit UI ──────────────────────────────────────────────────────────
+# ── UI Setup ────────────────────────────────────────────────────────
 st.set_page_config(layout="centered")
-st.title("🧠 DeepTalks (Inference API)")
-st.subheader("Your personal AI Companion")
+st.title("🧠 DeepTalks")
+st.subheader("Your personal AI Companion", divider="grey")
 
 if "history" not in st.session_state:
-    st.session_state.history = []  # list of (role, text)
+    st.session_state.history = []
 
-# Render chat history
-for role, text in st.session_state.history:
-    st.chat_message("user" if role == "You" else "assistant").write(text)
+# ── Show Previous Messages ─────────────────────────────────────────
+for role, message in st.session_state.history:
+    st.chat_message("user" if role == "You" else "assistant").write(message)
 
-# User input
-user_input = st.chat_input("Type your message…")
+# ── New User Input ─────────────────────────────────────────────────
+user_input = st.chat_input("Your message…")
 if user_input:
     st.chat_message("user").write(user_input)
     st.session_state.history.append(("You", user_input))
 
-    # Build context from last turns
+    # Prepare context
     recent = st.session_state.history[-CONTEXT_TURNS*2:]
-    context = "\n".join(msg for _, msg in recent)
-    prompt = (
-        f"{SYSTEM_PROMPT}\n"
-        f"Context:\n{context}\n\n"
-        f"User: {user_input}\nAssistant:"
-    )
+    context = "\n".join([f"{role}: {text}" for role, text in recent])
+    prompt = f"{SYSTEM_PROMPT}\n\nContext:\n{context}\n\nUser: {user_input}\nAssistant:"
 
     with st.spinner("Thinking…"):
         try:
-            answer = query_hf(prompt)
+            response = client.text_generation(
+                prompt,
+                max_new_tokens=MAX_NEW_TOKENS,
+                temperature=TEMPERATURE,
+                top_p=TOP_P,
+                do_sample=True,
+                return_full_text=False,
+            )
+
+            # Clean up
+            reply = response.strip()
+            for marker in ["User:", "Assistant:"]:
+                if marker in reply:
+                    reply = reply.split(marker)[0].strip()
+
+            if not reply:
+                reply = "I’m sorry, I didn’t catch that. Could you rephrase?"
+
         except Exception as e:
+            reply = "I’m sorry, something went wrong."
             st.error(f"Inference API error: {e}")
-            answer = "I’m sorry, something went wrong."
 
-    st.chat_message("assistant").write(answer)
-    st.session_state.history.append(("Bot", answer))
+    st.chat_message("assistant").write(reply)
+    st.session_state.history.append(("Bot", reply))