Update app.py

app.py

@@ -6,6 +6,7 @@ from itertools import islice
 from datetime import datetime
 import re  # for parsing <think> blocks
 import gradio as gr
+from typing import Dict, Union
 import torch
 from transformers import pipeline, TextIteratorStreamer
 from transformers import AutoTokenizer
@@ -245,184 +246,180 @@ def update_default_prompt(enable_search):
     today = datetime.now().strftime('%Y-%m-%d')
     return f"You are a helpful assistant. Today is {today}."
 
+# ------------------------------
+# Gradio UI
+# ------------------------------
 
-def toggle_theme(current_theme):
-    """Toggle between light and dark themes"""
-    if current_theme == "light":
-        return "dark", "☀️ Light Mode"
-    else:
-        return "light", "🌙 Dark Mode"
+# UI strings translations
+UI_TEXTS: Dict[str, Dict[str, Union[str, Dict[str, str]]]] = {
+    "en": {
+        "title": "Yee-R1 Demo",
+        "description": "Yee AI Data Security Expert",
+        "select_model": "Select Model",
+        "enable_search": "Enable Web Search",
+        "system_prompt": "System Prompt",
+        "generation_parameters": "Generation Parameters",
+        "max_tokens": "Max Tokens",
+        "temperature": "Temperature",
+        "top_k": "Top-K",
+        "top_p": "Top-P",
+        "repeat_penalty": "Repetition Penalty",
+        "web_search_settings": "Web Search Settings",
+        "max_results": "Max Results",
+        "max_chars_result": "Max Chars/Result",
+        "search_timeout": "Search Timeout (s)",
+        "clear_chat": "Clear Chat",
+        "cancel_generation": "Cancel Generation",
+        "chat_placeholder": "Type your message and press Enter...",
+        "theme_label": "Select Theme",
+        "language_label": "Select Language",
+        "theme_light": "Light",
+        "theme_dark": "Dark",
+        "language_en": "English",
+        "language_zh": "Chinese",
+    },
+    "zh": {
+        "title": "小熠演示",
+        "description": "小熠AI数据安全专家",
+        "select_model": "选择模型",
+        "enable_search": "启用网络搜索",
+        "system_prompt": "系统提示词",
+        "generation_parameters": "生成参数",
+        "max_tokens": "最大生成长度",
+        "temperature": "温度",
+        "top_k": "Top-K",
+        "top_p": "Top-P",
+        "repeat_penalty": "重复惩罚",
+        "web_search_settings": "网络搜索设置",
+        "max_results": "最大结果数",
+        "max_chars_result": "每个结果最大字符数",
+        "search_timeout": "搜索超时 (秒)",
+        "clear_chat": "清空聊天",
+        "cancel_generation": "取消生成",
+        "chat_placeholder": "请输入消息，按回车发送...",
+        "theme_label": "选择主题",
+        "language_label": "选择语言",
+        "theme_light": "明亮模式",
+        "theme_dark": "暗黑模式",
+        "language_en": "English",
+        "language_zh": "中文",
+    }
+}
 
+def get_ui_text(language, key):
+    return UI_TEXTS.get(language, UI_TEXTS["en"]).get(key, "")
 
