fix: persist hf token between uploads

src/heretic/main.py

@@ -166,6 +166,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,
+    # and since this program will often be run on rented or shared GPU servers,
+    # it's 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()
@@ -761,16 +767,12 @@ 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:
+                            if not hf_token:
+                                hf_token = prompt_password("Hugging Face access token:")
+                            if not hf_token:
                                 continue
 
-                            user = huggingface_hub.whoami(token)
+                            user = huggingface_hub.whoami(hf_token)
                             fullname = user.get(
                                 "fullname",
                                 user.get("name", "unknown user"),
@@ -801,23 +803,23 @@ 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 doesn't exist locally, it can be assumed
@@ -842,7 +844,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}[/].")