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)