Update app.py

app.py

@@ -1,33 +1,55 @@
 import os
 import threading
-import gradio as gr
-from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
-
-MODEL_NAMES = ["LeCarnet-3M", "LeCarnet-8M", "LeCarnet-21M"]
-HF_TOKEN = os.environ.get("HUGGINGFACEHUB_API_TOKEN") or os.getenv("HUGGINGFACEHUB_API_TOKEN")
-MEDIA_PATH = "media/le-carnet.png"
-
-models = {}
-tokenizers = {}
+from collections import defaultdict
 
-for name in MODEL_NAMES:
-    hub_id = f"MaxLSB/{name}"
-    tokenizers[name] = AutoTokenizer.from_pretrained(hub_id, token=HF_TOKEN)
-    models[name] = AutoModelForCausalLM.from_pretrained(hub_id, token=HF_TOKEN)
-    models[name].eval()
-
-
-def respond(prompt: str, chat_history, selected_model: str, max_tokens: int, temperature: float, top_p: float):
-    tokenizer = tokenizers[selected_model]
-    model = models[selected_model]
+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["HUGGINGFACEHUB_API_TOKEN"]
+
+# Preload models and tokenizers
+loaded_models = defaultdict(dict)
+
+for name, path in model_name_to_path.items():
+    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()
+
+def respond(
+    prompt: str,
+    chat_history,
+    model_name: str,
+    max_tokens: int,
+    temperature: float,
+    top_p: float,
+):
+    # Select the appropriate model and tokenizer
+    tokenizer = loaded_models[model_name]["tokenizer"]
+    model = loaded_models[model_name]["model"]
+
+    # Tokenize input
     inputs = tokenizer(prompt, 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,
@@ -38,60 +60,38 @@ def respond(prompt: str, chat_history, selected_model: str, max_tokens: int, tem
         eos_token_id=tokenizer.eos_token_id,
     )
 
+    # Run generation in a background thread
     thread = threading.Thread(target=model.generate, kwargs=generate_kwargs)
     thread.start()
 
-    prefix = f"<img src='{MEDIA_PATH}' width='24' style='display:inline; vertical-align:middle; margin-right:6px;'/> <strong>{selected_model}</strong>: "
+    # Stream results
     accumulated = ""
-    first = True
     for new_text in streamer:
-        if first:
-            accumulated = prefix + new_text
-            first = False
-        else:
-            accumulated += new_text
+        accumulated += new_text
         yield accumulated
 
-
-with gr.Blocks(css=".gr-chatbox {height: 600px !important;}") as demo:
-    gr.Markdown("## LeCarnet")
-
-    with gr.Row():
-        with gr.Column(scale=4):
-            with gr.Row():
-                toggle_btn = gr.Button("Show/hide parameters", elem_id="toggle-btn")
-            chat = gr.ChatInterface(
-                fn=respond,
-                additional_inputs=[
-                    gr.Dropdown(MODEL_NAMES, value="LeCarnet-8M", label="Model"),
-                    gr.Slider(1, 512, value=512, step=1, label="Max new tokens"),
-                    gr.Slider(0.1, 2.0, value=0.7, step=0.1, label="Temperature"),
-                    gr.Slider(0.1, 1.0, value=0.9, step=0.05, label="Top‑p"),
-                ],
-                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,
-                type="messages",
-            )
-
-        with gr.Column(scale=1, visible=True, elem_id="settings-panel"):
-            pass  # Inputs are already defined in ChatInterface
-
-    demo.load(
-        js="""
-        () => {
-            const toggleBtn = document.querySelector('#toggle-btn button') || document.querySelector('#toggle-btn');
-            const panel = document.querySelector('#settings-panel');
-            toggleBtn.addEventListener('click', () => {
-                panel.style.display = (panel.style.display === 'none') ? 'flex' : 'none';
-            });
-        }
-        """
-    )
-
+# Create Gradio Chat Interface
+demo = gr.ChatInterface(
+    fn=respond,
+    additional_inputs=[
+        gr.Dropdown(
+            choices=["LeCarnet-3M", "LeCarnet-8M", "LeCarnet-21M"],
+            value="LeCarnet-8M",
+            label="Model",
+        ),
+        gr.Slider(1, 512, value=512, step=1, label="Max New Tokens"),
+        gr.Slider(0.1, 2.0, value=0.7, step=0.1, label="Temperature"),
+        gr.Slider(0.1, 1.0, value=0.9, step=0.05, label="Top-p"),
+    ],
+    title="LeCarnet",
+    description="Select a model and enter text to get started.",
+    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,
+)
 
 if __name__ == "__main__":
     demo.queue(default_concurrency_limit=10, max_size=10).launch(ssr_mode=False, max_threads=10)