File size: 3,198 Bytes
144ccda
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
64a313a
144ccda
 
 
 
 
911fd62
d9956d6
64a313a
 
d9956d6
 
144ccda
4e3e82b
 
 
 
 
144ccda
4e3e82b
7e7ef19
 
4e3e82b
 
144ccda
911fd62
7e7ef19
c18735b
7e7ef19
 
911fd62
144ccda
c18735b
7e7ef19
 
 
 
 
c18735b
7e7ef19
144ccda
7e7ef19
4e3e82b
7e7ef19
4e3e82b
 
 
7e7ef19
4e3e82b
 
 
 
 
 
 
 
 
 
 
 
e028986
4e3e82b
c18735b
 
8f49090
4e3e82b
8f49090
c034f10
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
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()