inference/load_model.py

"""
Helion-OSC Sharded Model Loader
Efficiently loads 116 safetensors shards (2.8GB each)
"""

import torch
import json
import os
from pathlib import Path
from typing import Dict, Optional, List
import logging
from tqdm import tqdm
from safetensors.torch import load_file
from transformers import AutoConfig, AutoTokenizer
import psutil

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


class ShardedModelLoader:
    """
    Loader for sharded safetensors model files
    Optimized for 116 shards of 2.8GB each
    """
    
    def __init__(self, model_path: str):
        """
        Initialize the sharded model loader
        
        Args:
            model_path: Path to the inference directory containing shards
        """
        self.model_path = Path(model_path)
        self.config_path = self.model_path / "model_config.json"
        self.index_path = self.model_path / "model.safetensors.index.json"
        
        # Load configuration
        logger.info(f"Loading configuration from {self.config_path}")
        with open(self.config_path, 'r') as f:
            self.config = json.load(f)
        
        # Load weight index
        logger.info(f"Loading weight index from {self.index_path}")
        with open(self.index_path, 'r') as f:
            self.index = json.load(f)
        
        self.metadata = self.index.get("metadata", {})
        self.weight_map = self.index.get("weight_map", {})
        
        logger.info(f"Model: {self.metadata.get('model_type', 'unknown')}")
        logger.info(f"Total shards: {self.metadata.get('total_shards', 0)}")
        logger.info(f"Total size: {self.metadata.get('total_size', 0) / 1e9:.2f} GB")
        logger.info(f"Total parameters: {self.config['architectures_info']['total_parameters']}")
        logger.info(f"Active parameters: {self.config['architectures_info']['active_parameters']}")
    
    def get_shard_path(self, shard_name: str) -> Path:
        """Get full path to a shard file"""
        return self.model_path / shard_name
    
    def get_available_memory(self) -> Dict[str, float]:
        """Get available system memory"""
        memory = psutil.virtual_memory()
        result = {
            "ram_total_gb": memory.total / 1e9,
            "ram_available_gb": memory.available / 1e9,
            "ram_percent_used": memory.percent
        }
        
        if torch.cuda.is_available():
            for i in range(torch.cuda.device_count()):
                gpu_mem = torch.cuda.get_device_properties(i).total_memory
                gpu_allocated = torch.cuda.memory_allocated(i)
                result[f"gpu_{i}_total_gb"] = gpu_mem / 1e9
                result[f"gpu_{i}_available_gb"] = (gpu_mem - gpu_allocated) / 1e9
        
        return result
    
    def load_shard(self, shard_name: str, device: str = "cpu") -> Dict[str, torch.Tensor]:
        """
        Load a single shard file
        
        Args:
            shard_name: Name of the shard file
            device: Device to load tensors to
            
        Returns:
            Dictionary of weight tensors
        """
        shard_path = self.get_shard_path(shard_name)
        
        if not shard_path.exists():
            raise FileNotFoundError(f"Shard not found: {shard_path}")
        
        logger.debug(f"Loading shard: {shard_name}")
        return load_file(str(shard_path), device=device)
    
    def load_sharded_weights(
        self,
        device: str = "cpu",
        low_memory: bool = False,
        show_progress: bool = True
    ) -> Dict[str, torch.Tensor]:
        """
        Load all sharded weights
        
        Args:
            device: Device to load weights to
            low_memory: Use memory-efficient loading
            show_progress: Show progress bar
            
        Returns:
            Dictionary of all model weights
        """
        logger.info("Loading sharded model weights...")
        
        # Check available memory
        mem_info = self.get_available_memory()
        logger.info(f"Available RAM: {mem_info['ram_available_gb']:.2f} GB")
        if "gpu_0_available_gb" in mem_info:
            logger.info(f"Available GPU 0: {mem_info['gpu_0_available_gb']:.2f} GB")
        
        # Get unique shard files
        shard_files = sorted(set(self.weight_map.values()))
        total_shards = len(shard_files)
        
        logger.info(f"Loading {total_shards} shard files...")
        
        all_weights = {}
        
        # Create progress bar
        pbar = tqdm(shard_files, disable=not show_progress, desc="Loading shards")
        
        for shard_name in pbar:
            pbar.set_description(f"Loading {shard_name}")
            
            # Load shard
            shard_weights = self.load_shard(shard_name, device=device)
            
            # Add to all weights
            all_weights.update(shard_weights)
            
            # Clear memory if low_memory mode
            if low_memory:
                del shard_weights
                if torch.cuda.is_available():
                    torch.cuda.empty_cache()
        
        logger.info(f"Loaded {len(all_weights)} weight tensors")
        return all_weights
    
    def get_layer_weights(self, layer_idx: int) -> List[str]:
        """
        Get all weight keys for a specific layer
        
        Args:
            layer_idx: Layer index
            
        Returns:
            List of weight keys for that layer
        """
        prefix = f"model.layers.{layer_idx}."
        return [k for k in self.weight_map.keys() if k.startswith(prefix)]
    
    def get_shard_for_weight(self, weight_key: str) -> Optional[str]:
        """
        Get shard file name for a specific weight
        
        Args:
            weight_key: Weight key/name
            
        Returns:
            Shard file name or None
        """
        return self.weight_map.get(weight_key)
    
    def verify_shards(self) -> Dict[str, bool]:
        """
        Verify all shard files exist
        
        Returns:
            Dictionary mapping shard names to existence status
        """
        logger.info("Verifying shard files...")
        
        shard_files = set(self.weight_map.values())
        verification = {}
        
        for shard_name in tqdm(sorted(shard_files), desc="Verifying"):
            shard_path = self.get_shard_path(shard_name)
            verification[shard_name] = shard_path.exists()
        
        missing = [s for s, exists in verification.items() if not exists]
        
        if missing:
            logger.warning(f"Missing {len(missing)} shard files:")
            for shard in missing[:10]:  # Show first 10
                logger.warning(f"  - {shard}")
            if len(missing) > 10:
                logger.warning(f"  ... and {len(missing) - 10} more")
        else:
            logger.info("✓ All shard files present")
        
        return verification
    
    def load_metadata(self) -> Dict:
        """Load model metadata"""
        return {
            "config": self.config,
            "index": self.index,
            "total_shards": self.metadata.get("total_shards", 0),
            "total_size_gb": self.metadata.get("total_size", 0) / 1e9,
            "architecture": self.config.get("architectures_info", {}),
            "num_layers": self.config.get("num_hidden_layers", 0),
            "hidden_size": self.config.get("hidden_size", 0),
            "vocab_size": self.config.get("vocab_size", 0)
        }


