feat: refactor download and upload scripts for improved argument handling (more notebook friendly :D)

scripts/download_checkpoint.py

@@ -1,21 +1,42 @@
 """Download model from HuggingFace Hub.
 This script downloads a model repository from HuggingFace Hub to local directory.
+
+Example:
+    python download_checkpoint.py --repo-id "org/model-name" --local-dir "models"
 """
 
+import argparse
 import os
 
 from dotenv import load_dotenv
 from huggingface_hub import snapshot_download
 
-load_dotenv(override=True)
 
-# Configuration
+def parse_args() -> argparse.Namespace:
-REPO_ID = "janhq/250403-runpod-qwen7b-r1-distil"
+    """Parse command line arguments.
-LOCAL_DIR = "downloaded_model"  # Where to save the model
+
-HF_TOKEN = os.getenv("HF_TOKEN")
+    Returns:
+        argparse.Namespace: Parsed arguments
+    """
+    parser = argparse.ArgumentParser(description="Download model from HuggingFace Hub")
+    parser.add_argument(
+        "--repo-id", type=str, default="janhq/250403-llama-3.2-3b-instruct-grpo", help="HuggingFace repository ID"
+    )
+    parser.add_argument("--local-dir", type=str, default="downloaded_model", help="Local directory to save model")
+
+    return parser.parse_args()
+
 
-# Files to ignore during download
+def main():
-IGNORE_PATTERNS = [
+    """Main function to download model."""
+    args = parse_args()
+    load_dotenv(override=True)
+
+    # Configuration
+    HF_TOKEN = os.getenv("HF_TOKEN")
+
+    # Files to ignore during download
+    IGNORE_PATTERNS = [
         "*.log",  # Log files
         "*.pyc",  # Python cache
         ".git*",  # Git files
@@ -25,13 +46,18 @@ IGNORE_PATTERNS = [
         "events.*",  # Tensorboard
         "wandb/*",  # Weights & Biases
         "runs/*",  # Training runs
-]
+    ]
 
-# Download the model
+    # Download the model
-snapshot_download(
+    snapshot_download(
         token=HF_TOKEN,
-    repo_id=REPO_ID,
+        repo_id=args.repo_id,
-    local_dir=LOCAL_DIR,
+        local_dir=args.local_dir,
-    # ignore_patterns=IGNORE_PATTERNS,
+        repo_type="model",
-)
+        # ignore_patterns=IGNORE_PATTERNS
-print(f"✅ Done: {REPO_ID} -> {LOCAL_DIR}")
+    )
+    print(f"✅ Done: {args.repo_id} -> {args.local_dir}")
+
+
+if __name__ == "__main__":
+    main()

scripts/upload_checkpoint.py

@@ -1,22 +1,40 @@
 """Upload local directory to HuggingFace Hub.
 This script uploads a specified local directory to HuggingFace Hub as a private repository.
-It uses API token from HuggingFace for authentication.
+
+Example:
+    python upload_checkpoint.py --local-dir "models/my_model" --repo-id "org/model-name"
 """
 
+import argparse
 import os
 
 from dotenv import load_dotenv
 from huggingface_hub import HfApi
 
-load_dotenv(override=True)
 
-# Configuration
+def parse_args() -> argparse.Namespace:
-LOCAL_DIR = "trainer_output_deepseek-ai_DeepSeek-R1-Distill-Qwen-7B_gpu0_20250403_050520"
+    """Parse command line arguments.
-REPO_ID = "janhq/250403-runpod-qwen7b-r1-distil"
+
-HF_TOKEN = os.getenv("HF_TOKEN")
+    Returns:
+        argparse.Namespace: Parsed arguments
+    """
+    parser = argparse.ArgumentParser(description="Upload model to HuggingFace Hub")
+    parser.add_argument("--local-dir", type=str, required=True, help="Local directory to upload")
+    parser.add_argument("--repo-id", type=str, required=True, help="HuggingFace repository ID")
+    parser.add_argument("--public", action="store_true", help="Make repository public (default: private)")
+    return parser.parse_args()
+
+
+def main():
+    """Main function to upload model."""
+    args = parse_args()
+    load_dotenv(override=True)
 
-# Files to ignore during upload
+    # Configuration
-IGNORE_PATTERNS = [
+    HF_TOKEN = os.getenv("HF_TOKEN")
+
+    # Files to ignore during upload
+    IGNORE_PATTERNS = [
         "*.log",  # Log files
         "*.pyc",  # Python cache
         ".git*",  # Git files
@@ -26,14 +44,18 @@ IGNORE_PATTERNS = [
         "events.*",  # Tensorboard
         "wandb/*",  # Weights & Biases
         "runs/*",  # Training runs
-]
+    ]
 
-api = HfApi(token=HF_TOKEN)
+    api = HfApi(token=HF_TOKEN)
-api.create_repo(repo_id=REPO_ID, private=True, exist_ok=True, repo_type="model")
+    api.create_repo(repo_id=args.repo_id, private=not args.public, exist_ok=True, repo_type="model")
-api.upload_folder(
+    api.upload_folder(
-    folder_path=LOCAL_DIR,
+        folder_path=args.local_dir,
-    repo_id=REPO_ID,
+        repo_id=args.repo_id,
         repo_type="model",
-    # ignore_patterns=IGNORE_PATTERNS,
+        #  ignore_patterns=IGNORE_PATTERNS
-)
+    )
-print(f"✅ Done: {LOCAL_DIR} -> {REPO_ID}")
+    print(f"✅ Done: {args.local_dir} -> {args.repo_id}")
+
+
+if __name__ == "__main__":
+    main()