Update app.py

app.py

@@ -28,8 +28,7 @@ def chat_with_model(message, history, model_choice):
     tokenizer, model = load_model(model_choice)
     device = model.device
     
-    prompt = "You are Creeper AI Chatbot, an intelligent and conversational AI assistant designed to help users with questions, problem-solving, and creative tasks. You communicate clearly, reason carefully, and explain your thoughts in an easy-to-understand way. Stay friendly, professional, and curious. If the user's request is ambiguous, ask clarifying questions before proceeding.\n"
-    prompt += f"User: {message}\nAssistant:"
+    prompt = f"User: {message}\nAssistant:"
     
     streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
     inputs = tokenizer(prompt, return_tensors="pt").to(device)
@@ -37,10 +36,12 @@ def chat_with_model(message, history, model_choice):
     generation_kwargs = dict(
         **inputs,
         streamer=streamer,
-        max_new_tokens=2048,
+        max_new_tokens=512,
         temperature=0.7,
         top_p=0.9,
         do_sample=True,
+        stop_strings=["User:", "Assistant:"],
+        tokenizer=tokenizer
     )
     
     thread = Thread(target=model.generate, kwargs=generation_kwargs)
@@ -48,6 +49,8 @@ def chat_with_model(message, history, model_choice):
     
     partial_text = ""
     for new_text in streamer:
+        if "User:" in new_text:
+            break
         partial_text += new_text
         yield [
             {"role": "user", "content": message},
@@ -55,23 +58,47 @@ def chat_with_model(message, history, model_choice):
         ]
 
 def create_demo():
-    with gr.Blocks(title="Creeper AI Chatbot") as demo:
-        gr.Markdown("## 🤖 Creeper AI Chatbot")
+    custom_theme = gr.themes.Soft(
+        primary_hue="green",
+        secondary_hue="emerald",
+        neutral_hue="slate",
+    ).set(
+        body_background_fill="*neutral_950",
+        block_background_fill="*neutral_900",
+        block_border_width="2px",
+        block_label_text_color="*primary_500",
+        button_primary_background_fill="*primary_600",
+        button_primary_text_color="white",
+    )
+
+    with gr.Blocks(theme=custom_theme, title="Creeper AI Chatbot") as demo:
+        gr.Markdown("# 🌿 Creeper AI Chatbot")
+        gr.Markdown("*Low-latency Liquid Intelligence*")
         
-        model_choice = gr.Dropdown(
-            label="Select Model",
-            choices=list(MODEL_NAMES.keys()),
-            value="LFM 1.2B"
-        )
+        with gr.Row():
+            model_choice = gr.Dropdown(
+                label="Model Core",
+                choices=list(MODEL_NAMES.keys()),
+                value="LFM 1.2B",
+                interactive=True
+            )
         
         chatbot = gr.Chatbot(
-            label="Chat Session",
+            label="Neural Interface",
             type="messages",
-            height=450
+            height=500,
+            show_copy_button=True
         )
         
-        msg = gr.Textbox(label="Your message", placeholder="Type something...")
-        clear = gr.Button("Clear Chat")
+        with gr.Row():
+            msg = gr.Textbox(
+                label="Input Command",
+                placeholder="Message Creeper AI...",
+                scale=4
+            )
+            submit_btn = gr.Button("Send", variant="primary", scale=1)
+        
+        clear = gr.Button("Reset Memory (Wipe)")
         
         def add_user_message(user_message):
             return "", [{"role": "user", "content": user_message}]
@@ -79,6 +106,9 @@ def create_demo():
         msg.submit(add_user_message, [msg], [msg, chatbot], queue=False).then(
             chat_with_model, [msg, chatbot, model_choice], chatbot
         )
+        submit_btn.click(add_user_message, [msg], [msg, chatbot], queue=False).then(
+            chat_with_model, [msg, chatbot, model_choice], chatbot
+        )
         
         clear.click(lambda: [], None, chatbot, queue=False)