Update app.py

app.py

@@ -32,102 +32,111 @@ def load_model_and_tokenizer(model_name):
     
     return model_cache[model_name], tokenizer_cache[model_name]
 
-def respond(
-    message,
-    history: list[tuple[str, str]],
-    model_choice,
-    system_message,
-    max_tokens,
-    temperature,
-    top_p,
-):
+# Define available models
+available_models = [
+    "GoofyLM/BrainrotLM-Assistant-362M",
+    "GoofyLM/BrainrotLM2-Assistant-362M"
+]
+
+def respond(message, chat_history, model_choice, system_message, max_tokens, temperature, top_p):
     # Load selected model and tokenizer
-    model_name = model_choice
-    model, tokenizer = load_model_and_tokenizer(model_name)
+    model, tokenizer = load_model_and_tokenizer(model_choice)
     
     # Build conversation messages
     messages = [{"role": "system", "content": system_message}]
-    for user_msg, assistant_msg in history:  
-        if user_msg:  
-            messages.append({"role": "user", "content": user_msg})  
-        if assistant_msg:  
-            messages.append({"role": "assistant", "content": assistant_msg})  
+    for user_msg, assistant_msg in chat_history:
+        messages.append({"role": "user", "content": user_msg})
+        if assistant_msg:  # This might be None during streaming
+            messages.append({"role": "assistant", "content": assistant_msg})
     
-    messages.append({"role": "user", "content": message})  
+    # Add the current message
+    messages.append({"role": "user", "content": message})
     
-    # Format prompt using chat template  
-    prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)  
-    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)  
+    # Format prompt using chat template
+    prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
     
-    # Set up streaming  
-    streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)  
+    # Set up streaming
+    streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
     
-    # Configure generation parameters  
-    do_sample = temperature > 0 or top_p < 1.0  
-    generation_kwargs = dict(  
-        **inputs,  
-        streamer=streamer,  
-        max_new_tokens=max_tokens,  
-        temperature=temperature,  
-        top_p=top_p,  
-        do_sample=do_sample,  
-        pad_token_id=tokenizer.pad_token_id  
-    )  
+    # Configure generation parameters
+    generation_kwargs = dict(
+        **inputs,
+        streamer=streamer,
+        max_new_tokens=max_tokens,
+        temperature=temperature,
+        top_p=top_p,
+        do_sample=(temperature > 0 or top_p < 1.0),
+        pad_token_id=tokenizer.pad_token_id
+    )
     
-    # Start generation in separate thread  
-    thread = Thread(target=model.generate, kwargs=generation_kwargs)  
-    thread.start()  
+    # Start generation in a separate thread
+    thread = Thread(target=model.generate, kwargs=generation_kwargs)
+    thread.start()
     
-    # Stream response  
-    response = ""  
-    for token in streamer:  
-        response += token  
-        yield response
-
-# Define available models
-available_models = [
-    "GoofyLM/BrainrotLM-Assistant-362M",
-    "GoofyLM/BrainrotLM2-Assistant-362M"
-]
+    # Stream the response
+    partial_message = ""
+    for new_token in streamer:
+        partial_message += new_token
+        yield chat_history + [(message, partial_message)]
+        
+    return chat_history + [(message, partial_message)]
 
-# Create Gradio interface
+# Create the Gradio interface
 with gr.Blocks() as demo:
     gr.Markdown("# BrainrotLM Chat Interface")
     
     with gr.Row():
-        with gr.Column(scale=4):
+        with gr.Column(scale=3):
             chatbot = gr.Chatbot(height=600)
-            msg = gr.Textbox(label="Message", placeholder="Type your message here...", lines=3)
-            clear = gr.Button("Clear")
-        
+            
+            with gr.Row():
+                msg = gr.Textbox(
+                    label="Message",
+                    placeholder="Type your message here...",
+                    lines=3,
+                    show_label=False
+                )
+                submit = gr.Button("Send", variant="primary")
+            
+            clear = gr.Button("Clear Conversation")
+            
         with gr.Column(scale=1):
             model_dropdown = gr.Dropdown(
                 choices=available_models,
                 value=available_models[0],
                 label="Select Model"
             )
+            
             system_message = gr.Textbox(
-                value="Your name is BrainrotLM, an AI assistant trained by GoofyLM.", 
+                value="Your name is BrainrotLM, an AI assistant trained by GoofyLM.",
                 label="System message",
                 lines=4
             )
+            
             max_tokens = gr.Slider(1, 512, value=72, label="Max new tokens")
             temperature = gr.Slider(0.1, 2.0, value=0.65, label="Temperature")
             top_p = gr.Slider(0.1, 1.0, value=0.95, label="Top-p (nucleus sampling)")
     
-    chat_interface = gr.ChatInterface(
-        fn=respond,
-        chatbot=chatbot,
-        textbox=msg,
-        clear_btn=clear,
-        additional_inputs=[
-            model_dropdown,
-            system_message,
-            max_tokens,
-            temperature,
-            top_p,
-        ],
+    # Set up event handlers
+    submit_event = msg.submit(
+        respond,
+        inputs=[msg, chatbot, model_dropdown, system_message, max_tokens, temperature, top_p],
+        outputs=chatbot
+    )
+    
+    submit_click = submit.click(
+        respond,
+        inputs=[msg, chatbot, model_dropdown, system_message, max_tokens, temperature, top_p],
+        outputs=chatbot
     )
+    
+    # Clear message box after sending
+    submit_event.then(lambda: "", None, msg)
+    submit_click.then(lambda: "", None, msg)
+    
+    # Clear conversation button
+    clear.click(lambda: None, None, chatbot)
 
 if __name__ == "__main__":
     demo.launch()