ajuste temporario teste

app.py

@@ -21,7 +21,7 @@ def carregar_modelo_e_tokenizador_mbart(modelo):
     return model, tokenizer
 
 # TODO:batch?
-def traduzir_en_pt(text):
+def traduzir_en_pt(text, modelo, tokenizador):
     inputs = tokenizador(text, return_tensors='pt')
     input_ids = inputs.input_ids
     attention_mask = inputs.attention_mask
@@ -41,20 +41,20 @@ def mostrar_dataset():
         with st.expander("Dataset:"):
             st.table(dataset)
 
-def carregar_modelo():
+def carregar_modelo(modelos_opcao):
     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
+    return model, tokenizer
 
-def traduzir_dataset(dataset):
+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"])
+            texto_traduzido= traduzir_en_pt(dataset.iloc[i]["texto"], modelo, tokenizador)
             st.write(f'Tradução: {texto_traduzido}')
 
             # adiciona traducao em nova coluna dataset
@@ -80,8 +80,8 @@ if dataset is not None:
 
 if dataset is not None:
     if st.button("Carregar modelo"):
-        modelo, tokenizador = carregar_modelo()
+        modelo, tokenizador = carregar_modelo(modelos_opcao)
 
 if st.button("Traduzir dataset") and modelo is not None:
-    traduzir_dataset(dataset)
+    traduzir_dataset(dataset, modelo, tokenizador)
     resultado()