Update app.py

app.py

@@ -1,70 +1,126 @@
 import gradio as gr
-import requests
+from transformers import AutoTokenizer, AutoModelForCausalLM
+import torch
 
-# --------- Configuration ---------
-HF_API_URL = "https://api-inference.huggingface.co/models/driaforall/mem-agent"
-HF_TOKEN = "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"
+# --------- Load model ---------
+tokenizer = AutoTokenizer.from_pretrained("driaforall/mem-agent", trust_remote_code=True)
+model = AutoModelForCausalLM.from_pretrained("driaforall/mem-agent", trust_remote_code=True)
+model.eval()
 
-USER_AVATAR = "https://i.imgur.com/YourUserAvatar.png"
-AI_AVATAR = "https://i.imgur.com/YourAIAvatar.png"
+# --------- Chat function ---------
+def chat_fn(history, message):
+    if not message.strip():
+        return history, gr.update(value="")  # ignore empty messages
 
-GOOGLE_VERIFICATION = "OXT03M7HRqpImYW2myqTjRmEP_f8wZi2SnbrO0JmGcc"
+    messages = [{"role": "user", "content": message}]
+    inputs = tokenizer.apply_chat_template(
+        messages,
+        add_generation_prompt=True,
+        tokenize=True,
+        return_dict=True,
+        return_tensors="pt"
+    ).to(model.device)
 
-# --------- AI call function ---------
-def call_ai(message):
-    headers = {
-        "Authorization": f"Bearer {HF_TOKEN}",
-        "Content-Type": "application/json"
-    }
-    payload = {"inputs": message}
-    try:
-        response = requests.post(HF_API_URL, headers=headers, json=payload)
-        data = response.json()
-        reply = "⚠️ No response from AI"
-        if isinstance(data, list) and len(data) > 0:
-            if "generated_text" in data[0]:
-                reply = data[0]["generated_text"]
-            elif "generated_texts" in data[0] and len(data[0]["generated_texts"]) > 0:
-                reply = data[0]["generated_texts"][0]
-            elif isinstance(data[0], str):
-                reply = data[0]
-        elif isinstance(data, str):
-            reply = data
-        return reply
-    except Exception as e:
-        return f"⚠️ Error: {e}"
+    outputs = model.generate(**inputs, max_new_tokens=60, do_sample=True, temperature=0.8)
+    reply = tokenizer.decode(outputs[0][inputs["input_ids"].shape[-1]:], skip_special_tokens=True)
 
-# --------- UI ---------
-with gr.Blocks(title="Mem-Agent AI Chat", theme="default") as demo:
-    # Inject Google verification meta tag
-    gr.HTML(f"""
+    history.append((message, reply))
+    return history, gr.update(value="")  # reset input box
+
+# --------- CSS for chat bubbles ---------
+css = """
+body {
+  background: linear-gradient(135deg, #141e30, #243b55);
+  font-family: 'Poppins', sans-serif;
+  color: white;
+}
+
+#chatbox {
+  border-radius: 20px;
+  padding: 10px;
+  background: rgba(0,0,0,0.3);
+  box-shadow: 0 0 20px rgba(0,0,0,0.5);
+  max-height: 500px;
+  overflow-y: auto;
+}
+
+.user-msg {
+  display: flex;
+  justify-content: flex-end;
+  margin: 8px 0;
+}
+.bot-msg {
+  display: flex;
+  justify-content: flex-start;
+  margin: 8px 0;
+}
+.avatar {
+  width: 40px;
+  height: 40px;
+  border-radius: 50%;
+  margin: 0 8px;
+}
+.user-bubble {
+  background: #2f88ff;
+  color: white;
+  border-radius: 18px 18px 4px 18px;
+  padding: 10px 14px;
+  max-width: 60%;
+}
+.bot-bubble {
+  background: #3a3a3a;
+  color: #f5f5f5;
+  border-radius: 18px 18px 18px 4px;
+  padding: 10px 14px;
+  max-width: 60%;
+}
+"""
+
+# --------- Render chat history ---------
+def render_chat(history):
+    chat_html = ""
+    for user_msg, bot_msg in history:
+        chat_html += f"""
+        <div class="user-msg">
+          <div class="user-bubble">{user_msg}</div>
+          <img class="avatar" src="https://i.imgur.com/4M34hi2.png" alt="User">
+        </div>
+        <div class="bot-msg">
+          <img class="avatar" src="https://i.imgur.com/8Km9tLL.png" alt="AI">
+          <div class="bot-bubble">{bot_msg}</div>
+        </div>
+        """
+    return chat_html
+
+# --------- Gradio UI ---------
+with gr.Blocks(css=css, theme=gr.themes.Soft()) as demo:
+    # Inject Google Site Verification meta tag
+    gr.HTML("""
     <head>
-        <meta name="google-site-verification" content="{GOOGLE_VERIFICATION}" />
+        <meta name="google-site-verification" content="OXT03M7HRqpImYW2myqTjRmEP_f8wZi2SnbrO0JmGcc" />
     </head>
     """)
 
-    gr.Markdown("## 💬 Mem-Agent AI Chat")
-
     with gr.Column():
-        # Updated for modern Gradio
-        chat_history = gr.Chatbot(elem_id="chatbot", type="messages")  
-
+        gr.HTML("<h1 style='text-align:center; color:#00fff7;'>💬 AI Chat (Discord Style)</h1>")
+        chatbox = gr.HTML("", elem_id="chatbox")
         with gr.Row():
-            user_input = gr.Textbox(
-                placeholder="Type your message here...",
-                label="Your Message"
+            msg = gr.Textbox(
+                placeholder="Type your message...",
+                show_label=False,
+                container=False,
             )
-            submit_btn = gr.Button("Send")
+            send = gr.Button("Send 🚀")
+        clear = gr.Button("Clear Chat 🧹")
+
+    state = gr.State([])
 
-    # Handle sending messages
-    def respond(message, chat):
-        ai_reply = call_ai(message)
-        chat.append({"role": "user", "content": f"{message}"})
-        chat.append({"role": "assistant", "content": f"{ai_reply}"})
-        return "", chat
+    def process_message(history, message):
+        history, _ = chat_fn(history, message)
+        return history, render_chat(history), ""
 
-    submit_btn.click(respond, [user_input, chat_history], [user_input, chat_history])
-    user_input.submit(respond, [user_input, chat_history], [user_input, chat_history])
+    send.click(process_message, [state, msg], [state, chatbox, msg])
+    msg.submit(process_message, [state, msg], [state, chatbox, msg])
+    clear.click(lambda: ([], "", ""), None, [state, chatbox, msg])
 
-# --------- Launch ---------
 demo.launch()