#!/usr/bin/env python3
"""Test loading quantized shard to check for MISMATCH."""

import gc
import torch
import shutil
from pathlib import Path
from transformers import AutoModelForCausalLM, AutoConfig


def test_load():
    # Copy config.json from original model
    test_dir = Path("/data/models/test_quantize")
    config_src = Path("/data/models/Ornith-1.0-35B") / "config.json"
    config_dst = test_dir / "config.json"
    if not config_dst.exists():
        shutil.copy2(config_src, config_dst)
        print(f"Copied config.json to {test_dir}")
    
    print("\nLoading quantized test shard...")
    try:
        model = AutoModelForCausalLM.from_pretrained(
            str(test_dir),
            device_map="cpu",
            torch_dtype=torch.float16,
            trust_remote_code=True,
            low_cpu_mem_usage=True,
        )
        print("✓ Model loaded successfully!")
        
        # Check for mismatched parameters
        print("\nChecking parameter shapes...")
        mismatch_count = 0
        for name, param in model.named_parameters():
            if hasattr(param, 'quant_state') and param.quant_state is not None:
                # Quantized parameter - check if it loaded correctly
                expected_shape = config.get_shape_for_parameter(name) if hasattr(config, 'get_shape_for_parameter') else None
                if expected_shape and tuple(param.shape) != expected_shape:
                    print(f"✗ MISMATCH: {name} - expected {expected_shape}, got {tuple(param.shape)}")
                    mismatch_count += 1
                else:
                    print(f"✓ {name} - {tuple(param.shape)}")
        
        if mismatch_count == 0:
            print("\n✅ No MISMATCH errors!")
        else:
            print(f"\n❌ {mismatch_count} MISMATCH errors found!")
            
    except Exception as e:
        print(f"✗ Failed to load: {e}")
        import traceback
        traceback.print_exc()


if __name__ == "__main__":
    test_load()