import gradio as gr
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM

# 指定用的模型名称
model_checkpoint = "Helsinki-NLP/opus-mt-en-zh"
tokenizer = AutoTokenizer.from_pretrained(model_checkpoint)
model = AutoModelForSeq2SeqLM.from_pretrained(model_checkpoint)

def translate(text):
    # 对输入进行 encode 处理
    inputs = tokenizer.encode(text, return_tensors="pt")

    # 使用模型进行翻译
    outputs = model.generate(inputs, max_length=128, num_beams=4, early_stopping=True)

    # 对输出进行 decode 处理
    output_text = tokenizer.decode(outputs[0], skip_special_tokens=True)

    return output_text

iface = gr.Interface(fn=translate, inputs="text", outputs="text", title="翻译", description="输入英语文本，获取中文翻译。")

iface.launch()