From 4bc7c8005ec65dc82b64bf008aef7932d1d4640c Mon Sep 17 00:00:00 2001 From: Extraltodeus Date: Tue, 14 Jul 2026 07:29:29 +0200 Subject: [PATCH] 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 --- api/app.py | 21 +++ config.py | 1 + core/model_manager.py | 18 ++- scripts/jwash_mcp.py | 299 ++++++++++++++++++++++++++++++++++++++++++ ui/src/App.jsx | 51 ++++++- 5 files changed, 381 insertions(+), 9 deletions(-) create mode 100644 scripts/jwash_mcp.py diff --git a/api/app.py b/api/app.py index ef5e3c7..1a544ed 100644 --- a/api/app.py +++ b/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")} diff --git a/config.py b/config.py index 6580736..8cb7156 100644 --- a/config.py +++ b/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 } diff --git a/core/model_manager.py b/core/model_manager.py index 0213116..d05cca2 100644 --- a/core/model_manager.py +++ b/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: diff --git a/scripts/jwash_mcp.py b/scripts/jwash_mcp.py new file mode 100644 index 0000000..478d451 --- /dev/null +++ b/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, 01 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") diff --git a/ui/src/App.jsx b/ui/src/App.jsx index 5cd2be9..680327e 100644 --- a/ui/src/App.jsx +++ b/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) - return () => clearInterval(id) - }, []) + if (!apiMonitor) { + 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 +
+ +

Paths

@@ -1591,6 +1619,16 @@ export default function App() {