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

import torch
from transformers import AutoModelForCausalLM, AutoConfig


def test_load():
    print("Loading original model config...")
    config = AutoConfig.from_pretrained("/data/models/Ornith-1.0-35B", trust_remote_code=True)
    
    print("\nLoading quantized test shard...")
    try:
        model = AutoModelForCausalLM.from_pretrained(
            "/data/models/test_quantize",
            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()