fix: persist hf token between uploads

src/heretic/main.py

@@ -10,6 +10,7 @@
 from importlib.metadata import version
 from os.path import commonprefix
 from pathlib import Path
+from typing import Any
 
 import huggingface_hub
 import optuna
@@ -171,6 +172,12 @@ def run():
         )
         return
 
+    # Keep Hugging Face credentials in memory for this process only.
+    # We don't use huggingface_hub.login() because that stores the token on disk.
+    # Since this program will often be run on rented or shared GPU servers,
+    # it is better to not persist credentials.
+    hf_token = huggingface_hub.get_token()
+
     # Adapted from https://github.com/huggingface/accelerate/blob/main/src/accelerate/commands/env.py
     if torch.cuda.is_available():
         count = torch.cuda.device_count()
@@ -598,6 +605,43 @@ def count_completed_trials() -> int:
     if count_completed_trials() == settings.n_trials:
         study.set_user_attr("finished", True)
 
+    def print_hf_user_info(user: dict[str, Any]):
+        fullname = user.get(
+            "fullname",
+            user.get("name", "unknown user"),
+        )
+        email = user.get("email", "no email found")
+        print(f"Logged in as [bold]{fullname} ({email})[/]")
+
+    def authenticate_hf(token: str | None) -> tuple[dict[str, Any], str]:
+        # If there is already an existing HF token,
+        # (either externally through hf auth login or a previous upload),
+        # allow the user to switch accounts if they want
+        if token:
+            try:
+                user = huggingface_hub.whoami(token)
+                print_hf_user_info(user)
+                if (
+                    prompt_select(
+                        "How do you want to proceed?",
+                        ["Use this account", "Switch account"],
+                    )
+                    == "Use this account"
+                ):
+                    return user, token
+            except Exception as error:
+                print(f"[red]Failed to validate the Hugging Face token: ({error})[/]")
+
+        # Otherwise, keep prompting the user for a valid token until they cancel.
+        while True:
+            token = prompt_password("Hugging Face access token:")
+            try:
+                user = huggingface_hub.whoami(token)
+                print_hf_user_info(user)
+                return user, token
+            except Exception as error:
+                print(f"[red]Failed to validate the Hugging Face token: ({error})[/]")
+
     while True:
         # If no trials at all have been evaluated, the study must have been stopped
         # by pressing Ctrl+C while the first trial was running. In this case, we just
@@ -765,22 +809,10 @@ def count_completed_trials() -> int:
                             print(f"Model saved to [bold]{save_directory}[/].")
 
                         case "Upload the model to Hugging Face":
-                            # We don't use huggingface_hub.login() because that stores the token on disk,
-                            # and since this program will often be run on rented or shared GPU servers,
-                            # it's better to not persist credentials.
-                            token = huggingface_hub.get_token()
-                            if not token:
-                                token = prompt_password("Hugging Face access token:")
-                            if not token:
-                                continue
-
-                            user = huggingface_hub.whoami(token)
-                            fullname = user.get(
-                                "fullname",
-                                user.get("name", "unknown user"),
-                            )
-                            email = user.get("email", "no email found")
-                            print(f"Logged in as [bold]{fullname} ({email})[/]")
+                            try:
+                                user, hf_token = authenticate_hf(hf_token)
+                            except KeyboardInterrupt:
+                                break
 
                             repo_id = prompt_text(
                                 "Name of repository:",
@@ -805,22 +837,22 @@ def count_completed_trials() -> int:
                                 model.model.push_to_hub(
                                     repo_id,
                                     private=private,
-                                    token=token,
+                                    token=hf_token,
                                 )
                             else:
                                 print("Uploading merged model...")
                                 merged_model = model.get_merged_model()
                                 merged_model.push_to_hub(
                                     repo_id,
                                     private=private,
-                                    token=token,
+                                    token=hf_token,
                                 )
                                 del merged_model
                                 empty_cache()
                                 model.tokenizer.push_to_hub(
                                     repo_id,
                                     private=private,
-                                    token=token,
+                                    token=hf_token,
                                 )
 
                             # If the model path exists locally and includes the
@@ -857,7 +889,7 @@ def count_completed_trials() -> int:
                                     )
                                     + card.text
                                 )
-                                card.push_to_hub(repo_id, token=token)
+                                card.push_to_hub(repo_id, token=hf_token)
 
                             print(f"Model uploaded to [bold]{repo_id}[/].")