Update app.py

app.py

@@ -27,20 +27,31 @@ for name, path in model_name_to_path.items():
     loaded_models[name]["model"] = AutoModelForCausalLM.from_pretrained(path, token=hf_token)
     loaded_models[name]["model"].eval()
 
-def respond(
-    prompt: str,
-    chat_history,
-    model_name: str,
-    max_tokens: int,
-    temperature: float,
-    top_p: float,
-):
-    # Select the appropriate model and tokenizer
+def respond(message, history, model_name, max_tokens, temperature, top_p):
+    """
+    Generate a response from the selected model, streaming the output and updating chat history.
+    
+    Args:
+        message (str): User's input message.
+        history (list): Current chat history as list of (user_msg, bot_msg) tuples.
+        model_name (str): Selected model name.
+        max_tokens (int): Maximum number of tokens to generate.
+        temperature (float): Sampling temperature.
+        top_p (float): Top-p sampling parameter.
+    
+    Yields:
+        list: Updated chat history with the user's message and streaming bot response.
+    """
+    # Append user's message to history with an empty bot response
+    history = history + [(message, "")]
+    yield history  # Display user's message immediately
+
+    # Select tokenizer and model
     tokenizer = loaded_models[model_name]["tokenizer"]
     model = loaded_models[model_name]["model"]
 
     # Tokenize input
-    inputs = tokenizer(prompt, return_tensors="pt")
+    inputs = tokenizer(message, return_tensors="pt")
 
     # Set up streaming
     streamer = TextIteratorStreamer(
@@ -60,65 +71,79 @@ def respond(
         eos_token_id=tokenizer.eos_token_id,
     )
 
-    # Run generation in a background thread
+    # Start generation in a background thread
     thread = threading.Thread(target=model.generate, kwargs=generate_kwargs)
     thread.start()
 
-    # Stream results
-    accumulated = ""
+    # Stream the response with model name prefix
+    accumulated = f"**{model_name}:** "
     for new_text in streamer:
         accumulated += new_text
-        yield accumulated
+        history[-1] = (message, accumulated)
+        yield history
 
-# Create Gradio Chat Interface
+def submit(message, history, model_name, max_tokens, temperature, top_p):
+    """
+    Handle form submission by calling respond and clearing the input box.
+    
+    Args:
+        message (str): User's input message.
+        history (list): Current chat history.
+        model_name (str): Selected model name.
+        max_tokens (int): Max tokens parameter.
+        temperature (float): Temperature parameter.
+        top_p (float): Top-p parameter.
+    
+    Yields:
+        tuple: (updated chat history, cleared user input)
+    """
+    for updated_history in respond(message, history, model_name, max_tokens, temperature, top_p):
+        yield updated_history, ""
+
+# Create the Gradio interface with Blocks
 with gr.Blocks() as demo:
-    # Custom title with logo
+    # Title and description
+    gr.Markdown("# LeCarnet")
+    gr.Markdown("Select a model on the right and type a message to chat.")
+
+    # Two-column layout
     with gr.Row():
-        gr.HTML(
-            '<div style="display: flex; align-items: center;">'
-            f'<img src="file/{os.path.abspath("media/le-carnet.png")}" style="height: 50px; margin-right: 10px;" />'
-            '<h1 style="margin: 0;">LeCarnet</h1>'
-            '</div>'
-        )
-    
-    # Chat interface
-    chatbot = gr.ChatInterface(
-        fn=respond,
-        title=None,  # Remove default title
-        description=None,  # Remove default description
+        # Left column: Chat interface
+        with gr.Column():
+            chatbot = gr.Chatbot(
+                avatar_images=(None, "media/le-carnet.png"),  # User avatar: None, Bot avatar: Logo
+                label="Chat"
+            )
+            user_input = gr.Textbox(placeholder="Type your message here...", label="Message")
+            submit_btn = gr.Button("Send")
+
+        # Right column: Model selection and parameters
+        with gr.Column():
+            model_dropdown = gr.Dropdown(
+                choices=["LeCarnet-3M", "LeCarnet-8M", "LeCarnet-21M"],
+                value="LeCarnet-8M",
+                label="Model"
+            )
+            max_tokens = gr.Slider(1, 512, value=512, step=1, label="Max New Tokens")
+            temperature = gr.Slider(0.1, 2.0, value=0.7, step=0.1, label="Temperature")
+            top_p = gr.Slider(0.1, 1.0, value=0.9, step=0.05, label="Top-p")
+
+    # Example prompts
+    examples = gr.Examples(
         examples=[
             ["Il était une fois un petit garçon qui vivait dans un village paisible."],
             ["Il était une fois une grenouille qui rêvait de toucher les étoiles chaque nuit depuis son étang."],
             ["Il était une fois un petit lapin perdu"],
         ],
-        cache_examples=False,
+        inputs=user_input,
     )
-    
-    # Sidebar for model selection and parameters
-    with gr.Column(elem_classes="sidebar", variant="panel"):
-        gr.Markdown("### Model Configuration")
-        model_dropdown = gr.Dropdown(
-            choices=["LeCarnet-3M", "LeCarnet-8M", "LeCarnet-21M"],
-            value="LeCarnet-8M",
-            label="Model",
-        )
-        max_tokens_slider = gr.Slider(1, 512, value=512, step=1, label="Max New Tokens")
-        temperature_slider = gr.Slider(0.1, 2.0, value=0.7, step=0.1, label="Temperature")
-        top_p_slider = gr.Slider(0.1, 1.0, value=0.9, step=0.05, label="Top-p")
-        
-        # Pass parameters to the chatbot
-        chatbot.load(
-            fn=lambda x, y, z, w: None,
-            inputs=[model_dropdown, max_tokens_slider, temperature_slider, top_p_slider],
-            outputs=None,
-        )
-        chatbot.config.update({
-            "model_name": model_dropdown,
-            "max_tokens": max_tokens_slider,
-            "temperature": temperature_slider,
-            "top_p": top_p_slider,
-        })
-
-# Launch the app
+
+    # Event handling for submit button
+    submit_btn.click(
+        fn=submit,
+        inputs=[user_input, chatbot, model_dropdown, max_tokens, temperature, top_p],
+        outputs=[chatbot, user_input],
+    )
+
 if __name__ == "__main__":
     demo.queue(default_concurrency_limit=10, max_size=10).launch(ssr_mode=False, max_threads=10)