LTA_openwebtext_dualt/scripts/eval_dirichlet20k_entropy_push_20260508.py

#!/usr/bin/env python3
"""Focused decode sweep around the best Dirichlet-20k entropy/PPL point."""

from __future__ import annotations

import importlib.util
import sys
from pathlib import Path


BASE = Path(__file__).with_name("eval_c1024_decode_sweep_20260507.py")
spec = importlib.util.spec_from_file_location("eval_c1024_decode_sweep_20260507", BASE)
if spec is None or spec.loader is None:
    raise RuntimeError(f"cannot import {BASE}")

base = importlib.util.module_from_spec(spec)
sys.modules[spec.name] = base
spec.loader.exec_module(base)


def fmt_float(x: float) -> str:
    return str(x).replace(".", "p")


def entropy_push_configs() -> list[base.DecodeConfig]:
    configs: list[base.DecodeConfig] = []

    # The previous best was Cmax=16, temp=1.3. Sweep tightly around it.
    cmax_values = [2, 4, 6, 8, 12, 16, 24, 32, 48]
    temps = [1.10, 1.20, 1.25, 1.30, 1.35, 1.40, 1.45]
    for cmax in cmax_values:
        for temp in temps:
            configs.append(
                base.DecodeConfig(
                    f"post_sem1_blend_c{cmax}_t{fmt_float(temp)}",
                    "post",
                    1.0,
                    1.0,
                    "blend",
                    endpoint_temp=temp,
                    concentration_max=float(cmax),
                )
            )

    # A few schedule probes at the known good low-C region.
    for cmax in [8, 12, 16, 24]:
        for support_power in [0.8, 1.2]:
            for temp in [1.25, 1.35]:
                configs.append(
                    base.DecodeConfig(
                        f"post_sem1_blend_c{cmax}_sp{fmt_float(support_power)}_t{fmt_float(temp)}",
                        "post",
                        support_power,
                        1.0,
                        "blend",
                        endpoint_temp=temp,
                        concentration_max=float(cmax),
                    )
                )

    # Semantic schedule probes: delayed semantic endpoint can increase diversity,
    # early semantic endpoint can stabilize PPL. Keep this small.
    for cmax in [12, 16, 24]:
        for semantic_power in [0.8, 1.3]:
            for temp in [1.30, 1.40]:
                configs.append(
                    base.DecodeConfig(
                        f"post_sem{fmt_float(semantic_power)}_blend_c{cmax}_t{fmt_float(temp)}",
                        "post",
                        1.0,
                        semantic_power,
                        "blend",
                        endpoint_temp=temp,
                        concentration_max=float(cmax),
                    )
                )

    return configs


base.default_configs = entropy_push_configs
base.main()