Update app.py

app.py

@@ -1,7 +1,8 @@
 import os
 import threading
 from collections import defaultdict
-
+from PIL import Image
+import tempfile
 import gradio as gr
 from transformers import (
     AutoModelForCausalLM,
@@ -9,17 +10,14 @@ from transformers import (
     TextIteratorStreamer,
 )
 
-# Define model paths
 model_name_to_path = {
     "LeCarnet-3M": "MaxLSB/LeCarnet-3M",
     "LeCarnet-8M": "MaxLSB/LeCarnet-8M",
     "LeCarnet-21M": "MaxLSB/LeCarnet-21M",
 }
 
-# Load Hugging Face token
 hf_token = os.environ["HUGGINGFACEHUB_API_TOKEN"]
 
-# Preload models and tokenizers
 loaded_models = defaultdict(dict)
 
 for name, path in model_name_to_path.items():
@@ -27,40 +25,24 @@ 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 resize_logo(input_path, size=(100, 100)):
+    with Image.open(input_path) as img:
+        img = img.resize(size, Image.LANCZOS)
+        temp_file = tempfile.NamedTemporaryFile(suffix=".png", delete=False)
+        img.save(temp_file.name, format="PNG")
+        return temp_file.name
+
 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
+    yield history
     tokenizer = loaded_models[model_name]["tokenizer"]
     model = loaded_models[model_name]["model"]
-
-    # Tokenize input
     inputs = tokenizer(message, return_tensors="pt")
-
-    # Set up streaming
     streamer = TextIteratorStreamer(
         tokenizer,
         skip_prompt=False,
         skip_special_tokens=True,
     )
-
-    # Configure generation parameters
     generate_kwargs = dict(
         **inputs,
         streamer=streamer,
@@ -70,115 +52,61 @@ def respond(message, history, model_name, max_tokens, temperature, top_p):
         top_p=top_p,
         eos_token_id=tokenizer.eos_token_id,
     )
-
-    # Start generation in a background thread
     thread = threading.Thread(target=model.generate, kwargs=generate_kwargs)
     thread.start()
-
-    # Stream the response with model name prefix
-    accumulated = f"**{model_name}:** "
+    accumulated = f"**{model_name}**\n\n"
     for new_text in streamer:
         accumulated += new_text
         history[-1] = (message, accumulated)
         yield history
 
 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, ""
 
-def select_model(model_name, current_model):
-    """
-    Update the selected model name when a model button is clicked.
-    
-    Args:
-        model_name (str): The model name to select.
-        current_model (str): The currently selected model.
-    
-    Returns:
-        str: The newly selected model name.
-    """
-    return model_name
+def start_with_example(example, model_name, max_tokens, temperature, top_p):
+    for updated_history in respond(example, [], model_name, max_tokens, temperature, top_p):
+        yield updated_history, ""
+
+resized_logo_path = resize_logo("media/le-carnet.png", size=(100, 100))
+
+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",
+]
 
-# Create the Gradio interface with Blocks
 with gr.Blocks(css=".gr-button {margin: 5px; width: 100%;} .gr-column {padding: 10px;}") as demo:
-    # Title and description
     gr.Markdown("# LeCarnet")
-    gr.Markdown("Select a model on the right and type a message to chat.")
-
-    # Two-column layout with specific widths
+    gr.Markdown("Select a model on the right and type a message to chat, or choose an example below.")
     with gr.Row():
-        # Left column: Chat interface (80% width)
         with gr.Column(scale=4):
+            dataset = gr.Dataset(components=[gr.Textbox(visible=False)], samples=[[ex] for ex in examples], type="values")
             chatbot = gr.Chatbot(
-                avatar_images=(None, "media/le-carnet.png"),  # User avatar: None, Bot avatar: Logo
+                avatar_images=(None, resized_logo_path),
                 label="Chat",
-                height=600,  # Increase chat height for larger display
+                height=600,
             )
             user_input = gr.Textbox(placeholder="Type your message here...", label="Message")
             submit_btn = gr.Button("Send")
-
-        # Right column: Model selection and parameters (20% width)
         with gr.Column(scale=1, min_width=200):
-            # State to track selected model
-            model_state = gr.State(value="LeCarnet-8M")
-
-            # Model selection buttons
-            gr.Markdown("**Select Model**")
-            btn_3m = gr.Button("LeCarnet-3M")
-            btn_8m = gr.Button("LeCarnet-8M")
-            btn_21m = gr.Button("LeCarnet-21M")
-
-            # Sliders for parameters
+            model_dropdown = gr.Dropdown(
+                choices=list(model_name_to_path.keys()),
+                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"],
-        ],
-        inputs=user_input,
-    )
-
-    # Event handling for submit button
     submit_btn.click(
         fn=submit,
-        inputs=[user_input, chatbot, model_state, max_tokens, temperature, top_p],
+        inputs=[user_input, chatbot, model_dropdown, max_tokens, temperature, top_p],
         outputs=[chatbot, user_input],
     )
-
-    # Event handling for model selection buttons
-    btn_3m.click(
-        fn=select_model,
-        inputs=[gr.State("LeCarnet-3M"), model_state],
-        outputs=model_state,
-    )
-    btn_8m.click(
-        fn=select_model,
-        inputs=[gr.State("LeCarnet-8M"), model_state],
-        outputs=model_state,
-    )
-    btn_21m.click(
-        fn=select_model,
-        inputs=[gr.State("LeCarnet-21M"), model_state],
-        outputs=model_state,
+    dataset.change(
+        fn=start_with_example,
+        inputs=[dataset, model_dropdown, max_tokens, temperature, top_p],
+        outputs=[chatbot, user_input],
     )
 
 if __name__ == "__main__":