run_eval.py

# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import argparse
import os
import pprint
import sys

sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from omegaconf import OmegaConf

from nemo_rl.algorithms.utils import get_tokenizer
from nemo_rl.data.datasets import AllTaskProcessedDataset, load_eval_dataset
from nemo_rl.data.datasets.eval_datasets import _is_multimodal_dataset
from nemo_rl.data.datasets.response_datasets import load_response_dataset
from nemo_rl.distributed.virtual_cluster import init_ray
from nemo_rl.environments.utils import create_env
from nemo_rl.evals.eval import MasterConfig, run_env_eval, setup
from nemo_rl.models.generation import configure_generation_config
from nemo_rl.utils.config import load_config


def parse_args():
    """Parse command line arguments."""
    parser = argparse.ArgumentParser(description="Run Evaluation with configuration")
    parser.add_argument(
        "--config", type=str, default=None, help="Path to YAML config file"
    )

    # Parse known args for the script
    args, remaining = parser.parse_known_args()

    # Convert remaining args to OmegaConf format
    overrides = OmegaConf.from_dotlist(remaining)

    return args, overrides


def setup_data(tokenizer, data_config, env_configs, is_multimodal=False):
    print("Setting up data...")

    # load dataset
    # TODO(#2840): consolidate onto load_response_dataset. Migration is in
    # progress -- remaining eval-only datasets (mmlu, gpqa, math, mmau) will
    # move into DATASET_REGISTRY, at which point this branch collapses.
    # DATASET_REGISTRY cannot be used as the gate yet because it also contains
    # datasets such as daily-omni that still require an eval-specific wrapper.
    if data_config["dataset_name"] in {"AIME2024", "AIME2025", "AIME2026"}:
        base_dataset = load_response_dataset(data_config)
        rekeyed_ds = base_dataset.dataset
    else:
        base_dataset = load_eval_dataset(data_config)
        rekeyed_ds = base_dataset.rekeyed_ds

    # Mirrors nemo_rl/data/utils.py: use data.env_name to look up the env
    # config block and determine the registered environment class.
    env_key = next(iter(env_configs))
    env_name = data_config.get("env_name", env_key)
    registered_env_name = "vlm" if is_multimodal else env_name
    env = create_env(env_name=registered_env_name, env_config=env_configs[env_name])

    dataset = AllTaskProcessedDataset(
        dataset=rekeyed_ds,
        tokenizer=tokenizer,
        default_task_data_spec=base_dataset.task_spec,
        task_data_processors=base_dataset.processor,
        task_data_preprocessors=getattr(base_dataset, "preprocessor", None),
        max_seq_length=data_config["max_input_seq_length"],
    )

    return dataset, env, tokenizer


def main():
    """Main entry point."""
    # Parse arguments
    args, overrides = parse_args()

    if not args.config:
        args.config = os.path.join(
            os.path.dirname(__file__), "configs", "evals", "eval.yaml"
        )

    config = load_config(args.config)
    print(f"Loaded configuration from: {args.config}")

    if overrides:
        override_conf = OmegaConf.from_cli()
        print(f"Overrides: {override_conf}")
        config = OmegaConf.merge(config, override_conf)

    config = OmegaConf.to_container(config, resolve=True)
    config = MasterConfig(**config)
    print("Applied CLI overrides")

    # Print config
    print("Final config:")
    pprint.pprint(config)

    # Init ray
    init_ray()

    # Setup tokenizer — get_tokenizer handles both text-only and multimodal
    is_multimodal = _is_multimodal_dataset(config.data["dataset_name"])
    tokenizer = get_tokenizer(config.tokenizer, get_processor=is_multimodal)
    config.generation = configure_generation_config(
        config.generation, tokenizer, is_eval=True
    )

    # Setup data
    (
        dataset,
        env,
        tokenizer,
    ) = setup_data(tokenizer, config.data, config.env, is_multimodal=is_multimodal)

    # Setup
    (
        vllm_generation,
        dataloader,
        master_config,
    ) = setup(config, tokenizer, dataset)

    # Run evaluation
    run_env_eval(
        vllm_generation,
        dataloader,
        env,
        master_config,
    )


if __name__ == "__main__":
    main()