reference/code/ELF/src/generation.py

import time
import os
import json
import itertools
import wandb
from tqdm import tqdm

import jax
import jax.numpy as jnp
import numpy as np
from functools import partial
from typing import Any, Dict
from jax.random import PRNGKey

from utils.logging_utils import log_for_0
from utils.checkpoint_utils import upload_output_dir_to_hf
from utils.train_utils import TrainState
from configs.config import Config, SamplingConfig

from utils.metrics_utils import Metrics as PPLMetrics, compute_bleu, compute_rouge
from jax.experimental.multihost_utils import sync_global_devices
from utils.data_utils import get_dataloader, get_pad_token_id
from utils.encoder_utils import encode_text
from utils.generation_utils import (
    mask_after_eos, shift_left,
    _setup_generation, _make_pmap_pair, _shard_timesteps, _shard_noise,
    _build_run_name,
)


# ============================================
# Generation Helper
# ============================================
def run_generation(
    state: Any,
    encoder_params: Dict,
    encoder_apply_fn,
    eval_dataset,
    tokenizer,
    config,
    rng: PRNGKey,
    local_batch_size: int,
) -> PRNGKey:
    """Run test generation."""
    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)} ---")
        rng, sample_rng = jax.random.split(rng)
        common_kwargs = dict(
            state=state,
            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_apply_fn,
                dataset=eval_dataset,
            )

    return rng


