evaluate.py

"""Local evaluation entry point for HyperCLIP checkpoints."""

import argparse
import json
import os
import pickle
import re
import time
import warnings
from pathlib import Path

import yaml
from cloudpathlib import CloudPath

from eval_utils.main import evaluate_model


warnings.filterwarnings("ignore", message="Length of IterableDataset")


def path_or_cloudpath(value):
    if re.match(r"^\w+://", value):
        return CloudPath(value)
    return Path(value)


def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument("--train_output_dir", required=True,
                        help="Training output directory containing info.pkl.")
    parser.add_argument("--output_dir", default=None,
                        help="Evaluation output directory; defaults to train_output_dir.")
    parser.add_argument("--data_dir", default=None,
                        help="Root containing the prepared evaluation datasets.")
    parser.add_argument("--batch_size", default=128, type=int)
    parser.add_argument("--all_tasks", action="store_true",
                        help="Evaluate every task in tasklist.yml instead of validation tasks only.")
    parser.add_argument("--use_model", default=None,
                        help='Evaluate an explicit "ARCH CHECKPOINT" pair.')
    return parser.parse_args()


def main():
    args = parse_args()
    args.train_output_dir = Path(args.train_output_dir)
    args.output_dir = Path(args.output_dir) if args.output_dir else args.train_output_dir

    if args.use_model is not None:
        model_arch, model_checkpoint = args.use_model.split(maxsplit=1)
        args.train_output_dir = args.output_dir
        args.output_dir.mkdir(parents=True, exist_ok=True)
        with open(args.train_output_dir / "info.pkl", "wb") as handle:
            pickle.dump({
                "scale_config": {"model": model_arch},
                "checkpoint": model_checkpoint,
            }, handle)

    train_info_filename = args.train_output_dir / "info.pkl"
    with open(train_info_filename, "rb") as handle:
        train_info = pickle.load(handle)

    results_filename = args.output_dir / "eval_results.jsonl"
    with open(Path(__file__).with_name("tasklist.yml")) as handle:
        tasks = yaml.safe_load(handle)
    if not args.all_tasks:
        tasks = {key: value for key, value in tasks.items() if "val_task" in value["tags"]}

    results = {}
    cached_train_info_filename = args.output_dir / "info.pkl"
    if args.output_dir.exists() and cached_train_info_filename.exists():
        with open(cached_train_info_filename, "rb") as handle:
            cached_train_info = pickle.load(handle)
        assert cached_train_info == train_info, (
            "The output directory contains results for a different training config."
        )
        if results_filename.exists():
            with open(results_filename) as handle:
                for line in handle:
                    result = json.loads(line)
                    if result["key"] in tasks:
                        results[result["dataset"]] = result
            print(f"Found {len(results)} cached result(s) in {results_filename}.")
    else:
        args.output_dir.mkdir(parents=True, exist_ok=True)
        with open(cached_train_info_filename, "wb") as handle:
            pickle.dump(train_info, handle)

    try:
        checkpoint_exists = path_or_cloudpath(str(train_info["checkpoint"])).exists()
    except Exception:
        checkpoint_exists = False
    if not checkpoint_exists and args.use_model is None:
        fallback = args.train_output_dir / "checkpoints" / "epoch_latest.pt"
        print("Checkpoint not found at", train_info["checkpoint"])
        print("Defaulting to", fallback)
        train_info["checkpoint"] = fallback

    start_time = time.time()
    for task_key, task in tasks.items():
        task_name = task.get("name", task_key)
        if task_name in results:
            print(f"Skipping cached task: {task_name}")
        else:
            print(f"Evaluating: {task_name}")
            metrics = evaluate_model(
                task_key,
                train_info,
                args.data_dir,
                task.get("size"),
                batch_size=args.batch_size,
            )
            metrics["main_metric"] = metrics.get(task.get("main_metric", "acc1"))
            results[task_name] = {
                "key": task_key,
                "dataset": task_name,
                "metrics": metrics,
            }
            with open(results_filename, "a") as handle:
                handle.write(json.dumps(results[task_name]) + "\n")

        score = results[task_name]["metrics"]["main_metric"]
        print(f"Score: {score:.4f}" if score is not None else "Score: no summary metric")

    elapsed = int(time.time() - start_time)
    print(f"Evaluation time: {elapsed // 3600}h {(elapsed % 3600) // 60}m {elapsed % 60}s")
    print("=== Final results ===")
    for result in results.values():
        print(f'{result["dataset"]}: {result["metrics"]["main_metric"]}')


if __name__ == "__main__":
    main()