#!pip install transformers # put transformers in the requirements.txt file

import gradio as gr

from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline

en_nl_tokenizer = AutoTokenizer.from_pretrained("Helsinki-NLP/opus-mt-en-nl")
en_nl_model = AutoModelForSeq2SeqLM.from_pretrained("Helsinki-NLP/opus-mt-en-nl")

#translation_pipeline = pipeline("text-classification", model=finbert_sentiment_prosus, tokenizer=finbert_sentiment_prosus_tokenizer)
en_nl_translator = pipeline("translation_en_to_nl", model = en_nl_model, tokenizer=en_nl_tokenizer) #might ignore "translation_en_to_nl"
# 


#gr.Interface.load("models/Helsinki-NLP/opus-mt-en-nl").launch()

def translate(df):
    #translate the input_text
    
    for i in range(len(df['English'])):
        input_string = str(df.loc[i, 'English'])
        if input_string == '':
            pass
        else:
            df.loc[i, 'Dutch'] = en_nl_translator(input_string)[0]['translation_text'] #extract from list, then take the value associated to key 'translation_text'
    return df

input_df = gr.Dataframe(
            headers=["English", "Dutch"],
            datatype=["str", "str"],
            row_count=5,
            col_count=(2, "fixed"),
        )

batch_translate = gr.Interface(
    fn = translate,
    inputs = input_df,  #[input1, input2]
    outputs = "dataframe",
    description = "Hela menneke, fill in the left column with the English text you want to translate"
)

if __name__ == "__main__":
    batch_translate.launch()