From 8fd82bf7311d2ef95ef0771db4722b1510d75b67 Mon Sep 17 00:00:00 2001
From: wassname <1103714+wassname@users.noreply.github.com>
Date: Fri, 10 Jul 2026 16:31:25 +0800
Subject: [PATCH] fit: tqdm progress bar + first-prompt trace + summary; config
 configures loguru on import

- Jacobian.fit wraps prompts in tqdm (jlens has no bar; safe since fit only
  enumerate/len's them), logs the full first prompt (special tokens on, SHOULD
  line) and a done-summary -- token-efficient-logging style, both tqdm intervals set
- config.py sets up loguru on import (compact single-char icons, routed through
  tqdm.write so bars survive), so every script/notebook importing config gets it
- notebooks drop their manual logger setup and import config in cell 1

Co-Authored-By: Claudypoo <288921227+claudypoo@users.noreply.github.com>
---
 config.py                  | 11 +++++++++++
 jsteer/jacobian.py         | 19 +++++++++++++++++--
 nbs/persona_steering.ipynb |  2 +-
 nbs/word_steering.ipynb    |  4 ++--
 scripts/fit.py             |  5 ++++-
 5 files changed, 35 insertions(+), 6 deletions(-)

diff --git a/config.py b/config.py
index 191b86c..b317218 100644
--- a/config.py
+++ b/config.py
@@ -9,6 +9,17 @@ template -- see chat_corpus for why.
 from pathlib import Path
 
 import torch
+from loguru import logger
+from tqdm.auto import tqdm
+
+# Configure loguru once, on import, so every script/notebook that imports config
+# gets the same compact format. Routed through tqdm.write so log lines don't
+# break a live progress bar (e.g. the fit bar).
+logger.remove()
+logger.add(lambda m: tqdm.write(m, end=""), colorize=True,
+           format="<level>{level.icon}</level> {message}", level="INFO")
+for _lvl, _icon in (("INFO", "I"), ("WARNING", "W"), ("ERROR", "E"), ("DEBUG", "D")):
+    logger.level(_lvl, icon=_icon)
 
 ROOT = Path(__file__).resolve().parent
 ART = ROOT / "artifacts"
