# Import our modules 
import gradio as gr
from turjuman import turjuman 
import logging
import os
from transformers import AutoTokenizer


logging.basicConfig(
    format="%(asctime)s | %(levelname)s | %(name)s | %(message)s",
    datefmt="%Y-%m-%d %H:%M:%S",
    level=os.environ.get("LOGLEVEL", "INFO").upper(),
)
logger = logging.getLogger("turjuman.translate")
cache_dir="./mycache"


# Get the turjuman object and its tokenizer
turj = turjuman.turjuman(logger, cache_dir)
tokenizer = AutoTokenizer.from_pretrained('UBC-NLP/AraT5-base-title-generation')


# The translate function
def translate(sent):
    beam_options = {"search_method":"beam", "seq_length": 300, "num_beams":5, "no_repeat_ngram_size":2, "max_outputs":1}
    targets = turj.translate(sent,**beam_options)
    #print(targets)
    ans = ""
    for target in targets:
        target = tokenizer.decode(target, skip_special_tokens=True, clean_up_tokenization_spaces=True)
        ans += target
    return ans

#print(translate('Здравствуй, друг'))

interface = gr.Interface(fn=translate, inputs=['text'], outputs=['text'],
             allow_flagging=False,
             title='Turjuman Multi-lingual Translation',
             description='Please write the sentence you are trying to translate',
             css="""
             .chatbox{display:flex; flex-direction-column}
             .user_msg, resp_msg {padding: 4px;margin-bottom:4px;border-radius:4px; width:80%}
             .user_msg {background-color:cornflowerblue; color:white; align-self:start}
             .resp_msg {background-color:lightgray; align-self: self-end}
             """)

interface.launch(inline=False)