import gradio as gr
from langchain_ollama import ChatOllama
from langchain_core.messages import HumanMessage, SystemMessage

# 1. 初始化 Ollama 模型
# 确保 Ollama 已启动且模型已拉取：ollama pull maternion/hy-mt2:1.8b
llm = ChatOllama(model="maternion/hy-mt2:1.8b")

def translate_text(text, source_lang, target_lang):
    if not text:
        return ""

    # 构建翻译提示词
    system_prompt = f"你是一个专业的翻译助手。请将以下文本从{source_lang}翻译成{target_lang}。"

    messages = [
        SystemMessage(content=system_prompt),
        HumanMessage(content=text)
    ]

    # 调用模型
    response = llm.invoke(messages)
    return response.content

# 2. 构建 Gradio 界面
languages = ["中文", "英语", "日语", "韩语", "法语", "德语"]

with gr.Blocks() as demo:
    gr.Markdown("## 腾讯混元翻译模型 (Hy-MT2)")
    with gr.Row():
        source_lang = gr.Dropdown(languages, label="源语言", value="韩语")
        target_lang = gr.Dropdown(languages, label="目标语言", value="中文")

    with gr.Row():
        input_text = gr.Textbox(label="输入文本", lines=5)
        output_text = gr.Textbox(label="翻译结果", lines=5, interactive=False)

    btn = gr.Button("翻译")
    btn.click(fn=translate_text, inputs=[input_text, source_lang, target_lang], outputs=output_text)

demo.launch()

import gradio as gr
from langchain_ollama import ChatOllama
from langchain_core.messages import HumanMessage, SystemMessage

# 1. 初始化 Ollama 模型
# 确保 Ollama 已启动且模型已拉取：ollama pull maternion/hy-mt2:1.8b
llm = ChatOllama(model="maternion/hy-mt2:1.8b")


def translate_text(text, source_lang, target_lang):
    if not text:
        return ""

    # 构建翻译提示词
    system_prompt = f"你是一个专业的翻译助手。请将以下文本从{source_lang}翻译成{target_lang}。"

    messages = [
        SystemMessage(content=system_prompt),
        HumanMessage(content=text)
    ]

    # 调用模型
    response = llm.invoke(messages)
    return response.content


if __name__ == "__main__":
    # 2. 构建 Gradio 界面
    languages = ["中文", "英语", "日语", "韩语", "法语", "德语"]

    with gr.Blocks() as demo:
        gr.Markdown("## 腾讯混元翻译模型 (Hy-MT2)")
        with gr.Row():
            source_lang = gr.Dropdown(languages, label="源语言", value="韩语")
            target_lang = gr.Dropdown(languages, label="目标语言", value="中文")

        with gr.Row():
            input_text = gr.TextArea(label="输入文本", lines=7, max_lines=20)
            output_text = gr.TextArea(label="翻译结果", lines=7, max_lines=20, interactive=False)

        btn = gr.Button("翻译")
        btn.click(fn=translate_text, inputs=[input_text, source_lang, target_lang], outputs=output_text)

    demo.launch()