import gradio as gr
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, T5Tokenizer, MT5ForConditionalGeneration

tokenizer = T5Tokenizer.from_pretrained("engmatic-earth/mt5-zh-ja-en-trimmed-fine-tuned-v1")
model = AutoModelForSeq2SeqLM.from_pretrained("engmatic-earth/mt5-zh-ja-en-trimmed-fine-tuned-v1")

def output(raw_input_sentence):
    target_sentence = ["en2ja: " + str(raw_input_sentence)]
    translated = model.generate(**tokenizer(target_sentence, return_tensors="pt", padding=True))
    tgt_text = [tokenizer.decode(t, skip_special_tokens=True) for t in translated]
    translated_phrase = tgt_text[0]
    translated_phrase = translated_phrase.replace(" ", "")
    return translated_phrase
    
interface = gr.Interface(fn=output, inputs=gr.inputs.Textbox(lines=3, placeholder="Write what you want to say in Japanese.")
                        ,outputs='text')
interface.launch()