download_models.py

"""Script to download and prepare models for HuggingFace Spaces."""

import os
import asyncio
import logging
from pathlib import Path
from huggingface_hub import HfApi, upload_file
from reasoning.model_manager import ModelManager

# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

async def download_and_prepare_models():
    """Download all models and prepare for Spaces."""
    try:
        # Initialize model manager
        model_dir = os.path.join(os.getcwd(), "models")
        manager = ModelManager(model_dir)
        
        # Create models directory
        os.makedirs(model_dir, exist_ok=True)
        
        # Download all models
        logger.info("Starting model downloads...")
        await manager.initialize_all_models()
        logger.info("All models downloaded successfully!")
        
        return True
        
    except Exception as e:
        logger.error(f"Error downloading models: {e}")
        return False

def upload_to_spaces(space_name: str = "agentic-system-models"):
    """Upload models to HuggingFace Spaces."""
    try:
        api = HfApi()
        model_dir = os.path.join(os.getcwd(), "models")
        
        # Create .gitattributes for LFS
        gitattributes_path = os.path.join(model_dir, ".gitattributes")
        with open(gitattributes_path, "w") as f:
            f.write("*.gguf filter=lfs diff=lfs merge=lfs -text")
        
        # Upload .gitattributes first
        api.upload_file(
            path_or_fileobj=gitattributes_path,
            path_in_repo=".gitattributes",
            repo_id=f"spaces/{space_name}",
            repo_type="space"
        )
        
        # Upload each model file
        for model_file in Path(model_dir).glob("*.gguf"):
            logger.info(f"Uploading {model_file.name}...")
            api.upload_file(
                path_or_fileobj=str(model_file),
                path_in_repo=f"models/{model_file.name}",
                repo_id=f"spaces/{space_name}",
                repo_type="space"
            )
            
        logger.info("All models uploaded to Spaces successfully!")
        return True
        
    except Exception as e:
        logger.error(f"Error uploading to Spaces: {e}")
        return False

if __name__ == "__main__":
    # Download models
    asyncio.run(download_and_prepare_models())
    
    # Upload to Spaces
    upload_to_spaces()