import warnings
from cryptography.utils import CryptographyDeprecationWarning

with warnings.catch_warnings():
    warnings.filterwarnings('ignore', category=CryptographyDeprecationWarning)
    import paramiko

import gradio as gr
#from transformers import pipeline
from transformers import PreTrainedTokenizerFast, AutoTokenizer
from transformers import PreTrainedTokenizerFast
from transformer_mt.modeling_transformer import TransfomerEncoderDecoderModel
from transformer_mt_roberta.modeling_transformer_final import TransfomerEncoderDecoderModel as mt_roberta

#translation_pipeline = pipeline('translation_en_to_fr')

# seting up translation transformer into Gradio
#def translator_fn(text_input):
#    results = translation_pipeline(text_input)
#    return results[0]['translation_text']

# def translator_fn_baseline(text_in):
#     source_tokenizer = PreTrainedTokenizerFast.from_pretrained("da_en_output_dir/da_tokenizer")
#     target_tokenizer = PreTrainedTokenizerFast.from_pretrained("da_en_output_dir/en_tokenizer")
#     model = TransfomerEncoderDecoderModel.from_pretrained("da_en_output_dir")
<<<<<<< HEAD
#
=======
# 
>>>>>>> adb80531e202c58b4ab91375bc391ab50bbc882f
#     input_ids = source_tokenizer.encode(text_in, return_tensors="pt")
#     output_ids = model.generate(
#         input_ids,
#         max_length=10,
#         bos_token_id=target_tokenizer.bos_token_id,
#         eos_token_id=target_tokenizer.eos_token_id,
#         pad_token_id=target_tokenizer.pad_token_id,
#      )
<<<<<<< HEAD
#
=======
# 
>>>>>>> adb80531e202c58b4ab91375bc391ab50bbc882f
#     return target_tokenizer.decode(output_ids[0])

def translator_fn_roberta(text_in):
    source_tokenizer_pretrained_roberta = AutoTokenizer.from_pretrained("flax-community/roberta-base-danish")
    target_tokenizer_pretrained_roberta = PreTrainedTokenizerFast.from_pretrained("da_en_output_dir/en_tokenizer")
    model_pretrained_roberta = mt_roberta.from_pretrained("da_en_RoBERTa_pretrained")

    input_ids_pretrained_roberta = source_tokenizer_pretrained_roberta.encode(text_in, return_tensors="pt")
    output_ids_pretrained_roberta = input_ids_pretrained_roberta.generate(
        input_ids_pretrained_roberta,
        max_length=10,
        bos_token_id=target_tokenizer_pretrained_roberta.bos_token_id,
        eos_token_id=target_tokenizer_pretrained_roberta.eos_token_id,
        pad_token_id=target_tokenizer_pretrained_roberta.pad_token_id,
     )
    return target_tokenizer_pretrained_roberta.decode(output_ids_pretrained_roberta[0])

iface = gr.Interface(fn=translator_fn_roberta,
                     inputs=gr.inputs.Textbox(lines=2, placeholder=None, label="Your Danish text goes here."),
                     outputs=['text'],   # a list should match the number of values returned by fn to have one input and 2 putputs.
                     description = "This App translates text from Danish to the English language.",
                     title = "Danish to English Translator App",
                     theme = "peach")

iface.launch(share=False, enable_queue=True)