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=['afr', 'amh', 'arb', 'ary', 'arz', 'asm', 'ast', 'azj', 'bel', 'ben', 'bos', 'bul', 'cat', 'ceb', 'ces', 'ckb', 'cmn', 'cmn_Hant', 'cym', 'dan', 'deu', 'ell', 'eng', 'est', 'eus', 'fin', 'fra', 'fuv', 'gaz', 'gle', 'glg', 'guj', 'heb', 'hin', 'hrv', 'hun', 'hye', 'ibo', 'ind', 'isl', 'ita', 'jav', 'jpn', 'kam', 'kan', 'kat', 'kaz', 'kea', 'khk', 'khm', 'kir', 'kor', 'lao', 'lit', 'ltz', 'lug', 'luo', 'lvs', 'mai', 'mal', 'mar', 'mkd', 'mlt', 'mni', 'mya', 'nld', 'nno', 'nob', 'npi', 'nya', 'oci', 'ory', 'pan', 'pbt', 'pes', 'pol', 'por', 'ron', 'rus', 'slk', 'slv', 'sna', 'snd', 'som', 'spa', 'srp', 'swe', 'swh', 'tam', 'tel', 'tgk', 'tgl', 'tha', 'tur', 'ukr', 'urd', 'uzn', 'vie', 'xho', 'yor', 'yue', 'zlm', 'zsm', 'zul'])
tgt_lang = gr.Dropdown(label="目标语言",choices=['afr', 'amh', 'arb', 'ary', 'arz', 'asm', 'ast', 'azj', 'bel', 'ben', 'bos', 'bul', 'cat', 'ceb', 'ces', 'ckb', 'cmn', 'cmn_Hant', 'cym', 'dan', 'deu', 'ell', 'eng', 'est', 'eus', 'fin', 'fra', 'fuv', 'gaz', 'gle', 'glg', 'guj', 'heb', 'hin', 'hrv', 'hun', 'hye', 'ibo', 'ind', 'isl', 'ita', 'jav', 'jpn', 'kam', 'kan', 'kat', 'kaz', 'kea', 'khk', 'khm', 'kir', 'kor', 'lao', 'lit', 'ltz', 'lug', 'luo', 'lvs', 'mai', 'mal', 'mar', 'mkd', 'mlt', 'mni', 'mya', 'nld', 'nno', 'nob', 'npi', 'nya', 'oci', 'ory', 'pan', 'pbt', 'pes', 'pol', 'por', 'ron', 'rus', 'slk', 'slv', 'sna', 'snd', 'som', 'spa', 'srp', 'swe', 'swh', 'tam', 'tel', 'tgk', 'tgl', 'tha', 'tur', 'ukr', 'urd', 'uzn', 'vie', 'xho', 'yor', 'yue', 'zlm', 'zsm', 'zul'])
output_text = gr.Textbox(lines=5,label="翻译文本",placeholder="翻译文本")
demo=gr.Interface(fn=greet, inputs=[text,src_lang,tgt_lang], outputs=output_text)
demo.launch(share=True)