ajustes download dataset e model cpu

app.py

@@ -15,7 +15,7 @@ def carregar_modelo_e_tokenizador_mbart(modelo):
 
     st.write(f'Carregando modelo {modelo}')
     tokenizer = MBart50TokenizerFast.from_pretrained(modelo)
-    model = MBartForConditionalGeneration.from_pretrained(modelo).to("cuda")
+    model = MBartForConditionalGeneration.from_pretrained(modelo)
 
     tokenizer.src_lang = 'en_XX'
     return model, tokenizer
@@ -23,8 +23,8 @@ def carregar_modelo_e_tokenizador_mbart(modelo):
 # 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')
+    input_ids = inputs.input_ids
+    attention_mask = inputs.attention_mask
     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)
 
@@ -62,5 +62,6 @@ if st.button(f"Traduzir {qtde_linhas_traduzir} linhas"):
         # adiciona traducao em nova coluna dataset
         dataset["traduzido"]= texto_traduzido
 
-
-st.write("Fim 👍")    
+st.write(dataset)
+st.download_button(label='Baixar dataset', data=dataset.to_csv(index=False))
+st.write("Fim 👍")