# ============================================
# Evaluation / Testing Utilities
# ============================================
def test_generation_uncond(
    state: TrainState,
    tokenizer,
    rng: PRNGKey,
    config: Config,
    sampling_config: SamplingConfig,
    num_samples: int = 64,
    batch_size: int = 64,
):
    """Test unconditional generation (multi-device pmap)."""
    sampling_method = sampling_config.sampling_method
    time_schedule = sampling_config.time_schedule
    log_for_0(f"Config: {sampling_config}")

    (state_unreplicated, model_apply_fn, model_params_replicated,
     d_model, num_local_devices, effective_batch_size) = _setup_generation(
        state, config, batch_size, "              UNCONDITIONAL GENERATION EXAMPLES",
    )

    pad_token_id = get_pad_token_id(tokenizer)
    eos_token_id = tokenizer.eos_token_id if tokenizer.eos_token_id is not None else 1

    cfg_list = [1]
    steps_list = sampling_config.num_sampling_steps
    self_cond_cfg_scales_list = sampling_config.self_cond_cfg_scales
    wandb_tables = {}
    ppl_metrics = None
    if config.online_eval:
        ppl_metrics = PPLMetrics(
            gen_ppl_eval_model_name_or_path=config.eval_ppl_model,
            eval_ppl_batch_size=config.eval_ppl_batch_size,
            eval_context_size=config.eval_ppl_max_length,
        )

    for num_sampling_steps, cfg_scale, self_cond_cfg_scale in itertools.product(
        steps_list, cfg_list, self_cond_cfg_scales_list
    ):
        log_for_0(f"\n--- Method: {sampling_method}, Steps: {num_sampling_steps}, "
                  f"CFG Scale: {cfg_scale}, SC-CFG: {self_cond_cfg_scale} ---")

        p_generate, p_decode_ids = _make_pmap_pair(
            model_apply_fn, config, sampling_config, cfg_scale, self_cond_cfg_scale,
        )

        all_generated = []
        generation_time = 0.0
        decode_time = 0.0
        num_batches = (num_samples + effective_batch_size - 1) // effective_batch_size
        samples_processed = 0

        for batch_idx in tqdm(range(num_batches), desc="Generating samples"):
            if samples_processed >= num_samples:
                break

            current_total_batch = min(effective_batch_size, num_samples - samples_processed)
            current_total_batch = (
                (current_total_batch + num_local_devices - 1) // num_local_devices
            ) * num_local_devices
            current_per_device = current_total_batch // num_local_devices

            batch_rng = jax.random.fold_in(rng, batch_idx * jax.process_count() + jax.process_index())
            noise_rng, t_rng = jax.random.split(batch_rng)
            device_rngs = jax.random.split(noise_rng, num_local_devices)

            t_steps_sharded = _shard_timesteps(
                t_rng, num_local_devices, num_sampling_steps, time_schedule, config,
            )
            z_sharded = _shard_noise(
                device_rngs, num_local_devices, current_per_device,
                config.max_length, d_model, config.denoiser_noise_scale,
            )

            gen_start = time.time()
            latent_sharded = p_generate(
                model_params=model_params_replicated, rng=device_rngs,
                z=z_sharded, t_steps=t_steps_sharded,
                cond_seq=None, cond_seq_mask=None,
            )
            latent_sharded.block_until_ready()
            generation_time += time.time() - gen_start

            dec_start = time.time()
            t_final_sharded = t_steps_sharded[:, -1]
            predicted_ids_sharded = p_decode_ids(
                z=latent_sharded, model_params=model_params_replicated, t_final_val=t_final_sharded,
            )
            predicted_ids_sharded.block_until_ready()
            decode_time += time.time() - dec_start

            predicted_ids = predicted_ids_sharded.reshape(-1, predicted_ids_sharded.shape[-1])
            predicted_ids = mask_after_eos(predicted_ids, eos_token_id=eos_token_id, pad_token_id=pad_token_id)

            for i in range(predicted_ids.shape[0]):
                if samples_processed >= num_samples:
                    break
                text = tokenizer.decode(np.array(predicted_ids[i]), skip_special_tokens=True)
                all_generated.append((samples_processed, text))
                samples_processed += 1

        log_for_0(f"Generation: {generation_time:.2f}s ({num_sampling_steps} steps) | Decode: {decode_time:.2f}s")
        log_for_0("-" * 70)

        epoch_val = int(state_unreplicated.epoch)
        step_val = int(state_unreplicated.step)
        name = _build_run_name(
            sampling_method, num_sampling_steps, cfg_scale, self_cond_cfg_scale,
            time_schedule, getattr(sampling_config, "sde_gamma", 0.0), suffix="uncond",
        )

        # Dummy op so all processes stay in sync before rank-0-only file write / PPL
        jax.random.normal(jax.random.PRNGKey(0))

        # Rank 0 writes one file to the same folder (no shards).
        out_path = os.path.join(config.output_dir, name, f"all_generated_{epoch_val}_{step_val}.jsonl")
        if jax.process_index() == 0:
            os.makedirs(os.path.join(config.output_dir, name), exist_ok=True)
            with open(out_path, "w", encoding="utf-8") as f:
                for tid, gen in all_generated:
                    f.write(json.dumps({"id": tid, "generated": gen}, ensure_ascii=False) + "\n")
            log_for_0(f"Saved {len(all_generated)} generated texts to {out_path}")
            upload_output_dir_to_hf(config.output_dir, config.hf_repo_id, reason="generation")

        sync_global_devices("after_save")
        ppl_results = None
        if config.online_eval:
            log_for_0("\n" + "=" * 70)
            log_for_0("              PPL EVALUATION (all hosts, logged on rank 0)")
            log_for_0("=" * 70)
            ppl_metrics.reset()

            # All hosts read the exact same samples so JAX launches stay aligned.
            with open(out_path, "r", encoding="utf-8") as f:
                text_samples = [json.loads(line)["generated"] for line in f]

            nonempty_samples = [s for s in text_samples if isinstance(s, str) and s.strip()]
            skipped = len(text_samples) - len(nonempty_samples)
            if skipped > 0:
                log_for_0(f"PPL eval: skipped {skipped} empty samples")
            if not nonempty_samples:
                log_for_0("PPL eval: all samples empty; skipping perplexity computation")
            else:
                ppl_results = ppl_metrics.record_generative_perplexity(
                    text_samples=nonempty_samples,
                    max_length=config.eval_ppl_max_length,
                    retokenize=True,
                )
                mean_ppl = ppl_results["ppl"]
                log_for_0(f"Perplexity: {mean_ppl:.4f}")
                log_for_0(f"Mean Entropy: {ppl_results['mean_entropy']:.4f}")
            log_for_0("=" * 70 + "\n")

        sync_global_devices("after_ppl")

        if jax.process_index() == 0:
            if ppl_results is not None:
                mean_ppl = ppl_results["ppl"]
                metrics_line = {
                    "epoch": epoch_val, "step": step_val,
                    "ppl": mean_ppl, "mean_entropy": ppl_results["mean_entropy"],
                }
                with open(os.path.join(config.output_dir, name, "metrics.jsonl"), "a", encoding="utf-8") as f:
                    f.write(json.dumps(metrics_line, ensure_ascii=False) + "\n")
                upload_output_dir_to_hf(config.output_dir, config.hf_repo_id, reason="generation metrics")

            if config.use_wandb:
                table = wandb.Table(columns=["sample_id", "text"])
                for tid, gen in all_generated[:min(10, len(all_generated))]:
                    table.add_data(tid, gen)
                wandb_tables[f"generated_samples_uncond_steps{num_sampling_steps}_cfg{cfg_scale}"] = table
                if ppl_results is not None:
                    wandb_tables.update({
                        f"generation/{name}/ppl": mean_ppl,
                        f"generation/{name}/mean_entropy": ppl_results["mean_entropy"],
                    })

    if jax.process_index() == 0 and config.use_wandb and wandb_tables:
        try:
            wandb.log(wandb_tables)
        except Exception as e:
            log_for_0(f"Warning: wandb.log failed: {e}")

    log_for_0("=" * 70 + "\n")


