Add MCP control server, live API-generation view, and repetition penalty

- scripts/jwash_mcp.py: MCP server (FastMCP/stdio), an HTTP client of the running
  J-Wash server so an external LLM can drive an already-loaded model. Tools:
  generate, find_token, list_layers, scale_token/replace_token (pure-weights,
  layers required), set_intensity, list_edits/reset_edits.
- api/app.py + ui: /api/generate now records the last exchange (surfaced in
  /api/status) and broadcasts on /ws when done; the UI shows a "generated via
  API/MCP" panel at the top of the chat and an Options "API monitor" toggle that
  swaps the 2s status poll for an event-driven refresh.
- sampling: repetition penalty (default 1.0, applied in model_manager._sample),
  exposed as a "rep" field in the chat controls; intentionally not exposed
  through the MCP.
- ui: remove the redundant "md" chat toggle (already available in Options).

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

api/app.py

@@ -88,6 +88,9 @@ fit_manager.on_progress = _broadcast_fit
 # concurrent HF downloads: one state per repo_id
 _downloads = {}
 _downloads_lock = threading.Lock()
+# last synchronous /api/generate exchange (CLI/MCP): surfaced in /api/status so
+# the UI can show what an external client generates, without persisting it
+_last_generation = None
 
 
 class LoadRequest(BaseModel):
@@ -203,6 +206,7 @@ def api_status():
         "interventions": interventions.summary(),
         "interventions_scale": interventions.global_scale,
         "interventions_mode": interventions.mode,
+        "last_generation": _last_generation,
     }
 
 
@@ -740,6 +744,23 @@ async def api_generate_sync(req: GenerateSyncRequest):
         raise HTTPException(500, str(exc))
     if done.get("error"):
         raise HTTPException(500, done["error"])
+    global _last_generation
+    _last_generation = {
+        "n": (_last_generation or {}).get("n", 0) + 1,
+        "prompt": next(
+            (m.get("content", "") for m in reversed(req.messages) if m.get("role") == "user"),
+            "",
+        ),
+        "text": done.get("text", ""),
+        "stats": done.get("stats"),
+    }
+    # nudge any watching UI to refresh once this API generation is done (used by
+    # the "API monitor" mode, which drops the 2s status poll for event-driven refresh)
+    for ws in list(_ws_locks):
+        try:
+            await _ws_send(ws, json.dumps({"type": "api_generation"}))
+        except Exception:
+            pass
     return {"text": done.get("text", ""), "stats": done.get("stats")}
 
 

config.py

@@ -30,6 +30,7 @@ DEFAULT_SAMPLING = {
     "top_k": 40,
     "max_tokens": 512,
     "seed": -1,  # -1 = random
+    "repetition_penalty": 1.0,  # 1.0 = off
 }
 
 

core/model_manager.py

@@ -418,7 +418,10 @@ def _resolve_revision(source):
     return None
 
 
-def _sample(logits, temperature, top_p, top_k, generator=None):
+def _sample(logits, temperature, top_p, top_k, generator=None, penalty=1.0, penalty_ids=None):
+    if penalty != 1.0 and penalty_ids is not None and penalty_ids.numel():
+        score = logits[penalty_ids]
+        logits[penalty_ids] = torch.where(score > 0, score / penalty, score * penalty)
     if temperature <= 0:
         return int(logits.argmax())
     probs = torch.softmax(logits / temperature, -1)
@@ -621,6 +624,11 @@ class ModelManager:
             top_k = int(sampling.get("top_k", config.DEFAULT_SAMPLING["top_k"]))
             max_tokens = int(sampling.get("max_tokens", config.DEFAULT_SAMPLING["max_tokens"]))
             seed = int(sampling.get("seed", config.DEFAULT_SAMPLING["seed"]))
