import gradio as gr import torch from typing import List, Dict, Generator, Tuple, Optional, Any from codon.motif import MotifA1, MotifA1Tokenizer from codon.utils.generate import chat, ChatChunk from codon.utils.tokens import PackedTokenizer LOCALIZATION = { 'English': { 'title': 'Motif-A1 Interactive Web Demo', 'subtitle': 'Explore the experimental 105.41M reasoning language model from the CodonProject. Built-in with native Chain-of-Thought (CoT) support.', 'model_control': 'Model Status', 'status_not_loaded': 'Loading model... Please wait...', 'status_loading': 'Loading model... Please wait...', 'status_loaded': 'Model and tokenizer successfully loaded!', 'gen_params': 'Generation Parameters', 'temp_label': 'Temperature', 'temp_info': 'Higher values increase creativity, lower values increase precision.', 'max_tokens_label': 'Max New Tokens', 'max_tokens_info': 'Maximum number of generated tokens.', 'top_p_label': 'Top P (Nucleus Sampling)', 'top_p_info': 'Set to 1.0 to disable. Keep below 1.0 for standard nucleus sampling.', 'top_k_label': 'Top K Sampling', 'top_k_info': 'Set to 0 to disable. Standard filter limit.', 'persona': 'Persona Setting', 'sys_prompt_label': 'System Prompt', 'sys_prompt_val': 'You are Motif-A1, an intelligent AI assistant developed by CodonProject. Think carefully step-by-step.', 'sys_prompt_placeholder': 'Type custom behavior or persona guidelines here...', 'chatbot_label': 'Motif-A1 Conversation Window', 'user_input_placeholder': 'Type your prompt here and press Enter (or click Send)...', 'submit_btn': 'Send', 'stop_btn': 'Stop', 'clear_btn': 'Clear History', 'thinking_active': 'Thinking Process (thinking...)', 'thinking_collapsed': 'Thinking Process (collapsed)', 'warning_not_loaded': 'Model is not loaded yet! Attempting auto-loading...', 'error_load_failed': 'Error: Model failed to load automatically.', 'error_gen_failed': 'Generation Error:' }, '中文': { 'title': 'Motif-A1 交互式网页演示', 'subtitle': '探索由 CodonProject 开发的实验性 105.41M 推理大语言模型。内置原生思维链 (CoT) 支持。', 'model_control': '模型状态', 'status_not_loaded': '正在加载模型，请稍候...', 'status_loading': '正在加载模型，请稍候...', 'status_loaded': '模型和分词器已成功加载！', 'gen_params': '生成参数', 'temp_label': '温度 (Temperature)', 'temp_info': '较高的值增加创造力，较低的值提高精准度。', 'max_tokens_label': '最大生成长度 (Max New Tokens)', 'max_tokens_info': '生成文本的最大 Token 数量。', 'top_p_label': 'Top P (核采样)', 'top_p_info': '设置为 1.0 以禁用。保持小于 1.0 进行核采样。', 'top_k_label': 'Top K 采样', 'top_k_info': '设置为 0 以禁用。标准过滤器限制。', 'persona': '角色设定', 'sys_prompt_label': '系统提示词', 'sys_prompt_val': '你是由 CodonProject 开发的智能 AI 助手 Motif-A1。请一步步仔细思考。', 'sys_prompt_placeholder': '在此处输入自定义行为或角色设定指南...', 'chatbot_label': 'Motif-A1 对话窗口', 'user_input_placeholder': '在此输入您的内容并按回车（或点击发送）...', 'submit_btn': '发送', 'stop_btn': '停止', 'clear_btn': '清空历史', 'thinking_active': '思考过程 (思考中...)', 'thinking_collapsed': '思考过程 (已折叠)', 'warning_not_loaded': '模型尚未加载！正在尝试自动加载...', 'error_load_failed': '错误：模型自动加载失败。', 'error_gen_failed': '生成错误：' } } class ModelManager: ''' Manager class to handle model and tokenizer loading and generation. Attributes: model (Optional[MotifA1]): The language model instance. tokenizer (Optional[PackedTokenizer]): The tokenizer instance. device (torch.device): The device used for computation. ''' def __init__(self) -> None: self.model: Optional[MotifA1] = None self.tokenizer: Optional[PackedTokenizer] = None self.device: torch.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') def load_model(self) -> str: ''' Loads the model and tokenizer from remote source. Returns: str: Loading status message. ''' try: self.model = MotifA1().from_remote().to(self.device) self.tokenizer = MotifA1Tokenizer().from_remote() return f'Model and tokenizer successfully loaded on {self.device}!' except Exception as e: return f'Error loading model: {str(e)}' def is_loaded(self) -> bool: ''' Checks if both the model and tokenizer are loaded. Returns: bool: True if loaded, False otherwise. ''' return self.model is not None and self.tokenizer is not None # Global instance of ModelManager model_manager = ModelManager() def handle_user_message( user_message: str, chat_history: List[Dict[str, Any]], clean_messages: List[Dict[str, str]] ) -> Tuple[str, List[Dict[str, Any]], List[Dict[str, str]]]: ''' Handles the user input, adding it to the history list and cleaning the input box. Args: user_message (str): The raw text entered by the user. chat_history (List[Dict[str, Any]]): Chatbot display history with role and content. clean_messages (List[Dict[str, str]]): Clean dialog list for model context. Returns: Tuple[str, List[Dict[str, Any]], List[Dict[str, str]]]: Cleaned text, updated display history, and clean history. ''' if not user_message.strip(): return '', chat_history, clean_messages chat_history.append({'role': 'user', 'content': user_message}) clean_messages.append({'role': 'user', 'content': user_message}) chat_history.append({'role': 'assistant', 'content': ''}) return '', chat_history, clean_messages def generate_response( chat_history: List[Dict[str, Any]], clean_messages: List[Dict[str, str]], system_prompt: str, temperature: float, max_new_tokens: int, top_p: float, top_k: int, language: str ) -> Generator[Tuple[List[Dict[str, Any]], List[Dict[str, str]]], None, None]: ''' Generates model response in a streaming way and updates the UI in real time. Args: chat_history (List[Dict[str, Any]]): Chat history with UI HTML formatting. clean_messages (List[Dict[str, str]]): Clean message history. system_prompt (str): Custom system instruction. temperature (float): Generation temperature. max_new_tokens (int): Max response length. top_p (float): Sampling probability threshold. top_k (int): Top-K sampling filter. language (str): UI language choice ('English' or '中文'). Yields: Tuple[List[Dict[str, Any]], List[Dict[str, str]]]: Streaming updates to chatbot and clean state. ''' loc = LOCALIZATION[language] if not model_manager.is_loaded(): gr.Warning(loc['warning_not_loaded']) status = model_manager.load_model() if not model_manager.is_loaded(): chat_history[-1]['content'] = f'{loc["error_load_failed"]} {status}' yield chat_history, clean_messages return tk: Optional[int] = int(top_k) if top_k > 0 else None tp: Optional[float] = float(top_p) if top_p < 1.0 else None messages_to_send: List[Dict[str, str]] = [] if system_prompt.strip(): messages_to_send.append({'role': 'system', 'content': system_prompt.strip()}) messages_to_send.extend(clean_messages) cot_buffer = '' response_buffer = '' has_cot = False try: generator = chat( model=model_manager.model, tokenizer=model_manager.tokenizer, device=model_manager.device, messages=messages_to_send, max_new_tokens=max_new_tokens, temperature=temperature, top_k=tk, top_p=tp ) for chunk in generator: if chunk.is_cot: has_cot = True cot_buffer += chunk.content else: response_buffer += chunk.content display_text = '' if has_cot: if chunk.is_cot: display_text += ( f'
\n' f'{loc["thinking_active"]}\n\n' f'{cot_buffer}\n\n' f'
' ) else: display_text += ( f'
\n' f'{loc["thinking_collapsed"]}\n\n' f'{cot_buffer}\n\n' f'
\n\n' f'{response_buffer}' ) else: display_text += response_buffer chat_history[-1]['content'] = display_text yield chat_history, clean_messages clean_messages.append({'role': 'assistant', 'content': response_buffer}) yield chat_history, clean_messages except Exception as e: error_msg = f'\n\n{loc["error_gen_failed"]} {str(e)}' chat_history[-1]['content'] += error_msg yield chat_history, clean_messages def clear_history() -> Tuple[List[Any], List[Any]]: ''' Clears both UI chat history and clean message history. Returns: Tuple[List[Any], List[Any]]: Two empty lists. ''' return [], [] def autoload_callback(language: str) -> Generator[str, None, None]: ''' Triggers automatically when the Web page is first opened to load remote model. Args: language (str): Default language choice. Yields: str: Automatic loading status message. ''' loc = LOCALIZATION[language] yield loc['status_loading'] status = model_manager.load_model() if model_manager.is_loaded(): yield loc['status_loaded'] else: yield f'{loc["error_load_failed"]} {status}' def toggle_ui_language(lang: str) -> Tuple[gr.update, ...]: ''' Switches UI element texts dynamically depending on language choice. Args: lang (str): 'English' or '中文'. Returns: Tuple[gr.update, ...]: Element updates. ''' loc = LOCALIZATION[lang] header_html = ( f'
' f'

{loc["title"]}

' f'

{loc["subtitle"]}

' f'
' ) status_msg = loc['status_loaded'] if model_manager.is_loaded() else loc['status_not_loaded'] return ( gr.update(value=header_html), gr.update(value=f'### {loc["model_control"]}'), gr.update(label=loc['model_control'], value=status_msg), gr.update(value=f'### {loc["gen_params"]}'), gr.update(label=loc['temp_label'], info=loc['temp_info']), gr.update(label=loc['max_tokens_label'], info=loc['max_tokens_info']), gr.update(label=loc['top_p_label'], info=loc['top_p_info']), gr.update(label=loc['top_k_label'], info=loc['top_k_info']), gr.update(value=f'### {loc["persona"]}'), gr.update(label=loc['sys_prompt_label'], value=loc['sys_prompt_val'], placeholder=loc['sys_prompt_placeholder']), gr.update(label=loc['chatbot_label']), gr.update(placeholder=loc['user_input_placeholder']), gr.update(value=loc['submit_btn']), gr.update(value=loc['stop_btn']), gr.update(value=loc['clear_btn']) ) # Setup UI Blocks with gr.Blocks(title='Motif-A1 Interactive Web Demo', theme=gr.themes.Soft()) as demo: # Header Section header = gr.HTML( '''