diff --git a/jsteer/jacobian.py b/jsteer/jacobian.py
index 96c1db9..9353a72 100644
--- a/jsteer/jacobian.py
+++ b/jsteer/jacobian.py
@@ -151,10 +151,25 @@ class Jacobian:
         makes the fit resumable (atomic writes)."""
         lm = from_hf(model, tok, compile=compile)
         source_layers = _resolve_layers(layers, lm.n_layers)
-        lens = _jlens_fit(lm, prompts, source_layers=source_layers,
+
+        # Full trace of the first fit prompt as jlens sees it (special tokens on).
+        # SHOULD: a chat fit opens with the template's <|im_start|>user and ends at
+        # the assistant/<think> start; plain text means the template was skipped.
+        ids0 = tok(prompts[0], add_special_tokens=True).input_ids
+        logger.info(f"fit on {len(prompts)} prompts, layers={source_layers} "
+                    f"(dim_batch={dim_batch}, max_seq_len={max_seq_len})")
+        logger.info(f"FIT PROMPT[0] ({len(ids0)} tok): {tok.decode(ids0)!r}")
+
+        # jlens.fit has no progress bar; wrap prompts so we get one (fit consumes
+        # them only via enumerate/len, so this is safe). Both tqdm intervals set
+        # (token-efficient-logging) to avoid CR-spam in non-tty logs.
+        bar = tqdm(prompts, desc="fit J", mininterval=30, maxinterval=30)
+        lens = _jlens_fit(lm, bar, source_layers=source_layers,
                           dim_batch=dim_batch, max_seq_len=max_seq_len,
                           checkpoint_path=checkpoint_path)
-        return cls(lens=lens)
+        jac = cls(lens=lens)
+        logger.info(f"fit done: {jac!r}")
+        return jac
 
     def save(self, path: str) -> None:
         self.lens.save(path)      # fp16 by default; jlens-compatible file
diff --git a/nbs/persona_steering.ipynb b/nbs/persona_steering.ipynb
index 0226d02..e466a7e 100644
--- a/nbs/persona_steering.ipynb
+++ b/nbs/persona_steering.ipynb
@@ -19,7 +19,7 @@
     }
    },
    "outputs": [],
-   "source": "# demo notebook authored by Claude\nimport sys\nfrom loguru import logger\n\nlogger.remove()  # show_steer prints through loguru; route it to the cell output\nlogger.add(sys.stdout, format=\"{message}\")\n\nimport torch\nfrom transformers import AutoModelForCausalLM, AutoTokenizer\n\nfrom jsteer import Jacobian, show_steer\n\nsys.path.insert(0, \"..\")  # repo root for config.py\nimport config\n\nMODEL = \"Qwen/Qwen3.5-4B\"  # 4B-class: demo material. 0.6B degenerates too easily.\ntok = AutoTokenizer.from_pretrained(MODEL)\nmodel = AutoModelForCausalLM.from_pretrained(MODEL, dtype=torch.bfloat16).to(\"cuda\").eval()\n\n# fit-or-load the cache for THIS model. chat_corpus wraps jlens's WikiText in the\n# chat template (closer to the distribution we steer in than raw documents); the\n# lambda means WikiText is only built on a cache MISS. dim_batch=4 fits 4B on a\n# 3090; checkpoint_path makes a multi-hour 4B fit resumable if it dies.\njac = Jacobian.fit_cached(model, tok, lambda: config.chat_corpus(tok, 128),\n                          config.cache_path(MODEL), layers=(0.3, 0.9), dim_batch=4,\n                          checkpoint_path=str(config.cache_path(MODEL, \"ckpt\")))\njac"
+   "source": "# demo notebook authored by Claude\nimport sys\nsys.path.insert(0, \"..\")  # repo root for config.py\nimport config  # configures loguru on import (compact format, tqdm-safe)\n\nimport torch\nfrom transformers import AutoModelForCausalLM, AutoTokenizer\n\nfrom jsteer import Jacobian, show_steer\n\nMODEL = \"Qwen/Qwen3.5-4B\"  # 4B-class: demo material. 0.6B degenerates too easily.\ntok = AutoTokenizer.from_pretrained(MODEL)\nmodel = AutoModelForCausalLM.from_pretrained(MODEL, dtype=torch.bfloat16).to(\"cuda\").eval()\n\n# fit-or-load the cache for THIS model. chat_corpus wraps jlens's WikiText in the\n# chat template (closer to the distribution we steer in than raw documents); the\n# lambda means WikiText is only built on a cache MISS. dim_batch=4 fits 4B on a\n# 3090; checkpoint_path makes a multi-hour 4B fit resumable if it dies.\njac = Jacobian.fit_cached(model, tok, lambda: config.chat_corpus(tok, 128),\n                          config.cache_path(MODEL), layers=(0.3, 0.9), dim_batch=4,\n                          checkpoint_path=str(config.cache_path(MODEL, \"ckpt\")))\njac"
   },
   {
    "cell_type": "markdown",
diff --git a/nbs/word_steering.ipynb b/nbs/word_steering.ipynb
index 6210fc6..aa499f0 100644
--- a/nbs/word_steering.ipynb
+++ b/nbs/word_steering.ipynb
@@ -19,7 +19,7 @@
     }
    },
    "outputs": [],
-   "source": "# demo notebook authored by Claude\nimport sys\nfrom loguru import logger\n\nlogger.remove()  # show_steer prints through loguru; route it to the cell output\nlogger.add(sys.stdout, format=\"{message}\")\n\nimport torch\nfrom transformers import AutoModelForCausalLM, AutoTokenizer\n\nfrom jsteer import Jacobian, show_steer\n\nMODEL = \"Qwen/Qwen3.5-4B\"  # 4B-class: demo material. 0.6B degenerates too easily.\ntok = AutoTokenizer.from_pretrained(MODEL)\nmodel = AutoModelForCausalLM.from_pretrained(MODEL, dtype=torch.bfloat16).to(\"cuda\").eval()"
+   "source": "# demo notebook authored by Claude\nimport sys\nsys.path.insert(0, \"..\")  # repo root for config.py\nimport config  # configures loguru on import (compact format, tqdm-safe)\n\nimport torch\nfrom transformers import AutoModelForCausalLM, AutoTokenizer\n\nfrom jsteer import Jacobian, show_steer\n\nMODEL = \"Qwen/Qwen3.5-4B\"  # 4B-class: demo material. 0.6B degenerates too easily.\ntok = AutoTokenizer.from_pretrained(MODEL)\nmodel = AutoModelForCausalLM.from_pretrained(MODEL, dtype=torch.bfloat16).to(\"cuda\").eval()"
   },
   {
    "cell_type": "markdown",
@@ -40,7 +40,7 @@
     }
    },
    "outputs": [],
-   "source": "# fit-or-load: builds the cache on first run for ANY model, loads it after.\n# chat_corpus wraps jlens's WikiText in the chat template, closer to the\n# distribution we steer in (chat + <think>) than raw documents. The lambda means\n# WikiText is only built on a cache MISS. dim_batch=4 fits a 4B on a 24GB 3090;\n# checkpoint_path makes a multi-hour 4B fit resumable if it dies.\nsys.path.insert(0, \"..\")  # repo root for config.py\nimport config\n\njac = Jacobian.fit_cached(model, tok, lambda: config.chat_corpus(tok, 128),\n                          config.cache_path(MODEL), layers=(0.3, 0.9), dim_batch=4,\n                          checkpoint_path=str(config.cache_path(MODEL, \"ckpt\")))\njac"
+   "source": "# fit-or-load: builds the cache on first run for ANY model, loads it after.\n# chat_corpus wraps jlens's WikiText in the chat template, closer to the\n# distribution we steer in (chat + <think>) than raw documents. The lambda means\n# WikiText is only built on a cache MISS. dim_batch=4 fits a 4B on a 24GB 3090;\n# checkpoint_path makes a multi-hour 4B fit resumable if it dies.\njac = Jacobian.fit_cached(model, tok, lambda: config.chat_corpus(tok, 128),\n                          config.cache_path(MODEL), layers=(0.3, 0.9), dim_batch=4,\n                          checkpoint_path=str(config.cache_path(MODEL, \"ckpt\")))\njac"
   },
   {
    "cell_type": "markdown",
diff --git a/scripts/fit.py b/scripts/fit.py
index 34865f4..5aaba89 100644
--- a/scripts/fit.py
+++ b/scripts/fit.py
@@ -49,7 +49,10 @@ def main() -> None:
                               layers=tuple(args.layers), dim_batch=args.dim_batch,
                               max_seq_len=args.max_seq_len,
                               checkpoint_path=str(config.cache_path(args.model, "ckpt")))
-    logger.info(f"{jac!r} -> {out}  ({(time.monotonic() - t0) / 60:.1f} min)")
+    # BLUF summary: what to read first.
+    logger.info(f"DONE fit -> {out}")
+    logger.info(f"     {jac!r}  |  {args.n_prompts} prompts, dim_batch={args.dim_batch}, "
+                f"{(time.monotonic() - t0) / 60:.1f} min")
 
 
 if __name__ == "__main__":