From e90de299f4a3486f68381cbc4427b3c8ab7360b5 Mon Sep 17 00:00:00 2001
From: wassname <1103714+wassname@users.noreply.github.com>
Date: Sun, 12 Jul 2026 10:38:56 +0800
Subject: [PATCH] tqdm: disable single-batch h_bar/vjp bars, equal intervals
 (token-efficient-logging)

Single-batch persona lists made instant 1/1 bars that print a completion
line regardless of interval; disable those. Both intervals 120 for the
multi-batch fit/pullback case.

Co-Authored-By: Claudypoo <288921227+claudypoo@users.noreply.github.com>
---
 jsteer/jacobian.py | 4 +++-
 jsteer/vjp.py      | 2 +-
 2 files changed, 4 insertions(+), 2 deletions(-)

diff --git a/jsteer/jacobian.py b/jsteer/jacobian.py
index 646af54..44cc4c2 100644
--- a/jsteer/jacobian.py
+++ b/jsteer/jacobian.py
@@ -124,8 +124,10 @@ def _h_bar_final(model, tok, prompts: list[str], *, batch_size: int = 8,
     lm = from_hf(model, tok)                                     # layout detection only
     target_layer = lm.n_layers - 1
     acc, n = None, 0
+    # single-batch persona lists (<= batch_size) make a useless 1/1 bar that still prints
+    # a completion line; disable those. Both intervals equal for the multi-batch case.
     for i in tqdm(range(0, len(prompts), batch_size), desc=f"h_bar {label}",
-                  mininterval=30, maxinterval=60):
+                  disable=len(prompts) <= batch_size, mininterval=120, maxinterval=120):
         batch = prompts[i:i + batch_size]
         enc = tok(batch, return_tensors="pt", padding=True, truncation=True,
                   max_length=max_length, padding_side="right").to(model.device)
diff --git a/jsteer/vjp.py b/jsteer/vjp.py
index 595ba31..c5603c2 100644
--- a/jsteer/vjp.py
+++ b/jsteer/vjp.py
@@ -57,7 +57,7 @@ def pullback_vjp(model, tok, prompts: list[str], layers, cotangent: Tensor, *,
     G = {l: torch.zeros(d, dtype=torch.float32, device=model.device) for l in layers}
     count = 0
     for i in tqdm(range(0, len(prompts), batch_size), desc="pullback_vjp",
-                  mininterval=30, maxinterval=60):
+                  disable=len(prompts) <= batch_size, mininterval=120, maxinterval=120):
         batch = prompts[i:i + batch_size]
         enc = tok(batch, return_tensors="pt", padding=True, truncation=True,
                   max_length=max_length, padding_side="right").to(model.device)