#!/usr/bin/env python3
"""Inspect the Ornith-1.0-35B model architecture"""

import torch
from transformers import AutoModelForCausalLM, AutoConfig

print("=" * 80)
print("Inspecting Ornith-1.0-35B model")
print("=" * 80)

# Load config
print("\n1. Loading config...")
config = AutoConfig.from_pretrained("/data/models/Ornith-1.0-35B", trust_remote_code=True)
print(f"   Config class: {type(config).__name__}")
print(f"   Model type: {config.model_type}")

# Load model to CPU
print("\n2. Loading model to CPU...")
model = AutoModelForCausalLM.from_pretrained(
    "/data/models/Ornith-1.0-35B",
    device_map="cpu",
    torch_dtype=torch.bfloat16,
    trust_remote_code=True,
    low_cpu_mem_usage=True,
)
print(f"   Model class: {type(model).__name__}")

# Check if model has quantize_4bit
print("\n3. Checking for quantization methods...")
has_quantize = hasattr(model, 'quantize_4bit')
print(f"   Has quantize_4bit(): {has_quantize}")

# List all model components
print("\n4. Model components:")
for name, module in model.named_modules():
    if len(name.split('.')) <= 2:  # Top-level and first-level
        print(f"   {name}: {type(module).__name__}")

# Check for BnB quantization support
print("\n5. Checking BnB support...")
try:
    from bitsandbytes.nn import Linear4bit, Linear8bitLt
    print("   ✓ BnB 4bit and 8bit modules available")
except ImportError:
    print("   ✗ BnB not installed")

# Check if we can use prepare_model_for_kbit_training
print("\n6. Checking PEFT support...")
try:
    from peft import prepare_model_for_kbit_training
    print("   ✓ prepare_model_for_kbit_training available")
except ImportError:
    print("   ✗ PEFT not installed")

print("\n" + "=" * 80)
print("Summary:")
print(f"   Model: {type(model).__name__}")
print(f"   Config: {type(config).__name__}")
print(f"   Has quantize_4bit(): {has_quantize}")
print("=" * 80)