ajustes fluxo

app.py

@@ -28,33 +28,27 @@ def traduzir_en_pt(text):
     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)
 
+## streamlit ##
+def carregar_dataset():
+    dataset = st.file_uploader("(coluna a ser traduzida deve ser nomeada como 'texto')",type="csv")
+    if dataset is not None:
+        st.write('Carregando dataset...')
+        return pd.read_csv(dataset)
 
-###################
-#### interface ####
-###################
-
-# Cabeçalho
-st.title('Tradutor de datasets (inglês para português)')
-# Carrega dataset
-
-with st.beta_container():
-    dataset = st.file_uploader("Carrege o dataset (coluna a ser traduzida deve ser nomeada como 'texto')", type=["csv"])
-    st.write('Carregando dataset...')
-
-with st.beta_expander():
+def mostrar_dataset():
     if dataset is not None:
         st.write('🎲 Dataset carregado com sucesso!')
-        dataset = pd.read_csv(dataset)
-        st.write(dataset)
-    
+        with st.expander("Dataset:"):
+            st.table(dataset)
 
-with st.beta_container():
+def carregar_modelo():
     modelo_selecionado = st.selectbox('Escolha um modelo', modelos_opcao)
     if st.button("Carregar modelo escolhido"):
         model, tokenizer = carregar_modelo_e_tokenizador_mbart(modelo_selecionado)
         st.write(f"🎰 Modelo {modelo_selecionado} carregado com sucesso!  🔥")
+    return model, tokenizer
 
-with st.beta_container():
+def traduzir_dataset(dataset):
     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):
@@ -64,11 +58,28 @@ with st.beta_container():
             st.write(f'Tradução: {texto_traduzido}')
 
             # adiciona traducao em nova coluna dataset
-            dataset["traduzido"]= texto_traduzido
+            dataset["traduzido"][i]= texto_traduzido
         st.write("Traduzido 👍")
-with st.beta_expander():
-    st.write(dataset)
-
-with st.beta_expander():
+def resultado():
+    with st.expander():
+        st.write(dataset)
     st.download_button(label='Baixar dataset', data=dataset.to_csv(index=False))
-    st.write("Fim 👍")
+    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"):
+        model, tokenizer = carregar_modelo()
+
+
+if st.button("Traduzir dataset") and model is not None:
+    traduzir_dataset(dataset)
+    resultado()