verify_models.py

"""
Proprietary / All Rights Reserved - Non-Commercial Use Only
Source-available for portfolio viewing only. Commercial use, unauthorized modification, reproduction, or distribution is strictly prohibited. All rights reserved.
"""
import glob
import os
import torch
from utils_v3_0 import setup_logging

logger = setup_logging(__name__)


def check_model(filename):
    logger.info(f"Checking {filename}...")
    if not os.path.exists(filename):
        logger.info(f"File {filename} not found.")
        return
    try:
        checkpoint = torch.load(filename, map_location="cpu", weights_only=False)
        logger.info(f"Loaded {filename}")
        if "model" in checkpoint:
            logger.info(f"  state_dict keys: {len(checkpoint['model'])}")
        else:
            logger.warning("  Missing 'model' key")
        if "config" in checkpoint:
            logger.info(f"  config: {checkpoint['config']}")
        if "step" in checkpoint:
            logger.info(f"  step: {checkpoint['step']}")
    except Exception as e:
        logger.error(f"Failed to load {filename}: {e}")


if __name__ == "__main__":
    logger.info("--- Verifying Saved Models ---")
    files = glob.glob("checkpoint_step_*.pt") + glob.glob("best_model.pt")
    if not files:
        logger.info("No models found.")
    for f in sorted(files):
        check_model(f)