upload_to_hub.py

import os
import argparse
import shutil
from huggingface_hub import HfApi, create_repo, upload_folder

def upload_model(args):
    """Upload the fine-tuned model to Hugging Face Hub."""
    print(f"Preparing to upload Heaven1-base Guardian model to {args.org}/{args.model_id}")
    
    # Create a temporary directory for preparing the model
    temp_dir = "temp_upload"
    os.makedirs(temp_dir, exist_ok=True)
    
    # Copy model files
    print("Copying model files...")
    if os.path.exists(args.model_path):
        for item in os.listdir(args.model_path):
            source = os.path.join(args.model_path, item)
            dest = os.path.join(temp_dir, item)
            if os.path.isdir(source):
                shutil.copytree(source, dest, dirs_exist_ok=True)
            else:
                shutil.copy2(source, dest)
    else:
        print(f"Warning: Model directory {args.model_path} not found.")
    
    # Copy documentation files
    print("Copying documentation files...")
    docs_files = ["README.md", "model_card.md", "Heaven1-guardian.png"]
    for file in docs_files:
        if os.path.exists(file):
            shutil.copy2(file, os.path.join(temp_dir, file))
    
    # Rename model_card.md to README.md for proper display on the Hub
    if os.path.exists(os.path.join(temp_dir, "model_card.md")):
        print("Using model_card.md as the main README for the Hub...")
        if os.path.exists(os.path.join(temp_dir, "README.md")):
            # If both exist, rename the original README to avoid overwriting
            os.rename(os.path.join(temp_dir, "README.md"), os.path.join(temp_dir, "DETAILED_README.md"))
        os.rename(os.path.join(temp_dir, "model_card.md"), os.path.join(temp_dir, "README.md"))
    
    # Initialize Hugging Face API
    api = HfApi()
    
    # Create repository if it doesn't exist
    try:
        print(f"Creating repository: {args.org}/{args.model_id}")
        create_repo(
            repo_id=f"{args.org}/{args.model_id}",
            token=args.token,
            private=args.private,
            repo_type="model",
            exist_ok=True,
        )
    except Exception as e:
        print(f"Repository creation error (it might already exist): {e}")
    
    # Upload model to Hugging Face Hub
    print(f"Uploading files to {args.org}/{args.model_id}...")
    response = upload_folder(
        folder_path=temp_dir,
        repo_id=f"{args.org}/{args.model_id}",
        token=args.token,
        repo_type="model",
        ignore_patterns=[".*", "__pycache__/*", "temp_upload/*"],
    )
    
    print(f"Upload complete! Model available at: https://huggingface.co/{args.org}/{args.model_id}")
    
    # Clean up
    if not args.keep_temp:
        print("Cleaning up temporary directory...")
        shutil.rmtree(temp_dir)
    
    return response

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Upload Heaven1-base Guardian model to Hugging Face Hub")
    parser.add_argument("--model_path", type=str, default="./heaven1-base-8b", 
                        help="Path to the fine-tuned model directory")
    parser.add_argument("--org", type=str, default="safecircleia", 
                        help="Organization name on Hugging Face Hub")
    parser.add_argument("--model_id", type=str, default="heaven1-base-guardian", 
                        help="Model ID for the repository")
    parser.add_argument("--token", type=str, required=True, 
                        help="Hugging Face authentication token")
    parser.add_argument("--private", action="store_true", 
                        help="Whether to make the repository private")
    parser.add_argument("--keep_temp", action="store_true", 
                        help="Keep temporary upload directory after completion")
    
    args = parser.parse_args()
    
    upload_model(args)