Spaces:
Runtime error
Runtime error
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() |