import gradio as gr
from transformers import AutoProcessor, SeamlessM4Tv2Model
import torch

# 加载模型和处理器
processor = AutoProcessor.from_pretrained("facebook/seamless-m4t-v2-large")
model = SeamlessM4Tv2Model.from_pretrained("facebook/seamless-m4t-v2-large")




def greet(text,src_lang,tgt_lang):
    # 输入文本
    text_inputs = processor(text=text, src_lang=src_lang, return_tensors="pt")
    # 生成翻译的文本
    output_tokens = model.generate(**text_inputs, tgt_lang=tgt_lang, generate_speech=False) 
    translated_text_from_text = processor.decode(output_tokens[0].tolist()[0], skip_special_tokens=True)
    return translated_text_from_text

text = gr.Textbox(lines=5,label="输入文本",placeholder="请输入文本")
src_lang = gr.Dropdown(label="源语言",choices=["en","zh"])
tgt_lang = gr.Dropdown(label="目标语言",choices=["en","zh"])
output_text = gr.Textbox(lines=5,label="翻译文本",placeholder="翻译文本")
demo=gr.Interface(fn=greet, inputs=[text,src_lang,tgt_lang], outputs=output_text)
demo.launch()