Update app.py

app.py

@@ -1,49 +1,34 @@
 import os
 import threading
-from collections import defaultdict
-
 import gradio as gr
-from transformers import (
-    AutoModelForCausalLM,
-    AutoTokenizer,
-    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.get("HUGGINGFACEHUB_API_TOKEN", "default_token")  # Use default to avoid errors
-
-# Preload models and tokenizers
-loaded_models = defaultdict(dict)
-
-for name, path in model_name_to_path.items():
-    try:
-        loaded_models[name]["tokenizer"] = AutoTokenizer.from_pretrained(path, token=hf_token)
-        loaded_models[name]["model"] = AutoModelForCausalLM.from_pretrained(path, token=hf_token)
-        loaded_models[name]["model"].eval()
-    except Exception as e:
-        print(f"Error loading {name}: {str(e)}")
-
-def respond(message, history, model_name, max_tokens, temperature, top_p):
-    history = history + [(message, "")]
-    yield history
-
-    tokenizer = loaded_models[model_name]["tokenizer"]
-    model = loaded_models[model_name]["model"]
+from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
 
-    inputs = tokenizer(message, return_tensors="pt")
+# Hugging Face token
+hf_token = os.environ["HUGGINGFACEHUB_API_TOKEN"]
 
-    streamer = TextIteratorStreamer(
-        tokenizer,
-        skip_prompt=False,
-        skip_special_tokens=True,
-    )
+# Global model & tokenizer
+tokenizer = None
+model = None
+
+# Load selected model
+def load_model(model_name):
+    global tokenizer, model
+    full_model_name = f"MaxLSB/{model_name}"
+    tokenizer = AutoTokenizer.from_pretrained(full_model_name, token=hf_token)
+    model = AutoModelForCausalLM.from_pretrained(full_model_name, token=hf_token)
+    model.eval()
+
+# Initialize default model
+load_model("LeCarnet-8M")
+
+# Streamer for real-time generation
+streamer = None
+
+# Streaming generation function
+def respond(message, max_tokens, temperature, top_p):
+    global streamer
+    inputs = tokenizer(message, return_tensors="pt")
+    streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
 
     generate_kwargs = dict(
         **inputs,
@@ -58,61 +43,80 @@ def respond(message, history, model_name, max_tokens, temperature, top_p):
     thread = threading.Thread(target=model.generate, kwargs=generate_kwargs)
     thread.start()
 
-    accumulated = ""  # Removed model name prefix
+    response = ""
     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):
-    for updated_history in respond(message, history, model_name, max_tokens, temperature, top_p):
-        yield updated_history, ""
-
-with gr.Blocks(css=".gr-button {margin: 5px; width: 100%;} .gr-column {padding: 10px;}") as demo:
-    gr.Markdown("# LeCarnet")
-    gr.Markdown("Select a model on the right and type a message to chat.")
-
+        response += new_text
+        yield response
+
+# User input handler
+def user(message, chat_history):
+    chat_history.append([message, None])
+    return "", chat_history
+
+# Bot response handler
+def bot(chatbot, max_tokens, temperature, top_p):
+    message = chatbot[-1][0]
+    response_generator = respond(message, max_tokens, temperature, top_p)
+    for response in response_generator:
+        chatbot[-1][1] = response
+        yield chatbot
+
+# Model selector handler
+def update_model(model_name):
+    load_model(model_name)
+    return []
+
+# Gradio UI
+with gr.Blocks(title="LeCarnet - Chat Interface") as demo:
     with gr.Row():
+        # Left column: Options
+        with gr.Column(scale=1, min_width=150):
+            gr.Markdown("### 🧠 Model Settings")
+            model_selector = gr.Dropdown(
+                choices=["LeCarnet-3M", "LeCarnet-8M", "LeCarnet-21M"],
+                value="LeCarnet-8M",
+                label="Select 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 Sampling")
+            clear_button = gr.Button("🗑️ Clear Chat")
+
+        # Right column: Chat + Image
         with gr.Column(scale=4):
+            gr.Markdown("### 🤖 LeCarnet Chatbot")
+            model_logo = gr.Image(
+                value="media/le-carnet.png",
+                label="Model Logo",
+                height=100,
+                width=100,
+                interactive=False
+            )
             chatbot = gr.Chatbot(
-                avatar_images=(None, "https://raw.githubusercontent.com/maxlsb/le-carnet/main/media/le-carnet.png"),   # Using URL for reliability
-                label="Chat",
-                height=600,
+                bubble_full_width=False,
+                height=500
+            )
+            msg_input = gr.Textbox(
+                placeholder="Type your message and press Enter...",
+                label="User Input"
             )
-            user_input = gr.Textbox(placeholder="Type your message here...", label="Message")
-            submit_btn = gr.Button("Send")
-            examples = gr.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,
-            )
-
-        with gr.Column(scale=1, min_width=200):
-            model_dropdown = gr.Dropdown(
-                choices=["LeCarnet-3M", "LeCarnet-8M", "LeCarnet-21M"],
-                value="LeCarnet-8M",
-                label="Select Model"
+                inputs=msg_input,
+                label="Example Prompts"
             )
-            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")
 
-    # Submit button click
-    submit_btn.click(
-        fn=submit,
-        inputs=[user_input, chatbot, model_dropdown, max_tokens, temperature, top_p],
-        outputs=[chatbot, user_input],
-    )
-    
-    # Enter key press
-    user_input.submit(
-        fn=submit,
-        inputs=[user_input, chatbot, model_dropdown, max_tokens, temperature, top_p],
-        outputs=[chatbot, user_input],
+    # Event handlers
+    model_selector.change(fn=update_model, inputs=[model_selector], outputs=[])
+    msg_input.submit(fn=user, inputs=[msg_input, chatbot], outputs=[msg_input, chatbot], queue=False).then(
+        fn=bot, inputs=[chatbot, max_tokens, temperature, top_p], outputs=[chatbot]
     )
+    clear_button.click(fn=lambda: None, inputs=None, outputs=chatbot, queue=False)
 
 if __name__ == "__main__":
-    demo.queue(default_concurrency_limit=10, max_size=10).launch(ssr_mode=False, max_threads=10)
+    demo.queue()
+    demo.launch()