Update app.py

app.py

@@ -1,142 +1,99 @@
 import os
-import uuid
-import time
-import json
-from transformers import AutoTokenizer, AutoModelForCausalLM
+import threading
 import gradio as gr
-import modelscope_studio.components.antd as antd
-import modelscope_studio.components.antdx as antdx
-import modelscope_studio.components.base as ms
-import modelscope_studio.components.pro as pro
-
-# Define model paths
-MODEL_PATHS = {
-    "LeCarnet-3M": "MaxLSB/LeCarnet-3M",
-    "LeCarnet-8M": "MaxLSB/LeCarnet-8M",
-    "LeCarnet-21M": "MaxLSB/LeCarnet-21M",
-}
-
-# Set HF token
-hf_token = os.environ.get("HUGGINGFACEHUB_API_TOKEN")
-if not hf_token:
-    raise ValueError("HUGGINGFACEHUB_API_TOKEN environment variable not set.")
-
-# Load tokenizer and model globally
-tokenizer = None
-model = None
-
-def load_model(model_name: str):
-    global tokenizer, model
-    if model_name not in MODEL_PATHS:
-        raise ValueError(f"Unknown model: {model_name}")
-    print(f"Loading {model_name}...")
-    repo = MODEL_PATHS[model_name]
-    tokenizer = AutoTokenizer.from_pretrained(repo, use_auth_token=hf_token)
-    model = AutoModelForCausalLM.from_pretrained(repo, use_auth_token=hf_token)
-    model.eval()
-    print(f"{model_name} loaded.")
-
-def generate_response(prompt, max_new_tokens=200):
-    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
-    outputs = model.generate(**inputs, max_new_tokens=max_new_tokens)
-    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    return response[len(prompt):].strip()
-
-# CSS for styling chatbot header with avatar
-css = """
-.chatbot-chat-messages .ant-pro-chat-message .ant-pro-chat-message-header {
-  display: flex;
-  align-items: center;
-}
-.chatbot-chat-messages .ant-pro-chat-message .ant-pro-chat-message-header img {
-  width: 20px;
-  height: 20px;
-  margin-right: 8px;
-  vertical-align: middle;
-}
-"""
-
-# Default settings
-DEFAULT_SETTINGS = {
-    "model": "LeCarnet-3M",
-    "sys_prompt": "",
-}
-
-# Initial state with one fixed conversation
-state = gr.State({
-    "conversation_id": "default",
-    "conversation_contexts": {
-        "default": {
-            "history": [],
-            "settings": DEFAULT_SETTINGS,
-        }
-    },
-})
-
-# Welcome message (optional)
-def welcome_config():
-    return {
-        "title": "LeCarnet Chatbot",
-        "description": "Start chatting below!",
-        "promptSuggestions": ["Hello", "Tell me a story", "How are you?"]
-    }
-
-with gr.Blocks(css=css) as demo:
-    with ms.Application(), antd.Row(gutter=[20, 20], wrap=False, elem_id="chatbot"):
-        # Right Column - Chat Interface
-        with antd.Col(flex=1, elem_style=dict(height="100%")):
-            with antd.Flex(vertical=True, gap="small", elem_classes="chatbot-chat"):
-                chatbot = pro.Chatbot(
-                    elem_classes="chatbot-chat-messages",
-                    height=0,
-                    welcome_config=welcome_config()
-                )
-                with antdx.Suggestion(items=["Hello", "How are you?", "Tell me something"]) as suggestion:
-                    with ms.Slot("children"):
-                        input = antdx.Sender(placeholder="Type your message here...")
-
-        current_state = state
-
-        def add_message(user_input, state_value):
-            history = state_value["conversation_contexts"]["default"]["history"]
-            settings = state_value["conversation_contexts"]["default"]["settings"]
-            selected_model = settings["model"]
-
-            # Add user message
-            history.append({"role": "user", "content": user_input, "key": str(uuid.uuid4())})
-            yield {"chatbot": gr.update(value=history)}
-
-            # Start assistant response
-            history.append({
-                "role": "assistant",
-                "content": [],
-                "key": str(uuid.uuid4()),
-                "header": f'<img src="/file=media/le-carnet.png" style="width:20px;height:20px;margin-right:8px;"> <span>{selected_model}</span>',
-                "loading": True
-            })
-            yield {"chatbot": gr.update(value=history)}
-
-            try:
-                # Generate model response
-                prompt = "\n".join([msg["content"] for msg in history if msg["role"] == "user"])
-                response = generate_response(prompt)
-
-                # Update assistant message
-                history[-1]["content"] = [{"type": "text", "content": response}]
-                history[-1]["loading"] = False
-                yield {"chatbot": gr.update(value=history)}
-            except Exception as e:
-                history[-1]["content"] = [{
-                    "type": "text",
-                    "content": f'<span style="color: red;">{str(e)}</span>'
-                }]
-                history[-1]["loading"] = False
-                yield {"chatbot": gr.update(value=history)}
-
-        input.submit(fn=add_message, inputs=[input, state], outputs=[chatbot])
-
-# Load default model on startup
-load_model(DEFAULT_SETTINGS["model"])
-
+from transformers import (
+    AutoModelForCausalLM,
+    AutoTokenizer,
+    TextIteratorStreamer,
+)
+
+# Configuration
+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"  # Relative path to logo
+
+# Pre-load all tokenizers and models
+models = {}
+tokenizers = {}
+for name in MODEL_NAMES:
+    hub_id = f"MaxLSB/LeCarnet-{name.split('-')[-1]}M"
+    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,
+):
+    """
+    Generate a streaming response from the chosen LeCarnet model,
+    prepending the logo and model name in the chat bubble.
+    """
+    tokenizer = tokenizers[selected_model]
+    model = models[selected_model]
+    inputs = tokenizer(prompt, return_tensors="pt")
+
+    streamer = TextIteratorStreamer(
+        tokenizer,
+        skip_prompt=False,
+        skip_special_tokens=True,
+    )
+
+    generate_kwargs = dict(
+        **inputs,
+        streamer=streamer,
+        max_new_tokens=max_tokens,
+        do_sample=True,
+        temperature=temperature,
+        top_p=top_p,
+        eos_token_id=tokenizer.eos_token_id,
+    )
+
+    # Start generation in background thread
+    thread = threading.Thread(target=model.generate, kwargs=generate_kwargs)
+    thread.start()
+
+    prefix = f"<img src='{MEDIA_PATH}' alt='logo' width='20' style='vertical-align: middle;'/> <strong>{selected_model}</strong>: "
+    accumulated = ""
+    first = True
+    for new_text in streamer:
+        if first:
+            # include prefix only once at start
+            accumulated = prefix + new_text
+            first = False
+        else:
+            accumulated += new_text
+        yield accumulated
+
+
+# Build Gradio ChatInterface
+with gr.Blocks() as demo:
+    gr.Markdown("# LeCarnet: Short French Stories")
+    with gr.Row():
+        with gr.Column():
+            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"),
+                ],
+                title="LeCarnet Chat",
+                description="Type the beginning of a sentence and watch the model finish it.",
+                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).launch()
+    demo.queue()
+    demo.launch()