scripts/upload_to_hf.py

"""
Upload Milk Spoilage Classification Model to Hugging Face

This script uploads the trained model and associated files to Hugging Face Hub.
It handles both the model repository and optionally the Gradio Space.

Usage:
    python upload_to_hf.py                    # Interactive login
    python upload_to_hf.py --token YOUR_TOKEN # Direct token login
"""

import os
import sys
from huggingface_hub import login, upload_folder, create_repo, HfApi


# Configuration
MODEL_REPO_ID = "chenhaoq87/MilkSpoilageClassifier"
SPACE_REPO_ID = "chenhaoq87/MilkSpoilageClassifier-Demo"

# Files to include in model repository
MODEL_FILES = [
    "model/model.joblib",
    "model/config.json", 
    "apps/huggingface/handler.py",
    "apps/huggingface/requirements.txt",
    "README.md"
]

# Files to include in Gradio Space
SPACE_FILES = [
    "apps/gradio/app.py",
    "model/model.joblib",
]

# Patterns to ignore when uploading
IGNORE_PATTERNS = [
    "*.csv",           # Data files
    "*.ipynb",         # Jupyter notebooks
    "*.png",           # Image files
    "__pycache__/*",   # Python cache
    ".git/*",          # Git files
    "*.pyc",           # Compiled Python
    "upload_to_hf.py", # This script
    "prepare_model.py", # Preparation script
]


def check_required_files():
    """Check if all required files exist before uploading."""
    print("Checking required files...")
    
    missing_files = []
    for file in MODEL_FILES:
        if not os.path.exists(file):
            missing_files.append(file)
    
    if missing_files:
        print(f"\n[ERROR] Missing required files: {missing_files}")
        print("Please run 'python scripts/prepare_model.py' first to generate these files.")
        return False
    
    print("[OK] All required files found!")
    return True


def upload_model_repo():
    """Upload model files to Hugging Face model repository."""
    print(f"\n[UPLOAD] Uploading to model repository: {MODEL_REPO_ID}")
    
    try:
        # Create repo if it doesn't exist
        api = HfApi()
        try:
            api.create_repo(repo_id=MODEL_REPO_ID, repo_type="model", exist_ok=True)
            print(f"[OK] Repository '{MODEL_REPO_ID}' ready")
        except Exception as e:
            print(f"Note: {e}")
        
        # Upload the folder
        upload_folder(
            folder_path=".",
            repo_id=MODEL_REPO_ID,
            repo_type="model",
            ignore_patterns=IGNORE_PATTERNS
        )
        
        print(f"[OK] Model uploaded successfully!")
        print(f"   View at: https://huggingface.co/{MODEL_REPO_ID}")
        return True
        
    except Exception as e:
        print(f"[ERROR] Error uploading model: {e}")
        return False


def create_space_requirements():
    """Create a separate requirements file for the Gradio Space."""
    space_requirements = """gradio>=4.0
scikit-learn>=1.0
joblib>=1.0
numpy>=1.20
"""
    with open("space_requirements.txt", "w") as f:
        f.write(space_requirements)
    return "space_requirements.txt"


def upload_gradio_space():
    """Upload Gradio app to Hugging Face Space."""
    print(f"\n[UPLOAD] Uploading to Gradio Space: {SPACE_REPO_ID}")
    
    try:
        api = HfApi()
        
        # Create Space if it doesn't exist
        try:
            api.create_repo(
                repo_id=SPACE_REPO_ID, 
                repo_type="space", 
                space_sdk="gradio",
                exist_ok=True
            )
            print(f"[OK] Space '{SPACE_REPO_ID}' ready")
        except Exception as e:
            print(f"Note: {e}")
        
        # Create space-specific requirements
        space_req_file = create_space_requirements()
        
        # Upload app.py
        api.upload_file(
            path_or_fileobj="apps/gradio/app.py",
            path_in_repo="app.py",
            repo_id=SPACE_REPO_ID,
            repo_type="space"
        )
        print("  [OK] Uploaded app.py")
        
        # Upload model file
        api.upload_file(
            path_or_fileobj="model/model.joblib",
            path_in_repo="model.joblib",
            repo_id=SPACE_REPO_ID,
            repo_type="space"
        )
        print("  [OK] Uploaded model.joblib")
        
        # Upload requirements as requirements.txt for Space
        api.upload_file(
            path_or_fileobj=space_req_file,
            path_in_repo="requirements.txt",
            repo_id=SPACE_REPO_ID,
            repo_type="space"
        )
        print("  [OK] Uploaded requirements.txt")
        
        # Clean up temporary file
        os.remove(space_req_file)
        
        print(f"[OK] Space uploaded successfully!")
        print(f"   View at: https://huggingface.co/spaces/{SPACE_REPO_ID}")
        return True
        
    except Exception as e:
        print(f"[ERROR] Error uploading Space: {e}")
        return False


def main():
    """Main function to upload all files to Hugging Face."""
    print("=" * 60)
    print("Hugging Face Upload Script")
    print("=" * 60)
    
    # Check files exist
    if not check_required_files():
        return
    
    # Check for token argument
    token = None
    if len(sys.argv) > 1:
        if sys.argv[1] == "--token" and len(sys.argv) > 2:
            token = sys.argv[2]
        elif sys.argv[1].startswith("hf_"):
            token = sys.argv[1]
    
    # Login to Hugging Face
    print("\n[LOGIN] Logging in to Hugging Face...")
    try:
        if token:
            login(token=token, add_to_git_credential=False)
        else:
            print("   (You'll be prompted to enter your token if not already logged in)")
            login()
        print("[OK] Login successful!")
    except Exception as e:
        print(f"[ERROR] Login failed: {e}")
        print("\nTry running with your token directly:")
        print("   python upload_to_hf.py --token hf_YOUR_TOKEN_HERE")
        print("\nGet your token at: https://huggingface.co/settings/tokens")
        return
    
    # Upload model repository
    model_success = upload_model_repo()
    
    # Ask about Space upload
    if model_success:
        print("\n" + "-" * 60)
        # Auto-yes if token was provided via command line
        if token:
            print("Creating Gradio Space (use --no-space to skip)...")
            if "--no-space" not in sys.argv:
                upload_gradio_space()
            else:
                print("Skipping Gradio Space upload.")
        else:
            response = input("Would you like to also create/update the Gradio Space? (y/n): ")
            if response.lower() in ['y', 'yes']:
                upload_gradio_space()
            else:
                print("Skipping Gradio Space upload.")
    
    # Summary
    print("\n" + "=" * 60)
    print("Upload Complete!")
    print("=" * 60)
    
    if model_success:
        print(f"\n[SUCCESS] Your model is available at:")
        print(f"   https://huggingface.co/{MODEL_REPO_ID}")
        print(f"\n[API] To use the Inference API:")
        print(f"   POST https://api-inference.huggingface.co/models/{MODEL_REPO_ID}")
        print(f"\n[EXAMPLE] API call:")
        print(f'''   curl -X POST \\
     -H "Authorization: Bearer YOUR_HF_TOKEN" \\
     -H "Content-Type: application/json" \\
     -d '{{"inputs": [[4.5, 5.2, 6.1, 3.2, 4.0, 4.8]]}}' \\
     https://api-inference.huggingface.co/models/{MODEL_REPO_ID}''')


if __name__ == "__main__":
    main()