def load_full_model(
    model_path: str,
    device: str = "cuda" if torch.cuda.is_available() else "cpu",
    low_memory: bool = False
):
    """
    Convenience function to load the full model
    
    Args:
        model_path: Path to inference directory
        device: Device to load model to
        low_memory: Use low memory loading
        
    Returns:
        Loaded model weights and metadata
    """
    loader = ShardedModelLoader(model_path)
    
    # Verify shards first
    verification = loader.verify_shards()
    missing = sum(1 for exists in verification.values() if not exists)
    
    if missing > 0:
        raise FileNotFoundError(
            f"Cannot load model: {missing} shard files are missing. "
            f"Please download all 116 shard files."
        )
    
    # Load weights
    weights = loader.load_sharded_weights(
        device=device,
        low_memory=low_memory,
        show_progress=True
    )
    
    # Load metadata
    metadata = loader.load_metadata()
    
    return weights, metadata


def inspect_model(model_path: str):
    """
    Inspect model structure without loading weights
    
    Args:
        model_path: Path to inference directory
    """
    loader = ShardedModelLoader(model_path)
    
    print("\n" + "="*80)
    print("HELION-OSC MODEL INSPECTION")
    print("="*80)
    
    metadata = loader.load_metadata()
    
    print(f"\nModel Type: {metadata['architecture'].get('model_description', 'N/A')}")
    print(f"Architecture: {metadata['architecture'].get('architecture_type', 'N/A')}")
    print(f"Total Parameters: {metadata['architecture'].get('total_parameters', 'N/A')}")
    print(f"Active Parameters: {metadata['architecture'].get('active_parameters', 'N/A')}")
    
    print(f"\nModel Configuration:")
    print(f"  Layers: {metadata['num_layers']}")
    print(f"  Hidden Size: {metadata['hidden_size']}")
    print(f"  Vocabulary Size: {metadata['vocab_size']}")
    print(f"  Attention Heads: {metadata['config'].get('num_attention_heads', 'N/A')}")
    print(f"  KV Heads: {metadata['config'].get('num_key_value_heads', 'N/A')}")
    
    print(f"\nMoE Configuration:")
    arch = metadata['architecture']
    print(f"  Number of Experts: {arch.get('num_experts', 'N/A')}")
    print(f"  Experts per Token: {arch.get('experts_per_token', 'N/A')}")
    print(f"  Shared Experts: {arch.get('num_shared_experts', 'N/A')}")
    
    print(f"\nStorage Information:")
    print(f"  Total Shards: {metadata['total_shards']}")
    print(f"  Total Size: {metadata['total_size_gb']:.2f} GB")
    print(f"  Shard Size: ~2.8 GB each")
    print(f"  Format: safetensors")
    print(f"  Precision: bfloat16")
    
    print(f"\nContext Length:")
    print(f"  Max Position Embeddings: {metadata['config'].get('max_position_embeddings', 'N/A')}")
    print(f"  RoPE Theta: {metadata['config'].get('rope_theta', 'N/A')}")
    
    print("\n" + "="*80)
    
    # Verify shards
    print("\nVerifying shard files...")
    verification = loader.verify_shards()
    present = sum(1 for exists in verification.values() if exists)
    total = len(verification)
    
    print(f"\nShard Status: {present}/{total} files present")
    
    if present == total:
        print("✓ All shard files are available")
    else:
        print(f"✗ Missing {total - present} shard files")


def main():
    """Main CLI interface"""
    import argparse
    
    parser = argparse.ArgumentParser(description="Helion-OSC Sharded Model Loader")
    parser.add_argument(
        "model_path",
        type=str,
        help="Path to inference directory"
    )
    parser.add_argument(
        "--action",
        choices=["inspect", "verify", "load"],
        default="inspect",
        help="Action to perform"
    )
    parser.add_argument(
        "--device",
        type=str,
        default="cuda" if torch.cuda.is_available() else "cpu",
        help="Device to load model to"
    )
    parser.add_argument(
        "--low-memory",
        action="store_true",
        help="Use low memory mode"
    )
    
    args = parser.parse_args()
    
    if args.action == "inspect":
        inspect_model(args.model_path)
    
    elif args.action == "verify":
        loader = ShardedModelLoader(args.model_path)
        loader.verify_shards()
    
    elif args.action == "load":
        logger.info("Loading full model...")
        weights, metadata = load_full_model(
            args.model_path,
            device=args.device,
            low_memory=args.low_memory
        )
        logger.info(f"Successfully loaded {len(weights)} weight tensors")
        logger.info(f"Model ready on {args.device}")


if __name__ == "__main__":
    main()