import streamlit as st
from transformers import pipeline
import pandas as pd

modelos_opcao =[
    "Narrativa/mbart-large-50-finetuned-opus-en-pt-translation",
    # "unicamp-dl/translation-en-pt-t5" # desempenho inferior ao MBART (porém, mais rápido)
]


# Carrega o modelo
def carregar_modelo_e_tokenizador_mbart(modelo):
    # https://huggingface.co/Narrativa/mbart-large-50-finetuned-opus-en-pt-translation
    from transformers import MBart50TokenizerFast, MBartForConditionalGeneration    

    st.write(f'Carregando modelo {modelo}')
    tokenizer = MBart50TokenizerFast.from_pretrained(modelo)
    model = MBartForConditionalGeneration.from_pretrained(modelo)

    tokenizer.src_lang = 'en_XX'
    return model, tokenizer

# TODO:batch?
def traduzir_en_pt(text, modelo, tokenizador):
    inputs = tokenizador(text, return_tensors='pt')
    input_ids = inputs.input_ids
    attention_mask = inputs.attention_mask
    output = modelo.generate(input_ids, attention_mask=attention_mask, forced_bos_token_id=tokenizador.lang_code_to_id['pt_XX'])
    return tokenizador.decode(output[0], skip_special_tokens=True)

## streamlit ##
def carregar_dataset():
    dataset = st.file_uploader("(coluna a ser traduzida deve ser nomeada como 'texto')",type="csv")
    if dataset is not None:
        return pd.read_csv(dataset)

def mostrar_dataset():
    if dataset is not None:
        st.write('🎲 Dataset carregado com sucesso!')
        with st.expander("Dataset:"):
            st.table(dataset)

def carregar_modelo(modelos_opcao):
    modelo_selecionado = st.selectbox('Escolha um modelo', modelos_opcao)
    if modelo_selecionado is not None:
        model, tokenizer = carregar_modelo_e_tokenizador_mbart(modelo_selecionado)
        st.write(f"🎰 Modelo {modelo_selecionado} carregado com sucesso!  🔥")
    return model, tokenizer

def traduzir_dataset(dataset, _modelo, _tokenizador):
    qtde_linhas_traduzir = st.slider('Quantidade de linhas a serem traduzidas', 1, len(dataset), 50)
    if st.button(f"Traduzir {qtde_linhas_traduzir} linhas"):
        for i in range(qtde_linhas_traduzir):
            st.write(f'🔡  Traduzindo linha {i+1}...')
            st.write(f'Texto: {dataset.iloc[i]["texto"]}')
            texto_traduzido= traduzir_en_pt(dataset.iloc[i]["texto"], _modelo, _tokenizador)
            st.write(f'Tradução: {texto_traduzido}')

            # adiciona traducao em nova coluna dataset
            dataset["traduzido"][i]= texto_traduzido
        st.write("Traduzido 👍")
def resultado():
    with st.expander():
        st.write(dataset)
    st.download_button(label='Baixar dataset', data=dataset.to_csv(index=False))
    st.write("Fim 👍")

###################
#### interface ####
###################

# Cabeçalho
st.title('Tradutor de datasets (inglês para português)')
# Carrega dataset
dataset = carregar_dataset()
if dataset is not None:
    mostrar_dataset()
    

if st.button("Carregar modelo"):
    _modelo, _tokenizador = carregar_modelo(modelos_opcao)
    print(f"{_modelo.name_or_path} carregado")

    if st.button("Traduzir dataset") and _modelo is not None:
        traduzir_dataset(dataset, _modelo, _tokenizador)
        resultado()