Update app.py

app.py

@@ -2,28 +2,55 @@ import gradio as gr
 from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
 import torch
 
+# Model setup
 model_name = "facebook/blenderbot-400M-distill"
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
 
-history = []
+device = "cuda" if torch.cuda.is_available() else "cpu"
+model.to(device)
+model.eval()
 
-def respond(message):
-    global history
+# Optional speed boost on GPU
+if device == "cuda":
+    model = model.half()
+
+persona = "You are a helpful, concise, friendly assistant."
+
+def respond(message, history):
     history.append({"role": "user", "content": message})
-    # Only use last 3 messages for context
-    last_msgs = history[-3:]
-    input_text = " ".join([m["content"] for m in last_msgs])
-    inputs = tokenizer(input_text, return_tensors="pt")
-    outputs = model.generate(**inputs, max_new_tokens=100)
+
+    # Build context from last 3 turns
+    context = persona + "\n"
+    for msg in history[-6:]:
+        role = "User" if msg["role"] == "user" else "Bot"
+        context += f"{role}: {msg['content']}\n"
+    context += "Bot:"
+
+    inputs = tokenizer(
+        context,
+        return_tensors="pt",
+        truncation=True,
+        max_length=512
+    ).to(device)
+
+    with torch.no_grad():
+        outputs = model.generate(
+            **inputs,
+            max_new_tokens=120,
+            do_sample=True,
+            temperature=0.7,
+            top_p=0.9,
+            repetition_penalty=1.1
+        )
+
     response_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
+
     history.append({"role": "assistant", "content": response_text})
-    return history
+    return history, history
 
 def reset_chat():
-    global history
-    history = []
-    return history
+    return [], []
 
 with gr.Blocks(css="""
 body {background-color: #000 !important; color: #fff !important;}
@@ -35,10 +62,24 @@ body {background-color: #000 !important; color: #fff !important;}
 .gr-button {background-color: #0ff !important; color: #000 !important; border-radius: 8px;}
 footer {display: none !important;}
 """) as demo:
+
+    state = gr.State([])
+
     chatbot = gr.Chatbot(label="DevMegaBlack")
     msg = gr.Textbox(placeholder="Say something...")
     reset_btn = gr.Button("Reset Chat")
-    msg.submit(respond, msg, chatbot)
-    reset_btn.click(reset_chat, [], chatbot)
 
+    msg.submit(
+        respond,
+        [msg, state],
+        [chatbot, state]
+    ).then(lambda: "", None, msg)
+
+    reset_btn.click(
+        reset_chat,
+        [],
+        [chatbot, state]
+    )
+
+demo.queue()
 demo.launch()