def test_generation_cond(
    state: TrainState,
    encoder_params: Dict,
    encoder_apply_fn,
    tokenizer,
    rng: PRNGKey,
    config: Config,
    sampling_config: SamplingConfig,
    dataset,
    num_samples: int = 64,
    batch_size: int = 64,
):
    """Test conditional generation (multi-device pmap)."""
    sampling_method = sampling_config.sampling_method
    time_schedule = sampling_config.time_schedule
    log_for_0(f"Config: {sampling_config}")

    (state_unreplicated, model_apply_fn, model_params_replicated,
     d_model, num_local_devices, effective_batch_size) = _setup_generation(
        state, config, batch_size, "              CONDITIONAL GENERATION EXAMPLES",
    )

    encode_latent_mean, encode_latent_std = config.latent_mean, config.latent_std

    pad_token_id = get_pad_token_id(tokenizer, config.pad_token)
    eos_token_id = tokenizer.eos_token_id

    dataloader = get_dataloader(
        dataset, batch_size=effective_batch_size,
        shuffle=False, num_workers=0, drop_last=False,
        max_seq_length=config.max_length, pad_token_id=pad_token_id,
        max_input_seq_length=config.max_input_length, distributed=False,
    )

    p_encode = jax.pmap(
        partial(
            encode_text, encoder_apply_fn=encoder_apply_fn,
            latent_mean=encode_latent_mean, latent_std=encode_latent_std,
        ),
        axis_name="batch",
    )

    wandb_tables = {}
    cfg_list = sampling_config.cfgs
    steps_list = sampling_config.num_sampling_steps
    self_cond_cfg_scales_list = sampling_config.self_cond_cfg_scales

    for num_sampling_steps, cfg_scale, self_cond_cfg_scale in itertools.product(
        steps_list, cfg_list, self_cond_cfg_scales_list
    ):
        log_for_0(f"\n--- Steps: {num_sampling_steps}, CFG Scale: {cfg_scale}, "
                  f"SC-CFG: {self_cond_cfg_scale} ---")

        p_generate, p_decode_ids = _make_pmap_pair(
            model_apply_fn, config, sampling_config, cfg_scale, self_cond_cfg_scale,
        )

        all_generated = []  # (id, original_text, generated_text, context_text)
        generation_time = 0.0
        decode_time = 0.0
        samples_processed = 0

        for batch_idx, batch in enumerate(dataloader):
            if samples_processed >= num_samples:
                break

            batch_size_current = batch["input_ids"].shape[0]

            if batch_size_current % num_local_devices != 0:
                pad_size = num_local_devices - (batch_size_current % num_local_devices)
                for key in batch:
                    if isinstance(batch[key], np.ndarray):
                        pad_arr = np.zeros((pad_size,) + batch[key].shape[1:], dtype=batch[key].dtype)
                        batch[key] = np.concatenate([batch[key], pad_arr], axis=0)
                    elif isinstance(batch[key], list):
                        batch[key] = batch[key] + [""] * pad_size

            actual_batch_size = batch["input_ids"].shape[0]
            current_per_device = actual_batch_size // num_local_devices

            batch_rng = jax.random.fold_in(rng, batch_idx * jax.process_count() + jax.process_index())
            noise_rng, t_rng = jax.random.split(batch_rng)
            device_rngs = jax.random.split(noise_rng, num_local_devices)

            t_steps_sharded = _shard_timesteps(
                t_rng, num_local_devices, num_sampling_steps, time_schedule, config,
            )

            # Encode full input sequence using encoder_attention_mask so cond tokens
            # only attend to cond tokens
            max_length = config.max_length
            input_ids = jnp.array(batch["input_ids"])
            encoder_attention_mask = jnp.array(batch["encoder_attention_mask"])
            cond_seq_mask_arr = jnp.array(batch["cond_seq_mask"])

            input_ids_sharded = input_ids.reshape(num_local_devices, current_per_device, -1)
            encoder_attention_mask_sharded = encoder_attention_mask.reshape(
                num_local_devices, current_per_device, max_length, max_length,
            )

            cond_seq_sharded = p_encode(
                input_ids=input_ids_sharded,
                attention_mask=encoder_attention_mask_sharded,
                encoder_params=encoder_params,
            )
            cond_seq_mask_sharded = cond_seq_mask_arr.reshape(num_local_devices, current_per_device, -1)

            z_sharded = _shard_noise(
                device_rngs, num_local_devices, current_per_device,
                max_length, d_model, config.denoiser_noise_scale,
            )

            gen_start = time.time()
            latent_sharded = p_generate(
                model_params=model_params_replicated, rng=device_rngs,
                z=z_sharded, t_steps=t_steps_sharded,
                cond_seq=cond_seq_sharded, cond_seq_mask=cond_seq_mask_sharded,
            )
            latent_sharded.block_until_ready()
            generation_time += time.time() - gen_start

            gen_length = config.max_length - config.max_input_length
            cond_len_per_sample = cond_seq_mask_arr.astype(jnp.int32).sum(axis=1)

            dec_start = time.time()
            t_final_sharded = t_steps_sharded[:, -1]
            predicted_ids_sharded = p_decode_ids(
                z=latent_sharded, model_params=model_params_replicated, t_final_val=t_final_sharded,
            )
            predicted_ids_sharded.block_until_ready()
            predicted_ids = predicted_ids_sharded.reshape(-1, predicted_ids_sharded.shape[-1])
            # Strip cond prefix; prefix tokens may contain EOS so mask only after stripping.
            predicted_ids = shift_left(predicted_ids, cond_len_per_sample, 0)[:, :gen_length]
            predicted_ids = mask_after_eos(predicted_ids, eos_token_id=eos_token_id, pad_token_id=pad_token_id)
            decode_time += time.time() - dec_start

            original_texts = [batch["target"][i] for i in range(actual_batch_size)]
            context_texts = [batch["input"][i] for i in range(actual_batch_size)]

            for i in range(min(batch_size_current, actual_batch_size)):
                if samples_processed >= num_samples:
                    break
                text = tokenizer.decode(np.array(predicted_ids[i]), skip_special_tokens=True)
                all_generated.append((samples_processed, original_texts[i], text, context_texts[i]))
                samples_processed += 1

        log_for_0(f"Generation: {generation_time:.2f}s ({num_sampling_steps} steps) | Decode: {decode_time:.2f}s")
        log_for_0("-" * 70)

        epoch_val = int(state_unreplicated.epoch)
        step_val = int(state_unreplicated.step)
        name = _build_run_name(
            sampling_method, num_sampling_steps, cfg_scale, self_cond_cfg_scale,
            time_schedule, getattr(sampling_config, "sde_gamma", 0.0), suffix="cond",
        )

        if jax.process_index() == 0:
            os.makedirs(os.path.join(config.output_dir, name), exist_ok=True)
            out_path = os.path.join(config.output_dir, name, f"all_generated_{epoch_val}_{step_val}.jsonl")
            with open(out_path, "w", encoding="utf-8") as f:
                for tid, orig, gen, ctx in all_generated:
                    f.write(json.dumps({"id": tid, "generated": gen}, ensure_ascii=False) + "\n")
            log_for_0(f"Saved {len(all_generated)} generated texts to {out_path}")
            upload_output_dir_to_hf(config.output_dir, config.hf_repo_id, reason="generation")

            cond_eval_results = None
            if config.online_eval and all_generated:
                hypotheses = [gen for _, _, gen, _ in all_generated]
                references = [orig for _, orig, _, _ in all_generated]
                bleu_score = compute_bleu(hypotheses, references)
                rouge_scores = compute_rouge(hypotheses, references)
                cond_eval_results = {"bleu": bleu_score, **rouge_scores}
                log_for_0(
                    f"BLEU: {bleu_score:.2f}  ROUGE-1: {rouge_scores['rouge1']:.2f}  "
                    f"ROUGE-2: {rouge_scores['rouge2']:.2f}  ROUGE-L: {rouge_scores['rougeL']:.2f}"
                )

            if config.use_wandb:
                table = wandb.Table(columns=["sample_id", "context", "original", "generated"])
                for tid, orig, gen, ctx in all_generated[:min(10, len(all_generated))]:
                    table.add_data(tid, ctx, orig, gen)
                wandb_tables[f"generated_samples_cond_steps{num_sampling_steps}_cfg{cfg_scale}"] = table
                if cond_eval_results is not None:
                    wandb_tables.update({
                        f"generation/{name}/bleu": cond_eval_results["bleu"],
                        f"generation/{name}/rouge1": cond_eval_results["rouge1"],
                        f"generation/{name}/rouge2": cond_eval_results["rouge2"],
                        f"generation/{name}/rougeL": cond_eval_results["rougeL"],
                    })
            if cond_eval_results is not None:
                metrics_line = {"epoch": epoch_val, "step": step_val, **cond_eval_results}
                with open(os.path.join(config.output_dir, name, "metrics.jsonl"), "a", encoding="utf-8") as f:
                    f.write(json.dumps(metrics_line, ensure_ascii=False) + "\n")
                upload_output_dir_to_hf(config.output_dir, config.hf_repo_id, reason="generation metrics")

    if jax.process_index() == 0 and config.use_wandb and wandb_tables:
        try:
            wandb.log(wandb_tables)
        except Exception as e:
            log_for_0(f"Warning: wandb.log failed: {e}")

    log_for_0("=" * 70 + "\n")