Motif-A1 Interactive Web Demo

Explore the experimental 105.41M reasoning language model from the CodonProject. Built-in with native Chain-of-Thought (CoT) support.

''' ) # Language Toggle Widget with gr.Row(): lang_radio = gr.Radio( choices=['English', '中文'], value='English', label='Language / 语言', interactive=True ) with gr.Row(): # Sidebar Panel (30% width) with gr.Column(scale=3): # Model Control Section with gr.Group(): model_control_hdr = gr.Markdown('### Model Status') status_box = gr.Textbox( value='Loading model... Please wait...', label='Model Status', interactive=False ) # Hyperparameters Section with gr.Group(): gen_params_hdr = gr.Markdown('### Generation Parameters') temp_slider = gr.Slider( minimum=0.1, maximum=1.5, value=0.3, step=0.05, label='Temperature', info='Higher values increase creativity, lower values increase precision.' ) max_tokens_slider = gr.Slider( minimum=64, maximum=2048, value=1024, step=64, label='Max New Tokens', info='Maximum number of generated tokens.' ) top_p_slider = gr.Slider( minimum=0.0, maximum=1.0, value=0.95, step=0.05, label='Top P (Nucleus Sampling)', info='Set to 1.0 to disable. Keep below 1.0 for standard nucleus sampling.' ) top_k_slider = gr.Slider( minimum=0, maximum=100, value=0, step=1, label='Top K Sampling', info='Set to 0 to disable. Standard filter limit.' ) # System Prompt Customization with gr.Group(): persona_hdr = gr.Markdown('### Persona Setting') sys_prompt_box = gr.Textbox( value='You are Motif-A1, an intelligent AI assistant developed by CodonProject. Think carefully step-by-step.', lines=3, label='System Prompt', placeholder='Type custom behavior or persona guidelines here...' ) # Chatbot Panel (70% width) with gr.Column(scale=7): chatbot = gr.Chatbot( label='Motif-A1 Conversation Window', height=650 ) with gr.Row(): user_input = gr.Textbox( placeholder='Type your prompt here and press Enter (or click Send)...', lines=3, scale=8, label='Your Prompt', show_label=False ) with gr.Column(scale=2, min_width=100): submit_btn = gr.Button('Send', variant='primary') stop_btn = gr.Button('Stop', variant='stop') clear_btn = gr.Button('Clear History') # State variables to persist conversation histories clean_state = gr.State(value=[]) # Core conversation event stream (Send button click) submit_event = submit_btn.click( fn=handle_user_message, inputs=[user_input, chatbot, clean_state], outputs=[user_input, chatbot, clean_state], queue=False ).then( fn=generate_response, inputs=[chatbot, clean_state, sys_prompt_box, temp_slider, max_tokens_slider, top_p_slider, top_k_slider, lang_radio], outputs=[chatbot, clean_state] ) # Support textbox submission (Enter without Shift) user_input_event = user_input.submit( fn=handle_user_message, inputs=[user_input, chatbot, clean_state], outputs=[user_input, chatbot, clean_state], queue=False ).then( fn=generate_response, inputs=[chatbot, clean_state, sys_prompt_box, temp_slider, max_tokens_slider, top_p_slider, top_k_slider, lang_radio], outputs=[chatbot, clean_state] ) # Stop button click triggers instant Cancellation of both possible submit flows stop_btn.click( fn=None, inputs=None, outputs=None, cancels=[submit_event, user_input_event] ) # Language Radio change trigger lang_radio.change( fn=toggle_ui_language, inputs=[lang_radio], outputs=[ header, model_control_hdr, status_box, gen_params_hdr, temp_slider, max_tokens_slider, top_p_slider, top_k_slider, persona_hdr, sys_prompt_box, chatbot, user_input, submit_btn, stop_btn, clear_btn ] ) # Clear screen event clear_btn.click( fn=clear_history, inputs=[], outputs=[chatbot, clean_state], queue=False ) # Web page auto-load trigger on first open demo.load( fn=autoload_callback, inputs=[lang_radio], outputs=[status_box] ) if __name__ == '__main__': # Launch local server demo.queue().launch()