#!/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()