gradio_demo.py

import argparse
import os
import socket
import numpy as np
import gradio as gr
from transformers import AutoConfig, AutoTokenizer
from ml_dtypes import bfloat16

from utils.infer_func import InferManager

DEFAULT_LANGUAGES = [
    "English",
    "Chinese",
    "Japanese",
    "Korean",
    "French",
    "German",
    "Spanish",
    "Italian",
    "Portuguese",
    "Russian",
    "Arabic",
    "Hindi",
    "Bengali",
    "Thai",
    "Vietnamese",
    "Indonesian",
    "Turkish",
    "Polish",
    "Dutch",
    "Swedish",
    "Danish",
    "Norwegian",
    "Finnish",
    "Greek",
    "Czech",
    "Hungarian",
    "Romanian",
    "Ukrainian",
    "Malay",
    "Filipino",
    "Urdu",
    "Hebrew",
    "Persian",
]


def build_prompt(source_text: str, target_language: str, use_zh_template: bool) -> str:
    if use_zh_template:
        return (
            f"将以下文本翻译为{target_language}，注意只需要输出翻译后的结果，不要额外解释：\n"
            f"{source_text}"
        )
    return (
        f"Translate the following segment into {target_language}, without additional explanation.\n"
        f"{source_text}"
    )


def _get_ipv4_address() -> str:
    try:
        s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        s.connect(("8.8.8.8", 80))
        ip = s.getsockname()[0]
        s.close()
        return ip
    except Exception:
        return "127.0.0.1"


