first commit

app.py

@@ -0,0 +1,66 @@
+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 👍")    

requirements.txt

@@ -0,0 +1,4 @@
+transformers
+setencepiece
+sacremoses
+pandas