Update app.py

app.py

@@ -1,34 +1,30 @@
 import os
-import threading
+import uuid
+import time
+import json
+from transformers import AutoTokenizer, AutoModelForCausalLM
 import gradio as gr
-from transformers import (
-    AutoModelForCausalLM,
-    AutoTokenizer,
-    TextIteratorStreamer,
-)
+import modelscope_studio.components.antd as antd
+import modelscope_studio.components.base as ms
+import modelscope_studio.components.pro as pro
 
-# Define your models
 MODEL_PATHS = {
     "LeCarnet-3M": "MaxLSB/LeCarnet-3M",
     "LeCarnet-8M": "MaxLSB/LeCarnet-8M",
     "LeCarnet-21M": "MaxLSB/LeCarnet-21M",
 }
 
-# Add your Hugging Face token
 hf_token = os.environ.get("HUGGINGFACEHUB_API_TOKEN")
 if not hf_token:
     raise ValueError("HUGGINGFACEHUB_API_TOKEN environment variable not set.")
 
-# Load tokenizers & models - only load one initially
 tokenizer = None
 model = None
 
 def load_model(model_name: str):
-    """Loads the specified model and tokenizer."""
     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)
@@ -36,98 +32,75 @@ def load_model(model_name: str):
     model.eval()
     print(f"{model_name} loaded.")
 
-# Initial model load
-initial_model = list(MODEL_PATHS.keys())[0]
-load_model(initial_model)
-
+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()
 
-def respond(
-    prompt: str,
-    chat_history: list,
-    model_choice: str,
-    max_tokens: int,
-    temperature: float,
-    top_p: float,
-):
-    global tokenizer, model
+DEFAULT_SETTINGS = {
+    "model": "LeCarnet-3M",
+    "sys_prompt": "",
+}
 
-    # Reload model if it's not the currently loaded one
-    current_path = getattr(model.config, "_name_or_path", None)
-    desired_path = MODEL_PATHS[model_choice]
-    if current_path != desired_path:
-        load_model(model_choice)
-
-    # Tokenize
-    inputs = tokenizer(prompt, return_tensors="pt")
-    streamer = TextIteratorStreamer(
-        tokenizer,
-        skip_prompt=False,
-        skip_special_tokens=True,
-    )
-
-    # Prepare generation kwargs
-    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,
-    )
-
-    # Launch generation in a background thread
-    thread = threading.Thread(target=model.generate, kwargs=generate_kwargs)
-    thread.start()
-
-    # Stream back to the UI
-    accumulated = ""
-    for new_text in streamer:
-        accumulated += new_text
-        yield accumulated
-
-
-# If you have custom CSS, define it here; otherwise set to None or remove the css= line below
-custom_css = None
-
-with gr.Blocks(css=custom_css, fill_width=True) as demo:
-    with gr.Row():
-        with gr.Column(scale=1):
-            model_dropdown = gr.Dropdown(
-                choices=list(MODEL_PATHS.keys()),
-                value=initial_model,
-                label="Choose Model",
-                interactive=True
-            )
-            max_tokens_slider = gr.Slider(
-                minimum=1, maximum=512, value=512, step=1, label="Max new tokens"
-            )
-            temperature_slider = gr.Slider(
-                minimum=0.1, maximum=2.0, value=0.7, step=0.1, label="Temperature"
-            )
-            top_p_slider = gr.Slider(
-                minimum=0.1, maximum=1.0, value=0.9, step=0.05, label="Top‑p"
-            )
-
-        with gr.Column(scale=3):
-            chatbot = gr.ChatInterface(
-                fn=respond,
-                additional_inputs=[
-                    model_dropdown,
-                    max_tokens_slider,
-                    temperature_slider,
-                    top_p_slider,
-                ],
-                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,
-                submit_btn="Generate",
-                avatar_images=(None, "media/le-carnet.png")
-            )
+# Initial state with one fixed conversation
+state = gr.State({
+    "conversation_id": "default",
+    "conversation_contexts": {
+        "default": {
+            "history": [],
+            "settings": DEFAULT_SETTINGS,
+        }
+    },
+})
+
+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)
+                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...")
+
+        # Internal State
+        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()), "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"])
 
 if __name__ == "__main__":
-    demo.queue()
-    demo.launch()
+    demo.queue(default_concurrency_limit=10).launch()