scripts/verify_merged_model.py

"""Verify that the merged model directory has all required files."""

from __future__ import annotations

import json
import sys
from pathlib import Path

# Add project root to Python path
project_root = Path(__file__).resolve().parents[1]
if str(project_root) not in sys.path:
    sys.path.insert(0, str(project_root))


def verify_merged_model(merged_dir: Path) -> bool:
    """Check if merged model directory has required files."""
    
    print(f"Checking merged model at: {merged_dir}")
    
    if not merged_dir.exists():
        print(f"✗ Directory does not exist: {merged_dir}")
        return False
    
    if not merged_dir.is_dir():
        print(f"✗ Path is not a directory: {merged_dir}")
        return False
    
    config_path = merged_dir / "config.json"
    if not config_path.exists():
        print(f"✗ Missing config.json at {config_path}")
        return False
    
    try:
        with config_path.open() as f:
            config = json.load(f)
        if "model_type" not in config:
            print(f"✗ config.json missing 'model_type' key")
            print(f"  Keys present: {list(config.keys())[:10]}")
            return False
        print(f"✓ config.json found with model_type: {config.get('model_type')}")
    except Exception as e:
        print(f"✗ Failed to parse config.json: {e}")
        return False
    
    # Check for model files
    safetensors_files = list(merged_dir.glob("*.safetensors"))
    if not safetensors_files:
        print(f"⚠ No .safetensors files found (might use .bin files)")
    else:
        print(f"✓ Found {len(safetensors_files)} safetensors file(s)")
    
    tokenizer_config = merged_dir / "tokenizer_config.json"
    if tokenizer_config.exists():
        print(f"✓ tokenizer_config.json found")
    else:
        print(f"⚠ tokenizer_config.json not found (may be in base model)")
    
    return True


def main() -> None:
    """Entry point."""
    
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--merged-dir",
        default="./models/merged",
        help="Path to merged model directory"
    )
    args = parser.parse_args()
    
    merged_path = Path(args.merged_dir).resolve()
    if verify_merged_model(merged_path):
        print("\n✓ Merged model structure looks good!")
        sys.exit(0)
    else:
        print("\n✗ Merged model structure has issues.")
        print("\nTo fix:")
        print("1. Re-run the merge script: python scripts/merge_model.py")
        print("2. Or use the base model by setting merged_path: null in agent/config.yaml")
        sys.exit(1)


if __name__ == "__main__":
    main()