Fix: app.py

app.py

@@ -26,12 +26,15 @@ TEST_DATA  = Path("spins_test.npy")
 # Subprocess helpers
 # ---------------------------------------------------------------------------
 
-def _stream(command: list[str]):
-    """Run a command and yield log lines in real time."""
-    log = ["$ " + " ".join(shlex.quote(p) for p in command), ""]
-    yield "\n".join(log)
+def _stream_into(cmd: list[str], log_lines: list[str]):
+    """Run cmd, append each stdout line to log_lines, yield log_lines after each line.
+
+    stderr is merged into stdout so tracebacks are always visible.
+    Yields the joined log after every line so callers can stream updates.
+    """
+    log_lines.append("$ " + " ".join(shlex.quote(p) for p in cmd))
     proc = subprocess.Popen(
-        command,
+        cmd,
         stdout=subprocess.PIPE,
         stderr=subprocess.STDOUT,
         text=True,
@@ -39,11 +42,11 @@ def _stream(command: list[str]):
     )
     assert proc.stdout is not None
     for line in proc.stdout:
-        log.append(line.rstrip())
-        yield "\n".join(log[-300:])
+        log_lines.append(line.rstrip())
+        yield "\n".join(log_lines[-300:])
     rc = proc.wait()
-    log += ["", f"— exited {rc} —"]
-    yield "\n".join(log[-300:])
+    log_lines.append(f"[exit {rc}]")
+    yield "\n".join(log_lines[-300:])
 
 
 # ---------------------------------------------------------------------------
@@ -88,10 +91,11 @@ def run_ce(mode, epochs, batch_size, lr, max_steps):
         cmd += ["--epochs", str(int(epochs))]
         if int(max_steps) > 0:
             cmd += ["--max-train-steps", str(int(max_steps))]
-    for log in _stream(cmd):
+    log_lines: list[str] = []
+    for log in _stream_into(cmd, log_lines):
         yield log, None
     ckpt = str(CE_CKPT) if CE_CKPT.exists() else None
-    yield log, ckpt
+    yield "\n".join(log_lines[-300:]), ckpt
 
 
 # ---------------------------------------------------------------------------
@@ -115,10 +119,11 @@ def run_vi(mode, steps, batch_size, lr, warm_start):
         cmd += ["--num-steps", "3", "--log-every", "1"]
     else:
         cmd += ["--num-steps", str(int(steps))]
-    for log in _stream(cmd):
+    log_lines: list[str] = []
+    for log in _stream_into(cmd, log_lines):
         yield log, None
     ckpt = str(VI_CKPT) if VI_CKPT.exists() else None
-    yield log, ckpt
+    yield "\n".join(log_lines[-300:]), ckpt
 
 
 # ---------------------------------------------------------------------------
@@ -127,23 +132,19 @@ def run_vi(mode, steps, batch_size, lr, warm_start):
 
 def run_eval(which, num_samples, seed):
     OUT.mkdir(parents=True, exist_ok=True)
-    log_lines = []
+    log_lines: list[str] = []
 
-    def emit(msg=""):
-        log_lines.append(msg)
-        return (
-            "\n".join(log_lines[-200:]),
-            None, None,   # CE figure, VI figure
-        )
+    def current_log():
+        return "\n".join(log_lines[-300:])
 
     run_ce_ = which in ("CE", "Both")
     run_vi_ = which in ("VI", "Both")
 
     if run_ce_ and not CE_CKPT.exists():
-        yield emit("⚠  CE checkpoint not found. Run CE training first.")
+        yield "⚠  CE checkpoint not found. Run CE training first.", None, None
         return
     if run_vi_ and not VI_CKPT.exists():
-        yield emit("⚠  VI checkpoint not found. Run VI training first.")
+        yield "⚠  VI checkpoint not found. Run VI training first.", None, None
         return
 
     # ── Generate samples ───────────────────────────────────────────────────
@@ -154,7 +155,7 @@ def run_eval(which, num_samples, seed):
         if (label == "CE" and not run_ce_) or (label == "VI" and not run_vi_):
             continue
         log_lines.append(f"\n── Generating {num_samples} {label} samples ──")
-        yield "\n".join(log_lines[-200:]), None, None
+        yield current_log(), None, None
         cmd = [
             sys.executable, "sample.py",
             "--checkpoint",  str(ckpt),
@@ -162,9 +163,8 @@ def run_eval(which, num_samples, seed):
             "--output",      str(out_path),
             "--seed",        str(int(seed)),
         ]
-        for chunk in _stream(cmd):
-            log_lines[-1:] = chunk.splitlines()[-10:]
-            yield "\n".join(log_lines[-200:]), None, None
+        for log in _stream_into(cmd, log_lines):
+            yield log, None, None
 
     # ── Run eval ──────────────────────────────────────────────────────────
     for ckpt, smpl, label in [
@@ -173,31 +173,34 @@ def run_eval(which, num_samples, seed):
     ]:
         if (label == "CE" and not run_ce_) or (label == "VI" and not run_vi_):
             continue
+        if not smpl.exists():
+            log_lines.append(f"⚠  {smpl} not found — sample generation may have failed.")
+            yield current_log(), None, None
+            continue
         log_lines.append(f"\n── Evaluating {label} model ──")
-        yield "\n".join(log_lines[-200:]), None, None
+        yield current_log(), None, None
         cmd = [
             sys.executable, "eval.py",
             "--checkpoint",  str(ckpt),
             "--test-data",   str(TEST_DATA),
             "--num-samples", str(int(num_samples)),
-            "--samples-file",str(smpl),
+            "--samples-file", str(smpl),
             "--seed",        str(int(seed)),
         ]
-        for chunk in _stream(cmd):
-            log_lines[-1:] = chunk.splitlines()[-20:]
-            yield "\n".join(log_lines[-200:]), None, None
+        for log in _stream_into(cmd, log_lines):
+            yield log, None, None
 
     # ── Build figures ──────────────────────────────────────────────────────
     ce_fig = _samples_figure(CE_SMPL, "CE samples") if run_ce_ else None
     vi_fig = _samples_figure(VI_SMPL, "VI samples") if run_vi_ else None
-    yield "\n".join(log_lines[-200:]), ce_fig, vi_fig
+    yield current_log(), ce_fig, vi_fig
 
 
 # ---------------------------------------------------------------------------
 # Gradio UI
 # ---------------------------------------------------------------------------
 
-with gr.Blocks(title="Ising Transformer", theme=gr.themes.Soft()) as demo:
+with gr.Blocks(title="Ising Transformer") as demo:
     gr.Markdown(
         "# 2D Ising Transformer\n"
         "Autoregressive transformer trained on the 2D Ising model at the critical "
@@ -284,4 +287,4 @@ with gr.Blocks(title="Ising Transformer", theme=gr.themes.Soft()) as demo:
 
 
 if __name__ == "__main__":
-    demo.queue(default_concurrency_limit=1).launch()
+    demo.queue(default_concurrency_limit=1).launch(theme=gr.themes.Soft())