import gradio as gr
from lesson1_prompt import get_prompt_template
from lesson2_pipeline import get_chat_pipeline

# 初始化组件
prompt_template = get_prompt_template()
chat = get_chat_pipeline()

# 翻译功能
def translate_text(style, text):
    try:
        # 格式化输入文本
        prompt = f"Translate the text that is delimited by triple backticks into a style that is {style}. text: ```{text}```"
        
        # 调用 Hugging Face pipeline
        response = chat(prompt)  # 直接传入字符串
        return response[0]['generated_text']  # 返回生成的文本
    except Exception as e:
        print(f"Error: {str(e)}")
        return f"发生错误: {str(e)}"

# 构建 Gradio 界面
with gr.Blocks() as demo:
    gr.Markdown("# StyleTrans: 翻译风格转换助手")
    gr.Markdown("使用 Hugging Face 模型实现多样风格翻译。")

    with gr.Row():
        style = gr.Textbox(label="翻译风格", placeholder="如：正式语调、幽默风格等")
        text = gr.Textbox(label="输入文本", placeholder="请输入需要转换风格的内容")

    output = gr.Textbox(label="翻译结果")

    translate_button = gr.Button("翻译")
    translate_button.click(translate_text, inputs=[style, text], outputs=output)

# 启动应用
demo.launch()