-def toggle_language(current_lang):
-    """Toggle between Chinese and English"""
-    if current_lang == "zh":
-        return "en"
-    else:
-        return "zh"
-
-
-def get_ui_text(lang):
-    """Get UI text based on language"""
-    texts = {
-        "zh": {
-            "title": "## Yee-R1 Demo",
-            "subtitle": "小熠（Yee）AI 数据安全专家",
-            "dark_mode": "🌙 暗黑模式",
-            "light_mode": "☀️ 明亮模式",
-            "lang_btn": "🌐 English",
-            "select_model": "选择模型",
-            "enable_search": "启用网络搜索",
-            "system_prompt": "系统提示词",
-            "gen_params": "### 生成参数",
-            "max_tokens": "最大令牌数",
-            "temperature": "温度",
-            "top_k": "Top-K",
-            "top_p": "Top-P", 
-            "repeat_penalty": "重复惩罚",
-            "search_settings": "### 网络搜索设置",
-            "max_results": "最大结果数",
-            "max_chars": "每个结果最大字符数",
-            "search_timeout": "搜索超时时间 (秒)",
-            "clear_chat": "清空对话",
-            "cancel_gen": "取消生成",
-            "placeholder": "输入您的消息并按回车..."
-        },
-        "en": {
-            "title": "## Yee-R1 Demo",
-            "subtitle": "Yee AI Data Security Expert",
-            "dark_mode": "🌙 Dark Mode",
-            "light_mode": "☀️ Light Mode",
-            "lang_btn": "🌐 中文",
-            "select_model": "Select Model",
-            "enable_search": "Enable Web Search",
-            "system_prompt": "System Prompt",
-            "gen_params": "### Generation Parameters",
-            "max_tokens": "Max Tokens",
-            "temperature": "Temperature",
-            "top_k": "Top-K",
-            "top_p": "Top-P",
-            "repeat_penalty": "Repetition Penalty",
-            "search_settings": "### Web Search Settings",
-            "max_results": "Max Results",
-            "max_chars": "Max Chars/Result",
-            "search_timeout": "Search Timeout (s)",
-            "clear_chat": "Clear Chat",
-            "cancel_gen": "Cancel Generation",
-            "placeholder": "Type your message and press Enter..."
-        }
-    }
-    return texts[lang]
+with gr.Blocks(title=get_ui_text("en", "title")) as demo:
+    theme_dropdown = gr.Dropdown(label=get_ui_text("en", "theme_label"),
+                                 choices=[get_ui_text("en", "theme_light"), get_ui_text("en", "theme_dark")],
+                                 value=get_ui_text("en", "theme_light"))
+    language_dropdown = gr.Dropdown(label=get_ui_text("en", "language_label"),
+                                   choices=[get_ui_text("en", "language_en"), get_ui_text("en", "language_zh")],
+                                   value=get_ui_text("en", "language_en"))
 
+    title_md = gr.Markdown(value=get_ui_text("en", "title"))
+    description_md = gr.Markdown(value=get_ui_text("en", "description"))
 
-# ------------------------------
-# Gradio UI
-# ------------------------------
-with gr.Blocks(title="Yee-R1-Demo", theme=gr.themes.Default()) as demo:
-    # States
-    theme_state = gr.State("light")
-    lang_state = gr.State("zh")
-    
-    # Header with controls
-    with gr.Row():
-        title_md = gr.Markdown("## Yee-R1 Demo")
-        with gr.Row(scale=0):
-            lang_btn = gr.Button("🌐 English", size="sm")
-            theme_btn = gr.Button("🌙 暗黑模式", size="sm")
-    
-    subtitle_md = gr.Markdown("小熠（Yee）AI 数据安全专家")
-    
     with gr.Row():
         with gr.Column(scale=3):
-            model_dd = gr.Dropdown(label="选择模型", choices=list(MODELS.keys()), value=list(MODELS.keys())[0])
-            search_chk = gr.Checkbox(label="启用网络搜索", value=True)
-            sys_prompt = gr.Textbox(label="系统提示词", lines=3, value=update_default_prompt(search_chk.value))
-            gen_params_md = gr.Markdown("### 生成参数")
-            max_tok = gr.Slider(64, 16384, value=4096, step=32, label="最大令牌数")
-            temp = gr.Slider(0.1, 2.0, value=0.6, step=0.1, label="温度")
+            model_dd = gr.Dropdown(label=get_ui_text("en", "select_model"), choices=list(MODELS.keys()), value=list(MODELS.keys())[0])
+            search_chk = gr.Checkbox(label=get_ui_text("en", "enable_search"), value=True)
+            sys_prompt = gr.Textbox(label=get_ui_text("en", "system_prompt"), lines=3, value=update_default_prompt(search_chk.value))
+            gr.Markdown(f"### {get_ui_text('en', 'generation_parameters')}")
+            max_tok = gr.Slider(64, 16384, value=4096, step=32, label=get_ui_text("en", "max_tokens"))
+            temp = gr.Slider(0.1, 2.0, value=0.6, step=0.1, label=get_ui_text("en", "temperature"))
             k = gr.Slider(1, 100, value=40, step=1, label="Top-K")
             p = gr.Slider(0.1, 1.0, value=0.9, step=0.05, label="Top-P")
