Update app.py

app.py

@@ -1,87 +1,56 @@
 import gradio as gr
 from transformers import AutoModelForCausalLM, AutoTokenizer
+import torch
 
-model_name = "microsoft/DialoGPT-small"
+# -----------------------------
+# Model setup
+# -----------------------------
+model_name = "microsoft/DialoGPT-small"  # lightweight, fast conversational model
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 model = AutoModelForCausalLM.from_pretrained(model_name)
 
+# -----------------------------
+# Conversation history
+# -----------------------------
 history = []
 
 def respond(message, chat_history):
-    if chat_history is None:
-        chat_history = []
-
+    """Generate AI response"""
+    # Append user message to history
+    history.append(message)
     
-    chat_history.append((message, None))
-
+    # Keep last 5 messages for context
+    input_text = " ".join(history[-5:])
     
-    input_ids = None
-    for user_msg, bot_msg in chat_history[-5:]:
-        if input_ids is None:
-            input_ids = tokenizer.encode(user_msg + tokenizer.eos_token, return_tensors="pt")
-        else:
-            input_ids = tokenizer.encode(user_msg + tokenizer.eos_token, return_tensors="pt", add_special_tokens=False).to(input_ids.device)
+    # Encode and generate
+    inputs = tokenizer.encode(input_text + tokenizer.eos_token, return_tensors="pt")
+    outputs = model.generate(inputs, max_length=200, pad_token_id=tokenizer.eos_token_id)
+    response = tokenizer.decode(outputs[:, inputs.shape[-1]:][0], skip_special_tokens=True)
     
+    # Update history and chat
+    history.append(response)
+    chat_history.append((message, response))
     
-    outputs = model.generate(input_ids, max_new_tokens=100, pad_token_id=tokenizer.eos_token_id)
-    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
-
+    return chat_history, ""  # clear the input box
+
+# -----------------------------
+# Gradio UI
+# -----------------------------
+with gr.Blocks(css="""
+body {background-color: #000 !important; color: #fff !important;}
+.gr-chatbot {background-color: #111 !important; border-radius: 12px; height: 100% !important;}
+.gr-chatbot .message.user {border-color: #0ff; background-color: transparent !important;}
+.gr-chatbot .message.bot {border-color: #aaa; background-color: transparent !important;}
+.gr-textbox textarea {background-color: transparent !important; color: #fff !important; border: 1px solid #555 !important;}
+.gr-textbox textarea::selection {background-color: #0ff !important; color: #000 !important;}
+.gr-button {display: none !important;}
+footer {display: none !important;}
+""") as demo:
     
-    chat_history[-1] = (message, response)
-
-    return chat_history, chat_history
-
-with gr.Blocks() as demo:
-    chatbot = gr.Chatbot(elem_id="chatbox")
+    chatbot = gr.Chatbot(label="DevMegaBlack AI Chat")
     msg = gr.Textbox(placeholder="Say something...")
-    msg.submit(respond, inputs=[msg, chatbot], outputs=[chatbot, chatbot])
+    
+    msg.submit(respond, [msg, chatbot], [chatbot, msg])
 
-demo.launch(css="""
-body, html, #chatbox, .gradio-container {
-    height: 100% !important;
-    margin: 0; padding: 0;
-    background-color: #000 !important; 
-    color: #fff !important;
-}
-#chatbox {
-    border-radius: 8px;
-    border: 1px solid #fff;
-    background-color: #111 !important;
-    overflow-y: auto;
-}
-.bubble {
-    padding: 10px;
-    border-radius: 12px;
-    margin: 5px;
-    max-width: 80%;
-    word-wrap: break-word;
-    background-color: transparent !important;
-    border: 1px solid #555;
-    color: #fff;
-}
-.user {
-    text-align: right;
-    margin-left: auto;
-    border-color: #0ff;
-}
-.assistant {
-    text-align: left;
-    margin-right: auto;
-    border-color: #aaa;
-}
-.gr-textbox textarea {
-    background-color: transparent !important;
-    color: #fff !important;
-    border: 1px solid #555 !important;
-}
-.gr-textbox textarea::selection {
-    background-color: #0ff !important;
-    color: #000 !important;
-}
-.gr-button {
-    display: none !important;
-}
-footer {
-    display: none !important;
-}
-""")
+# Launch with full height box
+demo.launch()