tools/plot_losses.py

#!/usr/bin/env python3
"""Continuously render train/eval loss PNGs from the current training log.

This intentionally uses Pillow instead of matplotlib so it works with the
package's current pinned environment.
"""

from __future__ import annotations

import argparse
import ast
import math
import os
import re
import time
from pathlib import Path
from typing import Iterable, List, Sequence, Tuple

from PIL import Image, ImageDraw, ImageFont


ANSI_RE = re.compile(r"\x1b\[[0-9;]*[A-Za-z]")


def clean_text(text: str) -> str:
    return ANSI_RE.sub("", text.replace("\r", "\n"))


def iter_inline_dicts(text: str) -> Iterable[dict]:
    for line in clean_text(text).splitlines():
        if "loss" not in line:
            continue
        for match in re.finditer(r"\{[^{}]*\}", line):
            raw = match.group(0)
            if "loss" not in raw:
                continue
            try:
                obj = ast.literal_eval(raw)
            except Exception:
                continue
            if isinstance(obj, dict):
                yield obj


def latest_run_log_region(text: str) -> str:
    markers = ("# STAGE START: CONFIGURATION", "# CONFIGURATION")
    idx = max(text.rfind(marker) for marker in markers)
    return text[idx:] if idx >= 0 else text


def parse_points(log_path: Path, logging_steps: int, save_eval_steps: int) -> Tuple[List[Tuple[int, float]], List[Tuple[int, float]]]:
    if not log_path.exists():
        return [], []
    text = latest_run_log_region(log_path.read_text(encoding="utf-8", errors="replace"))
    train: List[Tuple[int, float]] = []
    evals: List[Tuple[int, float]] = []
    for obj in iter_inline_dicts(text):
        if "loss" in obj and "eval_loss" not in obj:
            try:
                value = float(obj["loss"])
            except Exception:
                continue
            step = int(obj.get("step") or len(train) * logging_steps + logging_steps)
            if math.isfinite(value):
                train.append((step, value))
        if "eval_loss" in obj:
            try:
                value = float(obj["eval_loss"])
            except Exception:
                continue
            step = int(
                obj.get("eval_global_step")
                or obj.get("global_step")
                or len(evals) * save_eval_steps + save_eval_steps
            )
            if math.isfinite(value):
                evals.append((step, value))
    # Trainer emits the same evaluation once through its log callback and once
    # through the explicit choice-metrics audit line; keep one point per step.
    return list(dict(train).items()), list(dict(evals).items())


def nice_bounds(values: Sequence[float]) -> Tuple[float, float]:
    if not values:
        return 0.0, 1.0
    lo = min(values)
    hi = max(values)
    if lo == hi:
        pad = max(abs(lo) * 0.1, 0.5)
        return lo - pad, hi + pad
    pad = (hi - lo) * 0.12
    return lo - pad, hi + pad


def draw_plot(points: Sequence[Tuple[int, float]], out_path: Path, title: str, ylabel: str) -> None:
    width, height = 1200, 720
    left, right, top, bottom = 95, 45, 65, 90
    img = Image.new("RGB", (width, height), "white")
    draw = ImageDraw.Draw(img)
    font = ImageFont.load_default()
    title_font = ImageFont.load_default()

    plot_w = width - left - right
    plot_h = height - top - bottom
    axis = (40, 40, 40)
    grid = (225, 225, 225)
    line = (30, 105, 210)
    text = (20, 20, 20)

    draw.text((left, 25), title, fill=text, font=title_font)
    draw.rectangle((left, top, left + plot_w, top + plot_h), outline=axis, width=2)

    if not points:
        msg = "No points yet. Waiting for Trainer logging/evaluation."
        draw.text((left + 25, top + plot_h // 2), msg, fill=(120, 120, 120), font=font)
        draw.text((left, height - 40), f"updated: {time.strftime('%Y-%m-%d %H:%M:%S UTC', time.gmtime())}", fill=(90, 90, 90), font=font)
        out_path.parent.mkdir(parents=True, exist_ok=True)
        tmp = out_path.with_suffix(out_path.suffix + ".tmp")
        img.save(tmp, format="PNG")
        os.replace(tmp, out_path)
        return

    xs = [p[0] for p in points]
    ys = [p[1] for p in points]
    xmin, xmax = min(xs), max(xs)
    if xmin == xmax:
        xmin = max(0, xmin - 1)
        xmax += 1
    ymin, ymax = nice_bounds(ys)

    def sx(x: float) -> float:
        return left + (x - xmin) / (xmax - xmin) * plot_w

    def sy(y: float) -> float:
        return top + plot_h - (y - ymin) / (ymax - ymin) * plot_h

    for i in range(6):
        y = top + i * plot_h / 5
        draw.line((left, y, left + plot_w, y), fill=grid)
        val = ymax - i * (ymax - ymin) / 5
        draw.text((10, y - 7), f"{val:.4g}", fill=text, font=font)
    for i in range(6):
        x = left + i * plot_w / 5
        draw.line((x, top, x, top + plot_h), fill=grid)
        val = int(round(xmin + i * (xmax - xmin) / 5))
        draw.text((x - 18, top + plot_h + 12), str(val), fill=text, font=font)

    coords = [(sx(x), sy(y)) for x, y in points]
    if len(coords) == 1:
        x, y = coords[0]
        draw.ellipse((x - 4, y - 4, x + 4, y + 4), fill=line)
    else:
        draw.line(coords, fill=line, width=3)
        for x, y in coords[-20:]:
            draw.ellipse((x - 3, y - 3, x + 3, y + 3), fill=line)

    last_step, last_loss = points[-1]
    draw.text((left, height - 65), "optimizer/global step", fill=text, font=font)
    draw.text((8, top - 24), ylabel, fill=text, font=font)
    draw.text(
        (left, height - 40),
        f"points={len(points)}  latest_step={last_step}  latest_loss={last_loss:.6g}  updated={time.strftime('%Y-%m-%d %H:%M:%S UTC', time.gmtime())}",
        fill=(70, 70, 70),
        font=font,
    )

    out_path.parent.mkdir(parents=True, exist_ok=True)
    tmp = out_path.with_suffix(out_path.suffix + ".tmp")
    img.save(tmp, format="PNG")
    os.replace(tmp, out_path)


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--log", default="logs/train_tmux.log")
    parser.add_argument("--out-dir", default="runs/DRU-RE-Yehia/plots")
    parser.add_argument("--interval", type=int, default=20)
    parser.add_argument("--logging-steps", type=int, default=20)
    parser.add_argument("--save-eval-steps", type=int, default=250)
    parser.add_argument("--once", action="store_true")
    args = parser.parse_args()

    log_path = Path(args.log)
    out_dir = Path(args.out_dir)
    while True:
        train, evals = parse_points(log_path, args.logging_steps, args.save_eval_steps)
        draw_plot(train, out_dir / "train_loss.png", "Train loss", "loss")
        draw_plot(
            evals,
            out_dir / "eval_loss.png",
            "Validation decision loss (NLL)",
            "eval_loss",
        )
        summary = {
            "train_points": len(train),
            "eval_points": len(evals),
            "latest_train": train[-1] if train else None,
            "latest_eval": evals[-1] if evals else None,
            "updated_utc": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
        }
        (out_dir / "loss_plot_summary.json").write_text(
            __import__("json").dumps(summary, ensure_ascii=False, indent=2) + "\n",
            encoding="utf-8",
        )
        if args.once:
            break
        time.sleep(max(5, args.interval))


if __name__ == "__main__":
    main()