import gradio as gr
import torch

from translate import Translator

# https://medium.com/analytics-vidhya/make-a-translate-app-with-huggingface-transformers-ce9203f84c79
# https://huggingface.co/docs/transformers/en/model_doc/mbart
title = "Translation Chatbot"
description = "A simple implementation of translating one language to another"
examples = [["UN Chief Says There Is No Military Solution in Syria","en_XX","ja_XX"]]

translator_obj = Translator()

def translate_sentence(sentence):
    return pipe(f'<-ja2zh-> {sentence}')[0]['translation_text']

def predict(input, 
    history=[],
    original_language="en_XX", 
    translated_language="ro_RO"):
    response = translator_obj.translate(input, original_language, translated_language)
    history.append((input, response))
    return history, history

if __name__ == "__main__":
    gr.Interface(
        fn=predict,
        title=title,
        description=description,
        examples=examples,
        inputs=[
        gr.Textbox(),
        "state",
        gr.Dropdown(
            [("English","en_XX"), ("French","fr_XX"), ("German","de_DE"), ("Japanese","ja_XX"), ("Russian","ru_RU")], value="en_XX", label="Input Language", info="Choose the language the input text is in."
        ),
        gr.Dropdown(
            [("French","fr_XX"), ("German","de_DE"), ("Japanese","ja_XX"), ("Russian","ru_RU"), ("English","en_XX")], value="fr_XX", label="Output Language", info="Choose the language to convert the text to."
        )
        ],
        outputs=[
        gr.Chatbot(),
        "state"
        ],
        theme='earneleh/paris',
    ).launch()