#!/usr/bin/env python """Evaluation script for trained ELF models: loads a checkpoint and generates text samples.""" import argparse import contextlib import copy import logging import os import sys # Initialize JAX distributed BEFORE importing other JAX modules import jax try: jax.distributed.initialize() except (RuntimeError, ValueError): pass # Single-host run, or already initialized. # Ensure repo root on sys.path so imports work when run as a script REPO_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) if REPO_ROOT not in sys.path: sys.path.insert(0, REPO_ROOT) import jax.numpy as jnp import optax from flax import jax_utils from transformers import AutoTokenizer from modules.t5_encoder import get_encoder from modules.model import ELF_models from utils.logging_utils import log_for_0 from utils.checkpoint_utils import load_encoder_checkpoint, load_checkpoint from utils.train_utils import TrainState from utils.data_utils import load_jsonl_dataset, load_dataset_split, get_pad_token_id from generation import test_generation_uncond, test_generation_cond from configs.config import load_config_from_yaml, apply_config_overrides, load_sampling_configs logging.basicConfig( format="%(levelname)s - %(name)s - %(message)s", handlers=[logging.StreamHandler(sys.stdout)], level=logging.INFO, force=True, ) logger = logging.getLogger(__name__) def parse_args(): parser = argparse.ArgumentParser(description="Evaluate trained ELF model by generating text samples") parser.add_argument("--config", type=str, required=True, help="Path to configuration YAML file") parser.add_argument( "--config_override", action="append", default=[], help="Override config values (field_name=value). Can be specified multiple times.", ) parser.add_argument("--seed", type=int, default=42, help="Random seed (used when --seeds is not specified)") parser.add_argument( "--seeds", type=str, default=None, help="Comma-separated list of seeds to evaluate (e.g. '42,123,456'). Overrides --seed.", ) parser.add_argument( "--checkpoint_path", type=str, required=True, help="Path to checkpoint file (e.g. outputs/elf_b-owt/checkpoint_19000) or HF repo id.", ) parser.add_argument( "--use_cpu", action="store_true", help="Host model init, train state template, and encoder/state replication on CPU", ) return parser.parse_args() def main(): args = parse_args() log_for_0("Loading configuration...") config = load_config_from_yaml(args.config) if args.config_override: config = apply_config_overrides(config, args.config_override) log_for_0(f"Applied {len(args.config_override)} config override(s)") num_devices = jax.device_count() num_local_devices = jax.local_device_count() num_hosts = jax.process_count() cpu_device = jax.local_devices(backend="cpu")[0] if args.use_cpu else None def cpu_ctx(): return jax.default_device(cpu_device) if args.use_cpu else contextlib.nullcontext() if config.global_batch_size is not None: log_for_0(f"Using global batch size for evaluation: {config.global_batch_size}") total_batch_size = config.global_batch_size local_batch_size = total_batch_size // num_hosts config.batch_size = local_batch_size elif config.batch_size is not None: log_for_0(f"Using batch size per device: {config.batch_size}") total_batch_size = config.batch_size * num_devices local_batch_size = config.batch_size * num_local_devices config.global_batch_size = total_batch_size else: raise ValueError("Either global_batch_size or batch_size must be specified") log_for_0(f"Config loaded from {args.config}") log_for_0(f"Model: {config.model}") log_for_0(f"Encoder Model: {config.encoder_model_name}") log_for_0(f"Encoder Checkpoint: {config.encoder_checkpoint}") log_for_0(f"Max length: {config.max_length}") log_for_0(f"Max input length: {config.max_input_length}") log_for_0(f"Num samples: {config.num_samples}") log_for_0(f"Sampling configs: {len(config.sampling_configs)} config(s)") seed_list = [int(s.strip()) for s in args.seeds.split(",")] if args.seeds is not None else [args.seed] log_for_0(f"Seeds to evaluate: {seed_list}") rng = jax.random.PRNGKey(config.seed) log_for_0("Loading tokenizer...") tokenizer = AutoTokenizer.from_pretrained(config.tokenizer_name or config.encoder_model_name) pad_token_id = get_pad_token_id(tokenizer, config.pad_token) log_for_0(f"Using {'EOS' if config.pad_token == 'eos' else 'PAD'} token for padding: {pad_token_id}") eval_dataset = None if config.eval_data_path is not None: log_for_0("Loading dataset for conditional generation...") if config.eval_data_path.endswith(".jsonl"): eval_dataset = load_jsonl_dataset( config.eval_data_path, tokenizer, input_key="input", output_key="output", ) else: eval_dataset = load_dataset_split(config.eval_data_path) log_for_0(f"Eval dataset size: {len(eval_dataset)}") # ============================================ # Load Encoder (frozen) # ============================================ log_for_0(f"Loading Encoder config: {config.encoder_model_name}...") encoder_config, encoder_model, _ = get_encoder(config.encoder_model_name, jnp.float32) encoder_params = load_encoder_checkpoint(config.encoder_checkpoint) log_for_0("encoder weights loaded.") # Multi-device eval passes encoder params directly into pmap, so replicate them # across local accelerator devices. encoder_params = jax_utils.replicate(encoder_params) log_for_0(f"Encoder d_model: {encoder_config.d_model}") # ============================================ # Create ELF Model # ============================================ log_for_0(f"Creating {config.model} model...") rng, init_rng, dropout_rng = jax.random.split(rng, 3) max_length = config.max_length with cpu_ctx(): # 2x dim if self_cond_prob > 0 to initialize self_cond_proj layer _text_enc_dim = encoder_config.d_model input_dim = 2 * _text_enc_dim if config.self_cond_prob > 0 else _text_enc_dim dummy_x = jnp.ones((1, max_length, input_dim)) dummy_t = jnp.ones((1,)) dummy_self_cond_cfg_scale = jnp.ones((1,)) if config.num_self_cond_cfg_tokens > 0 else None log_for_0(f"Dummy x shape: {dummy_x.shape}") log_for_0(f"Dummy t shape: {dummy_t.shape}") vocab_size = tokenizer.vocab_size model = ELF_models[config.model]( text_encoder_dim=encoder_config.d_model, max_length=max_length, attn_drop=config.attn_dropout, proj_drop=config.proj_dropout, num_time_tokens=config.num_time_tokens, num_self_cond_cfg_tokens=config.num_self_cond_cfg_tokens, vocab_size=vocab_size, num_model_mode_tokens=config.num_model_mode_tokens, bottleneck_dim=config.bottleneck_dim, ) log_for_0("Initializing ELF model...") init_args = dict( x=dummy_x, t=dummy_t, deterministic=True, self_cond_cfg_scale=dummy_self_cond_cfg_scale, ) with cpu_ctx(): elf_params = model.init(init_rng, **init_args) log_for_0("\n" + model.tabulate(init_rng, **init_args)) log_for_0("ELF initialization complete") total_params = sum(x.size for x in jax.tree_util.tree_leaves(elf_params)) log_for_0(f"ELF parameters: {total_params:,}") # ============================================ # Create Train State Template # ============================================ optimizer = optax.adamw(learning_rate=1e-4) with cpu_ctx(): state = TrainState.create( apply_fn=model.apply, params=elf_params["params"], tx=optimizer, dropout_rng=dropout_rng, ema_params1=copy.deepcopy(elf_params["params"]), ) # ============================================ # Determine checkpoints to evaluate # ============================================ if config.sampling_configs_path: config.sampling_configs = load_sampling_configs(config.sampling_configs_path) log_for_0(f"Loading checkpoint from: {args.checkpoint_path}") state, _ = load_checkpoint(args.checkpoint_path, state) state_replicated = jax_utils.replicate(state) for seed_idx, seed_val in enumerate(seed_list): if len(seed_list) > 1: log_for_0(f"\n{'#' * 70}") log_for_0(f"Seed {seed_idx + 1}/{len(seed_list)}: {seed_val}") log_for_0(f"{'#' * 70}") seed_rng = jax.random.PRNGKey(seed_val) original_output_dir = config.output_dir if len(seed_list) > 1: config.output_dir = os.path.join(original_output_dir, f"seed_{seed_val}") for sc_idx, sc in enumerate(config.sampling_configs): if len(config.sampling_configs) > 1: log_for_0(f"\n--- Sampling config {sc_idx + 1}/{len(config.sampling_configs)} ---") seed_rng, sample_rng = jax.random.split(seed_rng) common_kwargs = dict( state=state_replicated, tokenizer=tokenizer, rng=sample_rng, config=config, sampling_config=sc, batch_size=local_batch_size, num_samples=config.num_samples, ) if eval_dataset is None: test_generation_uncond(**common_kwargs) else: test_generation_cond( **common_kwargs, encoder_params=encoder_params, encoder_apply_fn=encoder_model.apply, dataset=eval_dataset, ) config.output_dir = original_output_dir log_for_0("\nEvaluation complete!") if __name__ == "__main__": main()