fix: persist hf token between uploads

src/heretic/main.py

@@ -171,6 +171,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 +604,14 @@ 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):
+        fullname = user.get(
+            "fullname",
+            user.get("name", "unknown user"),
+        )
+        email = user.get("email", "no email found")
+        print(f"Logged in as [bold]{fullname} ({email})[/]")
+
     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 +779,50 @@ 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
 
-                            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 hf_token:
+                                try:
+                                    user = huggingface_hub.whoami(hf_token)
+                                except Exception:
+                                    print(
+                                        "[yellow]Warning: Failed to validate the existing Hugging Face token. It might be expired or invalid.[/]"
+                                    )
+
+                            if user:
+                                _print_hf_user_info(user)
+
+                                choice = prompt_select(
+                                    "How do you want to proceed?",
+                                    [
+                                        "Use this account",
+                                        "Switch account",
+                                    ],
+                                )
+
+                                if choice == "Switch account":
+                                    user = None
+                                    hf_token = None
+
+                            while not user:
+                                hf_token = prompt_password(
+                                    "Hugging Face access token:", None
+                                )
+                                if not hf_token:
+                                    break
+
+                                try:
+                                    user = huggingface_hub.whoami(hf_token)
+                                    _print_hf_user_info(user)
+
+                                except Exception:
+                                    print(
+                                        "[red]Invalid token or authentication failed.[/]"
+                                    )
+                                    hf_token = None
+
+                            if not user:
+                                continue
 
                             repo_id = prompt_text(
                                 "Name of repository:",
@@ -805,23 +847,24 @@ 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,
-                                )
+
+                            model.tokenizer.push_to_hub(
+                                repo_id,
+                                private=private,
+                                token=hf_token,
+                            )
 
                             # If the model path exists locally and includes the
                             # card, use it directly. If the model path doesn't
@@ -857,7 +900,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

@@ -133,12 +133,17 @@ def prompt_path(message: str) -> str:
         return questionary.path(message, only_directories=True).ask()
 
 
-def prompt_password(message: str) -> str:
+def prompt_password(message: str, hf_token: str | None) -> str:
     if is_notebook():
         print()
+        # getpass doesn't support a default.
         return getpass.getpass(message)
     else:
-        return questionary.password(message).ask()
+        if hf_token:
+            # Set the existing token, but allow user to switch if desired.
+            return questionary.password(message, default=hf_token).ask()
+        else:
+            return questionary.password(message).ask()
 
 
 def format_duration(seconds: float) -> str: