Spaces:
Runtime error
Runtime error
File size: 2,386 Bytes
144ccda |
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 |
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).to("cuda")
tokenizer.src_lang = 'en_XX'
return model, tokenizer
# TODO:batch?
def traduzir_en_pt(text):
inputs = tokenizer(text, return_tensors='pt')
input_ids = inputs.input_ids.to('cuda')
attention_mask = inputs.attention_mask.to('cuda')
output = model.generate(input_ids, attention_mask=attention_mask, forced_bos_token_id=tokenizer.lang_code_to_id['pt_XX'])
return tokenizer.decode(output[0], skip_special_tokens=True)
###################
#### interface ####
###################
# Cabeçalho
st.title('Tradutor de datasets (inglês para português)')
# Carrega dataset
dataset = st.file_uploader("Carrege o dataset (coluna a ser traduzida deve ser nomeada como 'texto')", type=["csv"])
st.write('Carregando dataset...')
if dataset is not None:
st.write('🎲 Dataset carregado com sucesso!')
dataset = pd.read_csv(dataset)
st.write(dataset)
modelo_selecionado = st.selectbox('Escolha um modelo', modelos_opcao)
if st.button("Carregar modelo escolhido"):
tokenizer, model = carregar_modelo_e_tokenizador_mbart(modelo_selecionado)
st.write(f"🎰 Modelo {modelo_selecionado} carregado com sucesso! 🔥")
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"])
st.write(f'Tradução: {texto_traduzido}')
# adiciona traducao em nova coluna dataset
dataset["traduzido"]= texto_traduzido
st.write("Fim 👍")
|