import gradio as gr
from transformers import MarianMTModel, MarianTokenizer

model_name_en_to_ja = "Helsinki-NLP/opus-mt-en-jap"
model_en_to_ja = MarianMTModel.from_pretrained(model_name_en_to_ja)
tokenizer_en_to_ja = MarianTokenizer.from_pretrained(model_name_en_to_ja)

model_name_ja_to_en = "Helsinki-NLP/opus-mt-jap-en"
model_ja_to_en = MarianMTModel.from_pretrained(model_name_ja_to_en)
tokenizer_ja_to_en = MarianTokenizer.from_pretrained(model_name_ja_to_en)

# 日本語と英語の判定はできるが、翻訳精度は高くないです。
# 日本語と英語を両方入力した場合も挙動がおかしいです。
def translate(text):
    # 文字列がアルファベットは判定
    if text.strip()[0].isalpha(): 
        inputs = tokenizer_en_to_ja(text, return_tensors="pt", padding=True)
        translated = model_en_to_ja.generate(**inputs)
        translated_text = tokenizer_en_to_ja.batch_decode(translated, skip_special_tokens=True)
    else:
        inputs = tokenizer_ja_to_en(text, return_tensors="pt", padding=True)
        translated = model_ja_to_en.generate(**inputs)
        translated_text = tokenizer_ja_to_en.batch_decode(translated, skip_special_tokens=True)
    return translated_text[0]

gr.Interface(fn=translate, inputs="textbox", outputs="textbox", title="日⇔英 相互翻訳").launch()