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

import gradio as gr

from transformers import AutoTokenizer, AutoModelForSeq2SeqLM

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 = df.loc[i, 'English']
        df.loc[i, 'Dutch'] = en_nl_translator(input_string)

    return df

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

batch_translate = gr.Interface(
    fn = translate,
    inputs = input_df,  #[input1, input2]
    outputs = "dataframe",
    description = "Fill in with text you want to translate"
)

batch_translate.launch()

"""
def translate(a):
    
    return a

demo = gr.Interface(
    fn = some_func,
    [
        gr.Dataframe(
            headers=["English", "Dutch"],
            datatype=["str", "str"],
            row_count=10,
            col_count=(2, "fixed"),
        ),
        gr.Dropdown(["M", "F", "O"]),
    ],
    "dataframe",
    description="Add the English text you want to translate",
)

if __name__ == "__main__":
    gr.Interface.load("models/Helsinki-NLP/opus-mt-en-nl").launch()
    #demo.launch()
"""