-            rp = gr.Slider(1.0, 2.0, value=1.2, step=0.1, label="重复惩罚")
-            search_settings_md = gr.Markdown("### 网络搜索设置")
-            mr = gr.Number(value=6, precision=0, label="最大结果数")
-            mc = gr.Number(value=600, precision=0, label="每个结果最大字符数")
-            st = gr.Slider(minimum=0.0, maximum=30.0, step=0.5, value=5.0, label="搜索超时时间 (秒)")
-            clr = gr.Button("清空对话")
-            cnl = gr.Button("取消生成")
+            rp = gr.Slider(1.0, 2.0, value=1.2, step=0.1, label=get_ui_text("en", "repeat_penalty"))
+            gr.Markdown(f"### {get_ui_text('en', 'web_search_settings')}")
+            mr = gr.Number(value=6, precision=0, label=get_ui_text("en", "max_results"))
+            mc = gr.Number(value=600, precision=0, label=get_ui_text("en", "max_chars_result"))
+            st = gr.Slider(minimum=0.0, maximum=30.0, step=0.5, value=5.0, label=get_ui_text("en", "search_timeout"))
+            clr = gr.Button(get_ui_text("en", "clear_chat"))
+            cnl = gr.Button(get_ui_text("en", "cancel_generation"))
         with gr.Column(scale=7):
-            chat = gr.Chatbot(type="messages", show_copy_all_button=True, height="50vh")
-            txt = gr.Textbox(placeholder="输入您的消息并按回车...")
+            chat = gr.Chatbot(type="messages", show_copy_all_button=True)
+            txt = gr.Textbox(placeholder=get_ui_text("en", "chat_placeholder"))
             dbg = gr.Markdown()
 
