Update app.py

app.py

@@ -1,14 +1,14 @@
 import os
 import streamlit as st
-from huggingface_hub import InferenceClient
+from huggingface_hub import InferenceApi
 
-# ── Config ─────────────────────────────────────────────────────────
-MODEL_ID = "sourize/phi2-memory-lora"
-HF_TOKEN = os.getenv("HF_TOKEN")  # Set this in your Space secrets
-CONTEXT_TURNS = 6
+# ── Configuration ──────────────────────────────────────────────────────────
+MODEL_ID       = "sourize/phi2-memory-lora"
+HF_TOKEN       = os.getenv("HF_TOKEN")   # set in your Space secrets
+CONTEXT_TURNS  = 6
 MAX_NEW_TOKENS = 128
-TEMPERATURE = 0.7
-TOP_P = 0.9
+TEMPERATURE    = 0.7
+TOP_P          = 0.9
 
 SYSTEM_PROMPT = (
     "You are a helpful assistant for DeepTalks with base Phi-2\n"
@@ -18,54 +18,73 @@ SYSTEM_PROMPT = (
     "If you don't know, say \"I don't know.\"\n"
 )
 
-client = InferenceClient(model=MODEL_ID, token=HF_TOKEN)
+# ── Initialize the low-level Inference API ─────────────────────────────────
+inference = InferenceApi(
+    repo_id=MODEL_ID,
+    token=HF_TOKEN,
+    task="text-generation"
+)
+
+def query_hf(prompt: str) -> str:
+    """
+    Call the text-generation endpoint via InferenceApi.
+    """
+    outputs = inference(
+        inputs=prompt,
+        parameters={
+            "max_new_tokens": MAX_NEW_TOKENS,
+            "temperature": TEMPERATURE,
+            "top_p": TOP_P,
+            "do_sample": True,
+            "return_full_text": False
+        }
+    )
+    # outputs is a list of dicts: [{"generated_text": "..."}]
+    text = outputs[0]["generated_text"].strip()
+    # strip any stray markers
+    for marker in ("User:", "Assistant:"):
+        if marker in text:
+            text = text.split(marker)[0].strip()
+    return text or "I don't know."
 
-# ── UI Setup ────────────────────────────────────────────────────────
+# ── Streamlit UI ──────────────────────────────────────────────────────────
 st.set_page_config(layout="centered")
-st.title("🧠 DeepTalks")
-st.subheader("Your personal AI Companion", divider="grey")
+st.title("🧠 DeepTalks (Inference API)")
+st.subheader("Your personal AI Companion")
 
 if "history" not in st.session_state:
-    st.session_state.history = []
+    st.session_state.history = []  # list of (role, text)
 
-# ── Show Previous Messages ─────────────────────────────────────────
-for role, message in st.session_state.history:
-    st.chat_message("user" if role == "You" else "assistant").write(message)
+# Render previous messages
+for role, msg in st.session_state.history:
+    st.chat_message("user" if role=="You" else "assistant").write(msg)
 
-# ── New User Input ─────────────────────────────────────────────────
-user_input = st.chat_input("Your message…")
+# New user input
+user_input = st.chat_input("Type your message…")
 if user_input:
+    # Display and record user
     st.chat_message("user").write(user_input)
     st.session_state.history.append(("You", user_input))
 
-    # Prepare context
+    # Build context of last 2*CONTEXT_TURNS entries
     recent = st.session_state.history[-CONTEXT_TURNS*2:]
-    context = "\n".join([f"{role}: {text}" for role, text in recent])
-    prompt = f"{SYSTEM_PROMPT}\n\nContext:\n{context}\n\nUser: {user_input}\nAssistant:"
+    context = "\n".join(text for _, text in recent)
 
+    prompt = (
+        f"{SYSTEM_PROMPT}\n\n"
+        f"Context:\n{context}\n\n"
+        f"User: {user_input}\n"
+        f"Assistant:"
+    )
+
+    # Call Hugging Face Inference API
     with st.spinner("Thinking…"):
         try:
-            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?"
-
+            reply = query_hf(prompt)
         except Exception as e:
-            reply = "I’m sorry, something went wrong."
             st.error(f"Inference API error: {e}")
+            reply = "I’m sorry, something went wrong."
 
+    # Display and record assistant
     st.chat_message("assistant").write(reply)
     st.session_state.history.append(("Bot", reply))