def create_demo(hf_model: str, axmodel_path: str, max_seq_len=2047):
    embeds_path = os.path.join(axmodel_path, "model.embed_tokens.weight.npy")
    if not os.path.exists(embeds_path):
        raise FileNotFoundError(f"Missing embeddings file: {embeds_path}")

    tokenizer = AutoTokenizer.from_pretrained(hf_model)
    config = AutoConfig.from_pretrained(hf_model, trust_remote_code=True)
    embeds = np.load(embeds_path)

    eos_token_id = None
    if isinstance(config.eos_token_id, list) and len(config.eos_token_id) > 1:
        eos_token_id = config.eos_token_id

    imer = InferManager(config, axmodel_path, max_seq_len=max_seq_len)

    def translate(
        text,
        target_language,
        use_zh_template,
        temperature,
        top_p,
        top_k,
        repetition_penalty,
        max_new_tokens,
    ):
        if not text or not text.strip():
            return ""

        prompt = build_prompt(text.strip(), target_language, use_zh_template)
        messages = [{"role": "user", "content": prompt}]
        input_ids = tokenizer.apply_chat_template(
            messages,
            tokenize=True,
            add_generation_prompt=False,
            return_tensors="pt",
        )
        token_ids = input_ids[0].cpu().numpy().tolist()
        prefill_data = np.take(embeds, token_ids, axis=0).astype(bfloat16)

        slice_len = 128
        token_ids = imer.prefill(
            tokenizer,
            token_ids,
            prefill_data,
            slice_len=slice_len,
            top_k=top_k,
            top_p=top_p,
            temperature=temperature,
            repetition_penalty=repetition_penalty,
        )
        output = imer.decode(
            tokenizer,
            token_ids,
            embeds,
            slice_len=slice_len,
            eos_token_id=eos_token_id,
            stream=False,
            top_k=top_k,
            top_p=top_p,
            temperature=temperature,
            repetition_penalty=repetition_penalty,
            max_new_tokens=max_new_tokens,
        )
        return output.strip()

    def translate_stream(
        text,
        target_language,
        use_zh_template,
        temperature,
        top_p,
        top_k,
        repetition_penalty,
        max_new_tokens,
    ):
        if not text or not text.strip():
            yield ""
            return

        prompt = build_prompt(text.strip(), target_language, use_zh_template)
        messages = [{"role": "user", "content": prompt}]
        input_ids = tokenizer.apply_chat_template(
            messages,
            tokenize=True,
            add_generation_prompt=False,
            return_tensors="pt",
        )
        token_ids = input_ids[0].cpu().numpy().tolist()
        prefill_data = np.take(embeds, token_ids, axis=0).astype(bfloat16)

        slice_len = 128
        token_ids = imer.prefill(
            tokenizer,
            token_ids,
            prefill_data,
            slice_len=slice_len,
            top_k=top_k,
            top_p=top_p,
            temperature=temperature,
            repetition_penalty=repetition_penalty,
        )

        for text_so_far in imer.decode_stream(
            tokenizer,
            token_ids,
            embeds,
            slice_len=slice_len,
            eos_token_id=eos_token_id,
            top_k=top_k,
            top_p=top_p,
            temperature=temperature,
            repetition_penalty=repetition_penalty,
            max_new_tokens=max_new_tokens,
        ):
            yield text_so_far

    with gr.Blocks(title="HY-MT1.5-1.8B_GPTQ_INT4 Multilingual Translation") as demo:
        gr.Markdown("## HY-MT1.5-1.8B_GPTQ_INT4 Multilingual Translation")

        with gr.Group():
            input_text = gr.Textbox(
                label="Input Text",
                placeholder="Please enter the text you want to translate...",
                lines=6,
            )

        with gr.Group():
            with gr.Row(equal_height=True):
                target_language = gr.Dropdown(
                    choices=DEFAULT_LANGUAGES,
                    value="English",
                    label="Target Language",
                )
                use_zh_template = gr.Checkbox(
                    label="Use Chinese Prompt Template",
                    value=False,
                )
        with gr.Group():
            with gr.Row(equal_height=True):
                temperature = gr.Slider(
                    minimum=0.1,
                    maximum=1.5,
                    value=0.7,
                    step=0.05,
                    label="Temperature",
                )
                top_p = gr.Slider(
                    minimum=0.1,
                    maximum=1.0,
                    value=0.6,
                    step=0.05,
                    label="Top-p",
                )
                top_k = gr.Slider(
                    minimum=1,
                    maximum=100,
                    value=20,
                    step=1,
                    label="Top-k",
                )

        with gr.Group():
            with gr.Row(equal_height=True):
                repetition_penalty = gr.Slider(
                    minimum=1.0,
                    maximum=1.5,
                    value=1.05,
                    step=0.01,
                    label="Repetition Penalty",
                )
                max_new_tokens = gr.Slider(
                    minimum=1,
                    maximum=1024,
                    value=512,
                    step=1,
                    label="Max New Tokens",
                )

        translate_btn = gr.Button("Translate", variant="primary")
        output_text = gr.Textbox(
            label="Translation Result",
            lines=6,
            interactive=False,
        )

        translate_btn.click(
            translate_stream,
            inputs=[
                input_text,
                target_language,
                use_zh_template,
                temperature,
                top_p,
                top_k,
                repetition_penalty,
                max_new_tokens,
            ],
            outputs=output_text,
        )

    return demo


def parse_args():
    parser = argparse.ArgumentParser(description="HY-MT1.5-1.8B_GPTQ_INT4 Gradio Demo")
    parser.add_argument(
        "--hf_model",
        type=str,
        default="./hymt1-5_tokenizer",
        help="Path to HuggingFace model",
    )
    parser.add_argument(
        "--axmodel_path",
        type=str,
        default="./hymt1-5_2k_ax650n_axmodel/",
        help="Path to compiled axmodel directory",
    )
    parser.add_argument(
        "--max_seq_len",
        type=int,
        default=2047,
        help="Maximum sequence length for model inference",
    )
    parser.add_argument("--server_name", type=str, default="0.0.0.0")
    parser.add_argument("--server_port", type=int, default=7860)
    return parser.parse_args()


if __name__ == "__main__":
    args = parse_args()
    app = create_demo(args.hf_model, args.axmodel_path, args.max_seq_len)
    ipv4 = _get_ipv4_address()
    print(f"* Running on local URL:  http://{ipv4}:{args.server_port}")
    app.launch(server_name=args.server_name, server_port=args.server_port)