make style

huggingface · Dec 8, 2023 · 43d3565 · 43d3565
1 parent f06236b
commit 43d3565
Show file tree

Hide file tree

Showing 7 changed files with 55 additions and 17 deletions.
diff --git a/src/autotrain/trainers/clm/__main__.py b/src/autotrain/trainers/clm/__main__.py
@@ -499,7 +499,13 @@ def train(config):
             logger.info("Pausing space...")
             api = HfApi(token=config.token)
             api.pause_space(repo_id=os.environ["SPACE_ID"])
-
+            success_message = f"Your training run was successful! [Check out your trained model here](https://huggingface.co/{config.username}/{config.project-name})"
+            api.create_discussion(
+                repo_id=os.environ["SPACE_ID"],
+                title="Your training has finished successfully ✅",
+                description=success_message,
+                repo_type="space",
+            )
         if "ENDPOINT_ID" in os.environ:
             # shut down the endpoint
             logger.info("Pausing endpoint...")

diff --git a/src/autotrain/trainers/dreambooth/__main__.py b/src/autotrain/trainers/dreambooth/__main__.py
@@ -336,6 +336,13 @@ def load_model_hook(models, input_dir):
         logger.info("Pausing space...")
         api = HfApi(token=config.token)
         api.pause_space(repo_id=os.environ["SPACE_ID"])
+        success_message = f"Your training run was successful! [Check out your trained model here](https://huggingface.co/{config.username}/{config.project-name})"
+        api.create_discussion(
+            repo_id=os.environ["SPACE_ID"],
+            title="Your training has finished successfully ✅",
+            description=success_message,
+            repo_type="space",
+        )
 
 
 if __name__ == "__main__":

diff --git a/src/autotrain/trainers/generic/__main__.py b/src/autotrain/trainers/generic/__main__.py
@@ -42,12 +42,13 @@ def run(config):
         logger.info("Pausing space...")
         api = HfApi(token=config.token)
         api.pause_space(repo_id=os.environ["SPACE_ID"])
-        success_message = f"Your training run was successfull! [Check out your trained model here](https://huggingface.co/{config.username}/{config.project-name})"
-        api.create_discussion(repo_id=os.environ['SPACE_ID'],
-                              title="Your training has finished successfully ✅",
-                              description=success_message,
-                              repo_type="space"
-                             )
+        success_message = f"Your training run was succesful! [Check out your trained model here](https://huggingface.co/{config.username}/{config.project-name})"
+        api.create_discussion(
+            repo_id=os.environ["SPACE_ID"],
+            title="Your training has finished successfully ✅",
+            description=success_message,
+            repo_type="space",
+        )
     if "ENDPOINT_ID" in os.environ:
         # shut down the endpoint
         logger.info("Pausing endpoint...")

diff --git a/src/autotrain/trainers/seq2seq/__main__.py b/src/autotrain/trainers/seq2seq/__main__.py
@@ -243,7 +243,13 @@ def train(config):
             logger.info("Pausing space...")
             api = HfApi(token=config.token)
             api.pause_space(repo_id=os.environ["SPACE_ID"])
-
+            success_message = f"Your training run was successful! [Check out your trained model here](https://huggingface.co/{config.username}/{config.project-name})"
+            api.create_discussion(
+                repo_id=os.environ["SPACE_ID"],
+                title="Your training has finished successfully ✅",
+                description=success_message,
+                repo_type="space",
+            )
         if "ENDPOINT_ID" in os.environ:
             # shut down the endpoint
             logger.info("Pausing endpoint...")

diff --git a/src/autotrain/trainers/tabular/__main__.py b/src/autotrain/trainers/tabular/__main__.py
@@ -334,7 +334,13 @@ def train(config):
         logger.info("Pausing space...")
         api = HfApi(token=config.token)
         api.pause_space(repo_id=os.environ["SPACE_ID"])
-
+        success_message = f"Your training run was successful! [Check out your trained model here](https://huggingface.co/{config.username}/{config.project-name})"
+        api.create_discussion(
+            repo_id=os.environ["SPACE_ID"],
+            title="Your training has finished successfully ✅",
+            description=success_message,
+            repo_type="space",
+        )
     if "ENDPOINT_ID" in os.environ:
         # shut down the endpoint
         logger.info("Pausing endpoint...")

diff --git a/src/autotrain/trainers/text_classification/__main__.py b/src/autotrain/trainers/text_classification/__main__.py
@@ -192,15 +192,25 @@ def train(config):
             logger.info("Pushing model to hub...")
             api = HfApi(token=config.token)
             api.create_repo(repo_id=config.repo_id, repo_type="model", private=True)
-            api.upload_folder(folder_path=config.project_name, repo_id=config.repo_id, repo_type="model")
+            api.upload_folder(
+                folder_path=config.project_name,
+                repo_id=config.repo_id,
+                repo_type="model",
+            )
 
     if PartialState().process_index == 0:
         if "SPACE_ID" in os.environ:
             # shut down the space
             logger.info("Pausing space...")
             api = HfApi(token=config.token)
             api.pause_space(repo_id=os.environ["SPACE_ID"])
-
+            success_message = f"Your training run was successful! [Check out your trained model here](https://huggingface.co/{config.username}/{config.project-name})"
+            api.create_discussion(
+                repo_id=os.environ["SPACE_ID"],
+                title="Your training has finished successfully ✅",
+                description=success_message,
+                repo_type="space",
+            )
         if "ENDPOINT_ID" in os.environ:
             # shut down the endpoint
             logger.info("Pausing endpoint...")

diff --git a/src/autotrain/utils.py b/src/autotrain/utils.py
@@ -280,16 +280,18 @@ def wrapper(*args, **kwargs):
             return func(*args, **kwargs)
         except Exception:
             if PartialState().process_index == 0:
-                error_message = f'''{func.__name__} has failed due to an exception: {traceback.format_exc()}'''
+                error_message = f"""{func.__name__} has failed due to an exception: {traceback.format_exc()}"""
                 logger.error(error_message)
                 if "SPACE_ID" in os.environ:
                     # shut down the space
                     logger.info("Pausing space...")
                     api = HfApi(token=os.environ["HF_TOKEN"])
                     api.pause_space(repo_id=os.environ["SPACE_ID"])
-                    api.create_discussion(repo_id=os.environ['SPACE_ID'],
-                                          title="Your training has failed ❌",
-                                          description=error_message,
-                                          repo_type="space"
-                                         )
+                    api.create_discussion(
+                        repo_id=os.environ["SPACE_ID"],
+                        title="Your training has failed ❌",
+                        description=error_message,
+                        repo_type="space",
+                    )
+
     return wrapper