+            # repetition penalty (HF-style, over prompt + generated); 1.0 = off.
+            # Deliberately NOT exposed through the MCP server.
+            repetition_penalty = float(
+                sampling.get("repetition_penalty", config.DEFAULT_SAMPLING["repetition_penalty"])
+            )
             # base model with a generic template: the model has no notion of dialogue
             # turns, so we cut as soon as it reopens one (User: or a new Assistant:)
             stop_seqs = (
@@ -676,13 +684,19 @@ class ModelManager:
             reply_ids = []
             emitted = ""
             started = time.perf_counter()
+            penalty_ids = input_ids[0].to(logits.device) if repetition_penalty != 1.0 else None
             for _ in range(max_tokens):
                 if stop_event.is_set():
                     break
-                next_id = _sample(logits[0].float(), temperature, top_p, top_k, generator)
+                next_id = _sample(logits[0].float(), temperature, top_p, top_k, generator,
+                                  repetition_penalty, penalty_ids)
                 if next_id in eos_ids:
                     break
                 reply_ids.append(next_id)
+                if penalty_ids is not None:
+                    penalty_ids = torch.cat(
+                        [penalty_ids, torch.tensor([next_id], device=penalty_ids.device)]
+                    )
                 text = tokenizer.decode(reply_ids, skip_special_tokens=True)
                 stop_hit = next((s for s in stop_seqs if s in text), None)
                 if stop_hit:

scripts/jwash_mcp.py

@@ -0,0 +1,299 @@
+"""MCP server for J-Wash — let an external LLM autonomously test token-direction
+edits on a model that is ALREADY loaded in the running J-Wash app (port 8381).
+
+It exposes only what is needed to experiment, and nothing else:
+
+  * generate                — (re)generate text from the current model
+  * scale_token / replace_token — apply a pure-weights token operation, any intensity
+  * set_intensity           — global multiplier over all edits (sweep the intensity)
+  * list_edits / reset_edits — inspect / clear the current edits
+
+This is a thin HTTP client of the J-Wash REST API (same server as scripts/jlab.py),
+spoken over MCP/stdio so any MCP client (Claude Desktop, another agent, ...) can
+drive a model you loaded yourself. By design it never loads models or lenses,
+never changes the sampling defaults beyond the call, and never exports anything:
+a model AND a Jacobian lens must already be loaded from the J-Wash UI.
+
+Token edits are always applied in a *pure-weights* mode (read projection, or W_U
+abliteration on Gemma-style models): the live preview matches an exported
+checkpoint exactly, so what the model tests here is what a baked model would do.
+
+Run it from an MCP client over stdio:
+
+    pip install mcp
+    python -X utf8 scripts/jwash_mcp.py
+
+Point it at a non-default J-Wash instance with an env var:
+
+    JWASH_BASE=http://127.0.0.1:8382
+"""
+import json
+import os
+import urllib.error
+import urllib.parse
+import urllib.request
+
+from mcp.server.fastmcp import FastMCP
+
+BASE = os.environ.get("JWASH_BASE", "http://127.0.0.1:8381").rstrip("/")
+
+mcp = FastMCP(
+    "j-wash",
+    instructions=(
+        "Drive a model ALREADY loaded in the running J-Wash app to test "
+        "token-direction edits. Typical loop: (1) `generate` a baseline reply; "
+        "(2) find the exact token with `find_token` and the layers to target with "
+        "`list_layers`; (3) apply edits with `scale_token`/`replace_token` (layers "
+        "are required) — always pure-weights, faithful to an exported checkpoint; "
+        "(4) `generate` again to see the effect, tuning each edit's `factor` or the "
+        "global `set_intensity`; (5) `reset_edits` to start over. A model AND a "
+        "Jacobian lens must be loaded from the J-Wash UI first; this server never "
+        "loads models, lenses, or exports checkpoints."
+    ),
+)
+
+
+# --- HTTP plumbing (stdlib only, like scripts/jlab.py) ----------------------
+
+def _call(method, path, body=None, timeout=600):
+    url = BASE + path
+    data = json.dumps(body).encode("utf-8") if body is not None else None
+    req = urllib.request.Request(url, data=data, method=method)
+    if data is not None:
+        req.add_header("Content-Type", "application/json")
+    try:
+        with urllib.request.urlopen(req, timeout=timeout) as resp:
+            return json.loads(resp.read().decode("utf-8"))
+    except urllib.error.HTTPError as exc:
+        detail = exc.read().decode("utf-8", "replace")
+        try:
+            detail = json.loads(detail).get("detail", detail)
+        except (json.JSONDecodeError, AttributeError):
+            pass
+        raise ValueError(f"J-Wash {method} {path} -> HTTP {exc.code}: {detail}")
+    except urllib.error.URLError as exc:
+        raise ValueError(
+            f"J-Wash server unreachable at {BASE} ({exc.reason}). Start it "
+            "(python -X utf8 run.py) and load a model + lens, then retry."
+        )
+
+
+def _status():
+    return _call("GET", "/api/status")
+
+
+def _resolve_token(text):
+    """The EXACT single token for ``text`` (a leading space is significant)."""
+    r = _call("GET", "/api/token-lookup?q=" + urllib.parse.quote(text.strip()))
+    cands = r.get("candidates", [])
+    for c in cands:
+        if c["str"] == text:
+            return c
+    listing = ", ".join(f"{c['id']}:{c['str']!r}" for c in cands) or "none"
+    raise ValueError(
+        f"No exact single-token match for {text!r}. A leading space is "
+        f"significant (mid-sentence words usually need one, e.g. ' model'). "
+        f"Candidates: {listing}"
+    )
+
+
+def _parse_layers(spec, n_layers):
+    """None -> server default band; 'all'/'none'/'19-31'/'3,5,7' -> explicit list."""
+    if spec is None:
+        return None
+    spec = spec.strip().lower()
+    if spec in ("", "none"):
+        return []
+    if spec == "all":
+        if not n_layers:
+            raise ValueError("layers='all' needs a loaded model to know the layer count")
+        return list(range(n_layers))
+    out = set()
+    for part in spec.split(","):
+        part = part.strip()
+        if not part:
+            continue
+        if "-" in part:
+            lo, hi = part.split("-", 1)
+            out.update(range(int(lo), int(hi) + 1))
+        else:
+            out.add(int(part))
+    return sorted(out)
+
+
+_PURE_WEIGHTS_MODES = ("readthrough", "exact", "abliteration")
+
+
+def _edits_summary():
+    st = _status()
+    mode = st.get("interventions_mode")
+    return {
+        "mode": mode,
+        "pure_weights": mode in _PURE_WEIGHTS_MODES,
+        "global_intensity": st.get("interventions_scale"),
+        "edits": [
+            {
+                "id": r["id"],
+                "token": r["token"],
+                "op": r["mode"],
+                "factor": r["factor"],
+                "replacement": r.get("replacement"),
+                "layers": r["layers"],
+            }
+            for r in (st.get("interventions") or [])
+        ],
+    }
+
+
+def _add_rule(token, op, factor, replacement, layers):
+    st = _status()
+    loaded = st.get("loaded")
+    if not loaded:
+        raise ValueError(
+            "No model loaded in J-Wash — load a model and a Jacobian lens from "
+            "the app first."
+        )
+    if not st.get("lens"):
+        raise ValueError(
+            "No Jacobian lens loaded — load one in the Lens tab of J-Wash before "
+            "editing tokens."
+        )
+    # Force a pure-weights mode: read projection, or W_U abliteration on
+    # architectures that normalize their writes (Gemma 2/3 style).
+    pure_mode = "abliteration" if loaded.get("rebase_supported") is False else "readthrough"
+    _call("PATCH", "/api/interventions", {"mode": pure_mode})
+
+    body = {"token_id": _resolve_token(token)["id"], "mode": op, "factor": float(factor)}
+    if op == "replace":
+        body["replacement_id"] = _resolve_token(replacement)["id"]
+    parsed = _parse_layers(layers, loaded.get("n_layers"))
+    if parsed is not None:
+        body["layers"] = parsed
+    _call("POST", "/api/interventions", body)
+    return _edits_summary()
+
+
+# --- MCP tools --------------------------------------------------------------
+
+@mcp.tool()
+def generate(prompt: str, system: str | None = None, max_tokens: int = 200,
+             temperature: float = 0.0, seed: int = 1234) -> str:
+    """(Re)generate a reply from the model currently loaded in J-Wash, with the
+    active token edits applied — call it again to regenerate.
+
+    At temperature 0 generation is deterministic, so the reply changes only when
+    the edits change: this is the clean way to compare behaviour before vs after
+    an edit. Raise `temperature` (or set `seed=-1` for a random seed) to sample
+    varied continuations instead. `system` is an optional system prompt.
+    """
+    messages = [{"role": "system", "content": system}] if system else []
+    messages.append({"role": "user", "content": prompt})
+    r = _call("POST", "/api/generate", {
+        "messages": messages,
+        "sampling": {"temperature": temperature, "max_tokens": max_tokens, "seed": seed},
+    }, timeout=1800)
+    return r.get("text", "")
+
+
+@mcp.tool()
+def scale_token(token: str, factor: float, layers: str) -> dict:
+    """Multiply a token's own direction by `factor` (pure-weights edit).
+
+    `factor` is the intensity: 0 removes the token's direction, 0<factor<1
+    attenuates it, factor>1 amplifies it. `token` is the exact token string — a
+    leading space is usually significant (e.g. ' model'); use `find_token` to get
+    it. `layers` is REQUIRED: it selects where the edit acts and an edit that
+    targets no layer does nothing — pass a 0-based range or list ('19-25', '20',
+    '20,24', or 'all') and call `list_layers` to see the model's layers. The mode
+    is forced to pure-weights so the effect matches an exported checkpoint.
+    """
+    return _add_rule(token, "scale", factor, None, layers)
+
+
+@mcp.tool()
+def replace_token(token: str, replacement: str, layers: str, factor: float = 1.0) -> dict:
+    """Rewrite `token`'s component onto `replacement`'s direction (pure-weights),
+    e.g. token=' model', replacement=' fish' to make the model talk like a fish.
+
+    You MUST pass `layers` — it selects the layers where the replacement is
+    applied, and WITHOUT it nothing happens. Give a 0-based range or list
+    ('19-25', '20,24', or 'all'); call `list_layers` for the model's layers and
+    `find_token` for the exact ' token' strings (both must be single tokens, a
+    leading space usually being significant). `factor` scales the strength
+    (1.0 = full). The mode is forced to pure-weights (faithful to an export).
+    """
+    return _add_rule(token, "replace", factor, replacement, layers)
+
+
+@mcp.tool()
+def set_intensity(scale: float) -> dict:
+    """Set the global multiplier applied to ALL active edits — sweep the overall
+    intensity without touching each rule: 0 disables every edit, 1 is nominal,
+    >1 pushes them harder. Returns the current edits.
+    """
+    _call("PATCH", "/api/interventions", {"scale": scale})
+    return _edits_summary()
+
+
+@mcp.tool()
+def list_edits() -> dict:
+    """Show the active token edits, the pure-weights mode in force, and the global
+    intensity — a read-only snapshot of the current experiment.
+    """
+    return _edits_summary()
+
+
+@mcp.tool()
+def reset_edits() -> dict:
+    """Remove every token edit, returning the model to its unedited behaviour.
+    Use it to start a fresh experiment. Returns the (now empty) edits.
+    """
+    _call("DELETE", "/api/interventions")
+    return _edits_summary()
+
+
+@mcp.tool()
+def find_token(text: str) -> dict:
+    """Look up the single-token forms of `text` so you can pick the exact token to
+    edit before calling scale_token/replace_token.
+
+    A leading space is significant (' model' and 'model' are different tokens), so
+    the lookup also tries the space-prefixed and capitalization variants and
+    returns those that are exactly one token. Use a returned `token` string
+    verbatim; words that split into several tokens can't be edited directly.
+    """
+    r = _call("GET", "/api/token-lookup?q=" + urllib.parse.quote(text.strip()))
+    return {
+        "query": text,
+        "candidates": [{"id": c["id"], "token": c["str"]} for c in r.get("candidates", [])],
+    }
+
+
+@mcp.tool()
+def list_layers() -> dict:
+    """List the layers you can target with `scale_token`/`replace_token`.
+
+    Returns the model's total layer count (indices are 0-based, so the valid
+    range is 0..n_layers-1) and the layers the loaded Jacobian lens actually
+    covers — those are the calibrated ones to edit; targeting a layer outside
+    them falls back to a less reliable logit-lens direction.
+    """
+    st = _status()
+    loaded = st.get("loaded")
+    if not loaded:
+        raise ValueError(
+            "No model loaded in J-Wash — load a model and a Jacobian lens from "
+            "the app first."
+        )
+    n = loaded.get("n_layers")
+    lens = st.get("lens") or {}
+    out = {"n_layers": n, "valid_range": f"0-{n - 1}" if n else None}
+    if lens.get("fitted_layers_all"):
+        out["lens_fitted_layers"] = lens["fitted_layers_all"]
+    if lens.get("tapped_layers"):
+        out["lens_tapped_layers"] = lens["tapped_layers"]
+    return out
+
+
+if __name__ == "__main__":
+    mcp.run(transport="stdio")

ui/src/App.jsx

@@ -45,7 +45,7 @@ async function jsonFetch(url, options) {
   return body
 }
 
-const SAMPLING_DEFAULT = { temperature: 0.7, top_p: 0.95, top_k: 40, max_tokens: 512, seed: -1 }
+const SAMPLING_DEFAULT = { temperature: 0.7, top_p: 0.95, top_k: 40, max_tokens: 512, seed: -1, repetition_penalty: 1 }
 
 // Human-readable name of the loaded lens (local path or Hub file) for "which lens do I have?".
 function lensName(meta) {
@@ -124,6 +124,11 @@ export default function App() {
   const [chatMd, setChatMd] = useState(() => localStorage.getItem('jlens_chat_md') !== '0')
   useEffect(() => { localStorage.setItem('jlens_chat_md', chatMd ? '1' : '0') }, [chatMd])
 
+  // "API monitor" mode: instead of polling status every 2s, refresh only when an
+  // API/MCP generation completes (the server pushes on /ws). Off = normal 2s poll.
+  const [apiMonitor, setApiMonitor] = useState(() => localStorage.getItem('jlens_api_monitor') === '1')
+  useEffect(() => { localStorage.setItem('jlens_api_monitor', apiMonitor ? '1' : '0') }, [apiMonitor])
+
   // server-side settings (Options tab)
   const [settings, setSettings] = useState(null)
   useEffect(() => {
@@ -274,11 +279,29 @@ export default function App() {
         }
       }).catch(() => localStorage.removeItem('jlens_conv'))
     }
+  }, [])
+
+  // status refresh — normally a 2s poll; in "API monitor" mode it is event-driven
+  // instead, refetched only when an API/MCP generation ends (pushed over /ws).
+  useEffect(() => {
     const tick = () => jsonFetch('/api/status').then(setStatus).catch(() => {})
     tick()
-    const id = setInterval(tick, 2000)
+    if (!apiMonitor) {
-    return () => clearInterval(id)
+      const id = setInterval(tick, 2000)
-  }, [])
+      return () => clearInterval(id)
+    }
+    let ws
+    let timer
+    let closed = false
+    const open = () => {
+      const proto = location.protocol === 'https:' ? 'wss' : 'ws'
+      ws = new WebSocket(`${proto}://${location.host}/ws`)
+      ws.onmessage = (ev) => { try { if (JSON.parse(ev.data)?.type === 'api_generation') tick() } catch { /* ignore */ } }
+      ws.onclose = () => { if (!closed) timer = setTimeout(open, 2000) }
+    }
+    open()
+    return () => { closed = true; clearTimeout(timer); if (ws) ws.close() }
+  }, [apiMonitor])
 
   useEffect(() => {
     if (conv?.id) localStorage.setItem('jlens_conv', String(conv.id))
@@ -1517,6 +1540,11 @@ export default function App() {
               onChange={(e) => patchSettings({ chat_markdown: e.target.checked })} /> render replies as markdown
           </label>
         </div>
+        <div className="row"><label title="while an external API/MCP client drives the model: stop the 2s status polling and refresh the view only when a generation finishes — keeps the UI fast">API monitor</label>
+          <label style={{ width: 'auto' }}>
+            <input type="checkbox" checked={apiMonitor} onChange={(e) => setApiMonitor(e.target.checked)} /> refresh only when an API generation ends
+          </label>
+        </div>
 
         <h2>Paths</h2>
         <div className="row">
@@ -1591,6 +1619,16 @@ export default function App() {
           <textarea value={system} onChange={(e) => setSystem(e.target.value)} placeholder="(none)" disabled={!!conv?.id} />
         </details>
 
+        {status?.last_generation?.text && (
+          <div style={{ border: '1px solid var(--border)', borderRadius: 6, padding: '6px 8px', marginBottom: 6, maxHeight: 160, overflow: 'auto' }}
+            title="latest text generated through the API (CLI/MCP) — not saved to the conversation">
+            <div style={{ color: 'var(--muted)', fontSize: 12, marginBottom: 4 }}>
+              generated via API/MCP{status.last_generation.prompt ? ` · ${status.last_generation.prompt.slice(0, 80)}` : ''}
+            </div>
+            <div style={{ whiteSpace: 'pre-wrap', fontSize: 13 }}>{status.last_generation.text}</div>
+          </div>
+        )}
+
         <div className="messages" ref={messagesRef}>
           {messages.map((m, i) => (
             <div
@@ -1709,15 +1747,14 @@ export default function App() {
               onChange={(e) => setSampling({ ...sampling, max_tokens: +e.target.value })} /></label>
             <label title="-1 = random; ≥ 0 = reproducible sampling">seed <input type="number" step="1" min="-1" value={sampling.seed}
               onChange={(e) => setSampling({ ...sampling, seed: Math.trunc(+e.target.value) })} /></label>
+            <label title="repetition penalty — 1 = off, >1 penalizes tokens already in the context (curbs loops)">rep <input type="number" step="0.05" min="1" max="2" value={sampling.repetition_penalty ?? 1}
+              onChange={(e) => setSampling({ ...sampling, repetition_penalty: +e.target.value })} /></label>
             <button onClick={onRegenerate} disabled={streaming || !messages.some((m) => m.role === 'assistant')}>Regenerate</button>
             <button onClick={onContinue}
               title="extend the last reply: the model picks up exactly where it stopped"
               disabled={streaming || messages[messages.length - 1]?.role !== 'assistant' || messages[messages.length - 1]?.id == null}>
               Continue</button>
             <button onClick={onEditLast} disabled={streaming || !messages.some((m) => m.role === 'user')}>Edit last</button>
-            <label title="render assistant replies as markdown">
-              <input type="checkbox" checked={chatMd} onChange={(e) => setChatMd(e.target.checked)} /> md
-            </label>
             {lensMeta && (
               <button
                 className={editorOpen ? 'ed-toggle-on' : ''}