app.py

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'<details open>\n'
                        f'<summary><b>{loc["thinking_active"]}</b></summary>\n\n'
                        f'<span style="color: #4a90e2; font-style: italic;">{cot_buffer}</span>\n\n'
                        f'</details>'
                    )
                else:
                    display_text += (
                        f'<details>\n'
                        f'<summary><b>{loc["thinking_collapsed"]}</b></summary>\n\n'
                        f'<span style="color: #6c7a89; font-style: italic;">{cot_buffer}</span>\n\n'
                        f'</details>\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'<div style="text-align: center; margin-bottom: 20px;">'
        f'<h1 style="color: #2e4057; font-weight: 800; font-size: 2.5rem; margin-bottom: 10px;">{loc["title"]}</h1>'
        f'<p style="color: #5f6c7d; font-size: 1.1rem;">{loc["subtitle"]}</p>'
        f'</div>'
    )
    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(
        '''
        <div style="text-align: center; margin-bottom: 20px;">
            <h1 style="color: #2e4057; font-weight: 800; font-size: 2.5rem; margin-bottom: 10px;">
                Motif-A1 Interactive Web Demo
            </h1>
            <p style="color: #5f6c7d; font-size: 1.1rem;">
                Explore the experimental 105.41M reasoning language model from the CodonProject. Built-in with native Chain-of-Thought (CoT) support.
            </p>
        </div>
        '''
    )

    # 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()