-    # Event handlers
-    search_chk.change(fn=update_default_prompt, inputs=search_chk, outputs=sys_prompt)
-    clr.click(fn=lambda: ([], "", ""), outputs=[chat, txt, dbg])
-    cnl.click(fn=cancel_generation, outputs=dbg)
-    
-    # Theme toggle functionality
-    def handle_theme_toggle(current_theme, current_lang):
-        new_theme, _ = toggle_theme(current_theme)
-        ui_text = get_ui_text(current_lang)
-        new_btn_text = ui_text["light_mode"] if new_theme == "dark" else ui_text["dark_mode"]
-        
-        if new_theme == "dark":
-            demo._theme = gr.themes.Monochrome()
-        else:
-            demo._theme = gr.themes.Default()
-        return new_theme, new_btn_text
-    
-    # Language toggle functionality
-    def handle_language_toggle(current_lang, current_theme):
-        new_lang = toggle_language(current_lang)
-        ui_text = get_ui_text(new_lang)
-        
-        # Update all UI text
-        updates = [
-            new_lang,  # lang_state
-            ui_text["lang_btn"],  # lang_btn
-            ui_text["light_mode"] if current_theme == "dark" else ui_text["dark_mode"],  # theme_btn
-            ui_text["title"],  # title_md
-            ui_text["subtitle"],  # subtitle_md
-            ui_text["select_model"],  # model_dd label
-            ui_text["enable_search"],  # search_chk label
-            ui_text["system_prompt"],  # sys_prompt label
-            ui_text["gen_params"],  # gen_params_md
-            ui_text["max_tokens"],  # max_tok label
-            ui_text["temperature"],  # temp label
-            ui_text["top_k"],  # k label
-            ui_text["top_p"],  # p label
-            ui_text["repeat_penalty"],  # rp label
-            ui_text["search_settings"],  # search_settings_md
-            ui_text["max_results"],  # mr label
-            ui_text["max_chars"],  # mc label
-            ui_text["search_timeout"],  # st label
-            ui_text["clear_chat"],  # clr
-            ui_text["cancel_gen"],  # cnl
-            ui_text["placeholder"]  # txt placeholder
-        ]
-        
-        return updates
-    
-    theme_btn.click(
-        fn=handle_theme_toggle,
-        inputs=[theme_state, lang_state],
-        outputs=[theme_state, theme_btn]
+    # Function to update UI texts when language or theme changes
+    def update_ui(language, theme):
+        texts = UI_TEXTS.get("en") if language == UI_TEXTS["en"]["language_en"] else UI_TEXTS.get("zh")
+        # Map passed language (English/Chinese) to keys
+        lang_key = "en" if language == UI_TEXTS["en"]["language_en"] else "zh"
+        # Update labels
+        return (texts["title"], texts["description"],
+                texts["select_model"], texts["enable_search"], texts["system_prompt"],
+                texts["generation_parameters"], texts["max_tokens"], texts["temperature"], texts["repeat_penalty"],
+                texts["web_search_settings"], texts["max_results"], texts["max_chars_result"], texts["search_timeout"],
+                texts["clear_chat"], texts["cancel_generation"], texts["chat_placeholder"],
+                gr.themes.Dark() if theme == texts["theme_dark"] else gr.themes.Default())
+
+    def toggle_language_and_theme(language, theme):
+        # Return updated texts and theme to update all UI elements
+        lang_key = "en" if language == UI_TEXTS["en"]["language_en"] else "zh"
+        texts = UI_TEXTS[lang_key]
+        return {
+            "title_md": texts["title"],
+            "description_md": texts["description"],
+            "model_dd_label": texts["select_model"],
+            "search_chk_label": texts["enable_search"],
+            "sys_prompt_label": texts["system_prompt"],
+            "max_tok_label": texts["max_tokens"],
+            "temp_label": texts["temperature"],
+            "rp_label": texts["repeat_penalty"],
+            "mr_label": texts["max_results"],
+            "mc_label": texts["max_chars_result"],
+            "st_label": texts["search_timeout"],
+            "clr_label": texts["clear_chat"],
+            "cnl_label": texts["cancel_generation"],
+            "txt_placeholder": texts["chat_placeholder"],
+            "theme_obj": gr.themes.Dark() if theme == texts["theme_dark"] else gr.themes.Default()
+        }
+
+    # Update UI text labels on language or theme change
+    language_dropdown.change(
+        fn=toggle_language_and_theme,
+        inputs=[language_dropdown, theme_dropdown],
+        outputs=[
+            title_md, description_md, model_dd, search_chk, sys_prompt,
+            max_tok, temp, rp, mr, mc, st, clr, cnl, txt
+        ],
+        _js="""
+        function updateLabels(resp) {
+            title_md.textContent = resp.title_md;
+            description_md.textContent = resp.description_md;
+            model_dd.label = resp.model_dd_label;
+            search_chk.label = resp.search_chk_label;
+            sys_prompt.label = resp.sys_prompt_label;
+            max_tok.label = resp.max_tok_label;
+            temp.label = resp.temp_label;
+            rp.label = resp.rp_label;
+            mr.label = resp.mr_label;
+            mc.label = resp.mc_label;
+            st.label = resp.st_label;
+            clr.textContent = resp.clr_label;
+            cnl.textContent = resp.cnl_label;
+            txt.placeholder = resp.txt_placeholder;
+            return;
+        }
+        """
     )
-    
-    lang_btn.click(
-        fn=handle_language_toggle,
-        inputs=[lang_state, theme_state],
+
+    theme_dropdown.change(
+        fn=toggle_language_and_theme,
+        inputs=[language_dropdown, theme_dropdown],
         outputs=[
-            lang_state, lang_btn, theme_btn, title_md, subtitle_md,
-            model_dd, search_chk, sys_prompt, gen_params_md,
-            max_tok, temp, k, p, rp, search_settings_md,
-            mr, mc, st, clr, cnl, txt
+            title_md, description_md, model_dd, search_chk, sys_prompt,
+            max_tok, temp, rp, mr, mc, st, clr, cnl, txt
         ]
     )
-     
+
+    search_chk.change(fn=update_default_prompt, inputs=search_chk, outputs=sys_prompt)
+    clr.click(fn=lambda: ([], "", ""), outputs=[chat, txt, dbg])
+    cnl.click(fn=cancel_generation, outputs=dbg)
     txt.submit(fn=chat_response,
                inputs=[txt, chat, sys_prompt, search_chk, mr, mc,
                        model_dd, max_tok, temp, k, p, rp, st],
                outputs=[chat, dbg])
-    
     demo.launch()