Update app.py

app.py

@@ -519,6 +519,7 @@ logging.basicConfig(level=logging.INFO,
                     ])
 logger = logging.getLogger(__name__)
 
+# Main function to handle the translation workflow
 # Main function to handle the translation workflow
 def main(dataset_url, model_type, output_dataset_name, range_specification, token: gr.OAuthToken | None, profile: gr.OAuthProfile | None):
     try:
@@ -527,24 +528,24 @@ def main(dataset_url, model_type, output_dataset_name, range_specification, toke
             return "### You must be logged in to use this service."
         
         if token:
-            logging.info("Logged in to Hugging Face")
+            logger.info("Logged in to Hugging Face")
 
             # Configuration and paths
             tokenizer_name = "facebook/wmt21-dense-24-wide-en-x"
             model_repo_name = "cstr/wmt21ct2_int8"  # Repository to download the model from
 
             # Download the model snapshot from Hugging Face
-            model_path = snapshot_download(repo_id=model_repo_name, token=token)
-            logging.info(f"Model downloaded to: {model_path}")
+            model_path = snapshot_download(repo_id=model_repo_name, token=token.token)
+            logger.info(f"Model downloaded to: {model_path}")
 
             # Load the CTranslate2 model
             translator = ctranslate2.Translator(model_path, device="auto")
-            logging.info("CTranslate2 model loaded successfully.")
+            logger.info("CTranslate2 model loaded successfully.")
 
             # Load the tokenizer
             tokenizer = transformers.AutoTokenizer.from_pretrained(tokenizer_name)
             tokenizer.src_lang = "en"
-            logging.info("Tokenizer loaded successfully.")
+            logger.info("Tokenizer loaded successfully.")
 
             # Define the task based on user input
             task = {
@@ -566,17 +567,19 @@ def main(dataset_url, model_type, output_dataset_name, range_specification, toke
                 output_dir=".",
                 output_repo_name=output_dataset_name,
                 raw_file_path=task["raw_file"],
-                token=token,
+                token=token.token,
                 range_specification=task["range_spec"],
                 model_type=task["model_type"],
                 translator=translator,
                 tokenizer=tokenizer,
             )
-            return "Dataset translation completed!"
+            logger.info("Dataset translation completed!")
+            return "Dataset translation completed!\n\n### Logs:\n" + log_stream.getvalue()
         else:
             return "Login failed. Please try again."
     except Exception as e:
-        logging.error(f"An error occurred in the main function: {e}")
+        logger.error(f"An error occurred in the main function: {e}")
+        # Ensure logs are flushed and captured
         return f"An error occurred: {e}\n\n### Logs:\n{log_stream.getvalue()}"
 
 # Gradio interface setup
@@ -622,10 +625,10 @@ with gr.Blocks(theme=theme) as demo:
             range_specification = gr.Textbox(label="Range Specification", lines=1, placeholder="e.g., 1-100")
 
         with gr.Column():
-            output = gr.Textbox(label="Output", lines=1)
+            output = gr.Markdown(label="Output")
 
     submit_btn = gr.Button("Translate Dataset", variant="primary")
-    submit_btn.click(main, inputs=[dataset_url, model_type, output_dataset_name, range_specification], outputs=output)
+    submit_btn.click(main, inputs=[dataset_url, model_type, output_dataset_name, range_specification, token_input], outputs=output)
 
     gr.Markdown(datasets_desc)