#!/usr/bin/env python3
"""Test loading all quantized shards."""

import torch
from pathlib import Path
import glob


def test_all_shards():
    model_path = "/data/models/Ornith-1.0-35B-nf4"
    shards = sorted(glob.glob(f"{model_path}/*.safetensors"))
    
    print(f"Found {len(shards)} shards\n")
    
    total_tensors = 0
    total_params = 0
    
    for i, shard_path in enumerate(shards):
        print(f"Loading shard {i+1}/{len(shards)}: {Path(shard_path).name}")
        try:
            ckpt = torch.load(shard_path, map_location="cpu", weights_only=False)
            shard_tensors = len([k for k in ckpt.keys() if isinstance(ckpt[k], torch.Tensor)])
            total_tensors += shard_tensors
            print(f"   ✓ {shard_tensors} tensors")
        except Exception as e:
            print(f"   ✗ Failed: {e}")
    
    print(f"\nTotal: {total_tensors} tensors across {len(shards)} shards")
    print("✅ All shards loaded successfully!")


if __name__ == "__main__":
    test_all_shards()