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,30 @@ 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 validate_hf_token(
+        token: str | None,
+        *,
+        invalid_message: str = "[red]Invalid token or authentication failed.[/]",
+    ) -> tuple[dict | None, str | None]:
+        """Validate HF token and return (user, token). Returns (None, None) on failure."""
+        if not token:
+            return None, None
+        try:
+            user = huggingface_hub.whoami(token)
+            print_hf_user_info(user)
+            return user, token
+        except Exception:
+            print(invalid_message)
+            return None, None
+
     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 +796,38 @@ 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 = None
+                            if hf_token:
+                                user, hf_token = validate_hf_token(
+                                    hf_token,
+                                    invalid_message="[red]Failed to validate the existing Hugging Face token. It might be expired or invalid.[/]",
+                                )
 
-                            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})[/]")
+                            if user:
+                                try:
+                                    choice = prompt_select(
+                                        "How do you want to proceed?",
+                                        [
+                                            "Use this account",
+                                            "Switch account",
+                                        ],
+                                    )
+
+                                    if choice == "Switch account":
+                                        user = None
+                                        hf_token = None
+                                except KeyboardInterrupt:
+                                    break
+
+                            while not user:
+                                hf_token = prompt_password("Hugging Face access token:")
+                                if not hf_token:
+                                    break
+
+                                user, hf_token = validate_hf_token(hf_token)
+
+                            if not user:
+                                continue
 
                             repo_id = prompt_text(
                                 "Name of repository:",
@@ -805,22 +852,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 +904,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}[/].")
 

src/heretic/utils.py

@@ -136,6 +136,7 @@ def prompt_path(message: str) -> str:
 def prompt_password(message: str) -> str:
     if is_notebook():
         print()
+        # getpass doesn't support a default.
         return getpass.getpass(message)
     else:
         return questionary.password(message).ask()