app.py

import gradio as gr
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

# 设置设备，如果有GPU则使用GPU
device = "cuda" if torch.cuda.is_available() else "cpu"

# 加载基础模型和分词器
tokenizer_base = AutoTokenizer.from_pretrained("tencent/Hunyuan-MT-7B", trust_remote_code=True)
model_base = AutoModelForCausalLM.from_pretrained("tencent/Hunyuan-MT-7B", trust_remote_code=True).to(device)

# 加载Chimera集成模型和分词器
tokenizer_chimera = AutoTokenizer.from_pretrained("tencent/Hunyuan-MT-Chimera-7B", trust_remote_code=True)
model_chimera = AutoModelForCausalLM.from_pretrained("tencent/Hunyuan-MT-Chimera-7B", trust_remote_code=True).to(device)


def translate(model_choice, text_to_translate, source_lang, target_lang):
    """
    根据选择的模型进行翻译
    """
    if model_choice == "Hunyuan-MT-7B (基础版)":
        tokenizer = tokenizer_base
        model = model_base
        # 基础版Prompt模板
        prompt = f"Translate the following text from {source_lang} to {target_lang}:\n{text_to_translate}"
    else: # Chimera-7B
        tokenizer = tokenizer_chimera
        model = model_chimera
        # Chimera版需要一个特殊的、包含候选翻译的模板，这里我们简化一下，
        # 实际应用中会先用基础模型生成多个候选。为简化Demo，我们直接套用基础模板。
        # 官方的Chimera用法更复杂，需要输入多个候选翻译进行精炼。
        prompt = f"Translate the following text from {source_lang} to {target_lang}:\n{text_to_translate}"

    inputs = tokenizer(prompt, return_tensors="pt").to(device)
    
    # 生成翻译结果
    output = model.generate(**inputs, max_new_tokens=256)
    
    # 解码并清理结果
    response = tokenizer.decode(output[0], skip_special_tokens=True)
    
    # 移除prompt部分，只返回翻译结果
    translated_text = response.replace(prompt, "").strip()
    
    return translated_text

# --- 创建Gradio界面 ---
with gr.Blocks() as demo:
    gr.Markdown("# 腾讯混元翻译模型体验Demo")
    gr.Markdown("选择一个模型，输入源语言、目标语言和待翻译的文本。")
    
    with gr.Row():
        model_selector = gr.Radio(
            ["Hunyuan-MT-7B (基础版)", "Hunyuan-MT-Chimera-7B (集成优化版)"], 
            label="选择模型", 
            value="Hunyuan-MT-7B (基础版)"
        )
    
    with gr.Row():
        source_language = gr.Textbox(label="源语言", value="Chinese")
        target_language = gr.Textbox(label="目标语言", value="English")

    input_text = gr.Textbox(label="输入文本", lines=5, placeholder="在这里输入需要翻译的文本...")
    output_text = gr.Textbox(label="翻译结果", lines=5)
    
    translate_button = gr.Button("开始翻译")
    
    translate_button.click(
        fn=translate,
        inputs=[model_selector, input_text, source_language, target_language],
        outputs=output_text
    )

# 启动应用
demo.launch()