import gradio as gr
from transformers import pipeline

# 加载稳定的中英互译模型
translator_zh2en = pipeline("translation", model="Helsinki-NLP/opus-mt-zh-en", device=-1)
translator_en2zh = pipeline("translation", model="Helsinki-NLP/opus-mt-en-zh", device=-1)

# 4种翻译场景
SCENARIOS = {
    "日常对话": "（口语化、简洁自然）",
    "文档翻译": "（正式、准确，保留专业术语）",
    "旅游场景": "（突出地点、时间等关键信息）",
    "商务沟通": "（礼貌、专业的商务语境）"
}

def translate(text, src_lang, tgt_lang, scenario):
    if not text.strip():
        return "请输入内容"
    prompt = f"{SCENARIOS[scenario]}\n{text}"
    if src_lang == "中文" and tgt_lang == "英文":
        return translator_zh2en(prompt)[0]["translation_text"]
    else:
        return translator_en2zh(prompt)[0]["translation_text"]

# 界面
with gr.Blocks(title="多场景翻译助手") as demo:
    gr.Markdown("# 多场景中英翻译工具")
    with gr.Row():
        src = gr.Dropdown(["中文", "英文"], label="源语言", value="中文")
        tgt = gr.Dropdown(["英文", "中文"], label="目标语言", value="英文")
        scene = gr.Dropdown(list(SCENARIOS.keys()), label="翻译场景", value="日常对话")
    input_box = gr.Textbox(label="输入", lines=4)
    output_box = gr.Textbox(label="输出", lines=4)
    gr.Button("翻译").click(translate, [input_box, src, tgt, scene], output_box)

if __name__ == "__main__":
    demo.launch()