Update app.py

app.py

@@ -5,19 +5,23 @@ import json
 from transformers import AutoTokenizer, AutoModelForCausalLM
 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
 
@@ -38,6 +42,21 @@ def generate_response(prompt, max_new_tokens=200):
     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": "",
@@ -54,17 +73,28 @@ state = gr.State({
     },
 })
 
+# 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)
+                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...")
 
-        # Internal State
         current_state = state
 
         def add_message(user_input, state_value):
@@ -77,7 +107,13 @@ with gr.Blocks(css=css) as demo:
             yield {"chatbot": gr.update(value=history)}
 
             # Start assistant response
-            history.append({"role": "assistant", "content": [], "key": str(uuid.uuid4()), "loading": True})
+            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: