server.py

"""M3 — live J-lens chat server.

Chat with an open-weight model and watch its workspace: every generated token
streams back together with the J-lens readout (top words per layer) at that
position; the prompt's full position × layer grid arrives right after prefill.

Endpoints:
  GET  /               the UI (static/index.html)
  GET  /api/info       model/lens metadata (layers, band, vocab hints)
  POST /api/generate   NDJSON stream: prefill grid event, then one event per
                       generated token, then a done event.

Run:
  python server.py [--model Qwen/Qwen3-1.7B] [--device cuda] [--port 8732]

Design notes:
- One global model + lens, one GPU: requests are serialized with a lock.
- Custom generation loop (KV cache) with residual hooks at the selected
  layers; per step the lens transport + unembed gives the readout. The lens
  adds one d_model x d_model and one unembed matmul per tracked layer per
  token — negligible next to the forward pass itself.
- `use_jacobian=false` in the request switches to the raw logit lens.
"""

from __future__ import annotations

import argparse
import asyncio
import json
import logging
import sys
import threading
from pathlib import Path

import torch
import uvicorn
from fastapi import FastAPI, Request
from fastapi.responses import FileResponse, StreamingResponse

logger = logging.getLogger("m3")

LENS_REPO = "neuronpedia/jacobian-lens"
LENS_REVISION = "qwen-n1000"
LENS_FILES = {
    "Qwen/Qwen3-1.7B": "qwen3-1.7b/jlens/Salesforce-wikitext/Qwen3-1.7B_jacobian_lens.pt",
    "Qwen/Qwen3.5-4B": "qwen3.5-4b/jlens/Salesforce-wikitext/Qwen3.5-4B_jacobian_lens_n1000.pt",
    "Qwen/Qwen3.5-2B-Base": "qwen3.5-2b-pt/jlens/Salesforce-wikitext/Qwen3.5-2B-Base_jacobian_lens.pt",
}

STATE: dict = {}
GPU_LOCK = threading.Lock()

# Prefill readout processes positions in slices this big, capping the
# transient [positions x vocab] logits buffers (~500 MB at 256 positions on a
# 250k vocab) regardless of conversation length.
import os

READOUT_CHUNK = int(os.environ.get("JSPACE_READOUT_CHUNK", "256"))


# --------------------------------------------------------------------------- #
# model plumbing
# --------------------------------------------------------------------------- #


def load(model_name: str, device: str, dtype: str = "auto") -> None:
    import jlens
    import transformers

    # bf16 everywhere by default except plain CPU, where fp32 is the safe
    # default; pass --dtype bf16 to halve CPU RAM (e.g. 4B in ~8 GB).
    if dtype == "auto":
        torch_dtype = torch.float32 if device == "cpu" else torch.bfloat16
    else:
        torch_dtype = {"bf16": torch.bfloat16, "fp32": torch.float32}[dtype]
    logger.info("loading %s (%s, %s)", model_name, device, torch_dtype)
    if device == "auto":
        hf = transformers.AutoModelForCausalLM.from_pretrained(
            model_name, dtype=torch_dtype, device_map="auto"
        )
    else:
        hf = transformers.AutoModelForCausalLM.from_pretrained(
            model_name, dtype=torch_dtype
        ).to(device)
    tok = transformers.AutoTokenizer.from_pretrained(model_name)
    model = jlens.from_hf(hf, tok)
    lens = jlens.JacobianLens.from_pretrained(
        LENS_REPO, filename=LENS_FILES[model_name], revision=LENS_REVISION
    )
    n = model.n_layers
    # heuristic workspace band (paper: ~1/3 of depth to a few layers before out)
    band = [round(n / 3), n - 3]
    from interventions import LensVectors
    STATE.update(model_name=model_name, model=model, tok=tok, lens=lens,
                 hf=hf, band=band, lens_vectors=LensVectors(model, lens),
                 display_mask=_display_mask(model, tok, model_name))
    logger.info("ready: n_layers=%d d_model=%d band=%s", n, model.d_model, band)


def _display_mask(model, tok, model_name: str) -> torch.Tensor:
    """Word-start display filter: readouts shown to the user are restricted to
    tokens that begin a plain word (leading space, alphabetic, len > 2).
    Fragments like 'itude' or bare punctuation are noise, not thoughts.
    Ranks for pinned words are always computed over the FULL vocabulary —
    the mask affects display only. Cached next to the server per model."""
    import re

    from jlens.vis import _meaningful_token_mask

    vocab = model._lm_head.weight.shape[0]
    cache = Path(__file__).parent / f"display_mask_{model_name.split('/')[-1]}.pt"
    if cache.exists():
        return torch.load(cache, weights_only=True).to(model.input_device)
    logger.info("building word-start display mask (one-time)")
    mask = _meaningful_token_mask(tok, vocab, torch.device("cpu"))
    word_re = re.compile(r"^[A-Za-z][A-Za-z'\-]+$")
    for tid in mask.nonzero().flatten().tolist():
        raw = tok.decode([tid])
        s = raw.strip()
        if not (raw.startswith(" ") and s.isascii() and word_re.match(s) and len(s) > 2):
            mask[tid] = False
    torch.save(mask, cache)
    logger.info("display vocabulary: %d word tokens", int(mask.sum()))
    return mask.to(model.input_device)


def default_layers() -> list[int]:
    """Readout layers: fitted layers, thinned to ~12, always incl. band edges."""
    lens, model = STATE["lens"], STATE["model"]
    fitted = [l for l in lens.source_layers]
    step = max(1, len(fitted) // 12)
    picked = sorted({*fitted[::step], STATE["band"][0], STATE["band"][1],
                     fitted[-1]} & set(fitted))
    return picked


@torch.no_grad()
def lens_readout(residuals: dict[int, torch.Tensor], layers: list[int],
                 use_jacobian: bool, topk: int, pinned_ids: list[int],
                 masked: bool = True):
    """residuals: {layer: [n_pos, d_model]} -> per-layer top-k (+probs) and
    pinned ranks. Display top-k is restricted to word-start tokens when
    `masked`; probabilities and pinned ranks always use the full vocabulary."""
    model, lens, tok = STATE["model"], STATE["lens"], STATE["tok"]
    mask = STATE["display_mask"]
    out = {}
    for layer in layers:
        h_all = residuals[layer].float()
        ids_parts, probs_parts, ranks_parts = [], [], []
        for start in range(0, h_all.shape[0], READOUT_CHUNK):
            h = h_all[start:start + READOUT_CHUNK]
            if use_jacobian and layer in lens.jacobians:
                h = lens.transport(h, layer)
            logits = model.unembed(h).float()  # [chunk, vocab]
            probs = torch.softmax(logits, dim=-1)
            display = logits.masked_fill(~mask.to(logits.device), float("-inf")) \
                if masked else logits
            top = display.topk(topk, dim=-1)
            ids_parts.append(top.indices.cpu())
            probs_parts.append(probs.gather(-1, top.indices).cpu())
            if pinned_ids:
                pin_vals = logits[:, pinned_ids]  # [chunk, n_pin]
                ranks_parts.append(
                    (logits.unsqueeze(-1) > pin_vals.unsqueeze(1)).sum(1).cpu())
            del logits, probs, display
        out[layer] = {
            "ids": torch.cat(ids_parts),
            "probs": torch.cat(probs_parts),
            "ranks": torch.cat(ranks_parts) if ranks_parts else None,
        }
    return out


def decode_batch(ids) -> list[str]:
    tok = STATE["tok"]
    return [tok.decode([t]) for t in ids]


def pack_positions(readout, layers, n_pos, pin_groups):
    """-> list over positions of {layer: {top: [words], ranks: [ints]}}.

    Ranks are collapsed per pinned *word*: min over its token variants
    (bare / leading-space / case variants)."""
    cols = []
    for p in range(n_pos):
        col = {}
        for layer in layers:
            entry = {
                "top": decode_batch(readout[layer]["ids"][p].tolist()),
                "p": [round(x, 5) for x in readout[layer]["probs"][p].tolist()],
            }
            if pin_groups:
                raw = readout[layer]["ranks"][p].tolist()
                entry["ranks"] = [min(raw[a:b]) for a, b in pin_groups]
            col[str(layer)] = entry
        cols.append(col)
    return cols


def single_token_ids(word: str) -> list[int]:
    tok = STATE["tok"]
    ids = []
    for v in {word, " " + word, word.capitalize(), " " + word.capitalize(),
              word.lower(), " " + word.lower()}:
        enc = tok.encode(v, add_special_tokens=False)
        if len(enc) == 1:
            ids.append(enc[0])
    return sorted(set(ids))


@torch.no_grad()
def generate_events(req: dict):
    """Synchronous generator of event dicts (runs in a worker thread)."""
    from jlens.hooks import ActivationRecorder

    import contextlib

    from interventions import InterventionHooks, WorkspaceOffHooks, prepare

    model, tok, hf = STATE["model"], STATE["tok"], STATE["hf"]
    layers = req.get("layers") or default_layers()
    topk = int(req.get("topk", 8))
    use_jacobian = bool(req.get("use_jacobian", True))
    max_new = min(int(req.get("max_new_tokens", 200)), 512)
    masked = bool(req.get("masked", True))
    pin_words = req.get("pinned", [])
    pinned_ids, pin_meta, pin_groups = [], [], []
    for w in pin_words:
        ids = single_token_ids(w)
        if not ids:
            continue
        start = len(pinned_ids)
        pinned_ids.extend(ids)
        pin_groups.append((start, len(pinned_ids)))  # slice into pinned_ids
        pin_meta.append({"word": w, "n_variants": len(ids)})

    if req.get("raw"):
        prompt = req["prompt"]
    else:
        messages = req["messages"]
        prompt = tok.apply_chat_template(
            messages, tokenize=False, add_generation_prompt=True,
            enable_thinking=bool(req.get("thinking", False)),
        )
    interventions = prepare(
        req.get("interventions", []), tok, STATE["lens_vectors"],
        tuple(STATE["band"]), STATE["lens"].source_layers,
    )
    ws_off = req.get("workspace_off")  # None | "light" | "medium" | "heavy"
    ws_ctx = (
        WorkspaceOffHooks(model, STATE["lens"], tuple(STATE["band"]), ws_off)
        if ws_off else contextlib.nullcontext()
    )

    input_ids = model.encode(prompt, max_length=2048)
    prompt_tokens = decode_batch(input_ids[0].tolist())
    n_prompt = len(prompt_tokens)

    yield {"type": "start", "n_prompt": n_prompt, "layers": layers,
           "band": STATE["band"], "pinned": pin_meta,
           "prompt_tokens": prompt_tokens,
           "workspace_off": ws_off,
           "interventions": [
               {"kind": iv.kind, "source": iv.source, "target": iv.target,
                "alpha": iv.alpha, "layers": list(iv.layer_range)}
               for iv in interventions]}

    eos_ids = {tok.eos_token_id}
    im_end = tok.convert_tokens_to_ids("<|im_end|>")
    if im_end is not None:
        eos_ids.add(im_end)

    try:
        yield from _generate_body(
            hf, model, tok, input_ids, layers, use_jacobian, topk, max_new,
            masked, pinned_ids, pin_groups, eos_ids, n_prompt, ws_ctx,
            interventions)
    finally:
        # hand the allocator's freed blocks back to the driver so this
        # process's VRAM footprint drops between requests (shared-GPU manners)
        if torch.cuda.is_available():
            torch.cuda.empty_cache()

    yield {"type": "done"}


@torch.no_grad()
def _generate_body(hf, model, tok, input_ids, layers, use_jacobian, topk,
                   max_new, masked, pinned_ids, pin_groups, eos_ids, n_prompt,
                   ws_ctx, interventions):
    from jlens.hooks import ActivationRecorder

    from interventions import InterventionHooks

    with GPU_LOCK, ws_ctx, InterventionHooks(model.layers, interventions):
        # ---- prefill: full grid over prompt positions
        # (intervention hooks are registered first, so the recorder — and
        # therefore the telemetry — sees the modified residual stream)
        with ActivationRecorder(model.layers, at=layers) as rec:
            out = hf(input_ids=input_ids, use_cache=True)
            residuals = {l: rec.activations[l][0].detach() for l in layers}
        past = out.past_key_values
        readout = lens_readout(residuals, layers, use_jacobian, topk, pinned_ids, masked)
        yield {"type": "prefill",
               "columns": pack_positions(readout, layers, n_prompt, pin_groups)}

        # ---- decode loop
        next_id = int(out.logits[0, -1].argmax())
        for step in range(max_new):
            if next_id in eos_ids:
                break
            step_ids = torch.tensor([[next_id]], device=input_ids.device)
            with ActivationRecorder(model.layers, at=layers) as rec:
                out = hf(input_ids=step_ids, past_key_values=past, use_cache=True)
                residuals = {l: rec.activations[l][0].detach() for l in layers}
            past = out.past_key_values
            readout = lens_readout(residuals, layers, use_jacobian, topk, pinned_ids, masked)
            col = pack_positions(readout, layers, 1, pin_groups)[0]
            yield {"type": "token", "text": tok.decode([next_id]),
                   "pos": n_prompt + step, "col": col}
            next_id = int(out.logits[0, -1].argmax())


# --------------------------------------------------------------------------- #
# app
# --------------------------------------------------------------------------- #

app = FastAPI(title="jspace live lens")
STATIC = Path(__file__).parent / "static"


@app.get("/")
def index():
    return FileResponse(STATIC / "index.html")


@app.get("/api/info")
def info():
    model, lens = STATE["model"], STATE["lens"]
    return {
        "model": STATE["model_name"],
        "n_layers": model.n_layers,
        "d_model": model.d_model,
        "fitted_layers": lens.source_layers,
        "default_layers": default_layers(),
        "band": STATE["band"],
    }


@app.post("/api/generate")
async def generate(request: Request):
    req = await request.json()

    def ndjson():
        try:
            for event in generate_events(req):
                yield json.dumps(event, ensure_ascii=False) + "\n"
        except Exception as exc:  # surface errors to the client
            logger.exception("generate failed")
            yield json.dumps({"type": "error", "message": str(exc)}) + "\n"

    # run the sync generator on a thread so the event loop stays responsive
    async def agen():
        loop = asyncio.get_running_loop()
        it = ndjson()
        while True:
            chunk = await loop.run_in_executor(None, lambda: next(it, None))
            if chunk is None:
                break
            yield chunk

    return StreamingResponse(agen(), media_type="application/x-ndjson")


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--model", default="Qwen/Qwen3-1.7B")
    parser.add_argument("--device", default="cuda",
                        help="cuda | cpu | auto (auto = GPU with CPU offload for large models)")
    parser.add_argument("--dtype", default="auto", choices=["auto", "bf16", "fp32"],
                        help="auto = bf16 on GPU, fp32 on CPU; bf16 halves CPU RAM")
    parser.add_argument("--vram-fraction", type=float, default=None,
                        help="cap this process's share of GPU memory (0-1); a "
                             "request that would exceed it fails cleanly instead "
                             "of starving other GPU tenants")
    parser.add_argument("--host", default="127.0.0.1",
                        help="bind address (use 0.0.0.0 inside a container)")
    parser.add_argument("--port", type=int, default=8732)
    args = parser.parse_args()
    sys.stdout.reconfigure(encoding="utf-8", errors="replace")
    sys.stderr.reconfigure(encoding="utf-8", errors="replace")
    logging.basicConfig(level=logging.INFO, format="%(asctime)s %(message)s",
                        datefmt="%H:%M:%S")
    if args.vram_fraction and torch.cuda.is_available():
        torch.cuda.set_per_process_memory_fraction(args.vram_fraction, 0)
        logger.info("VRAM capped at %.0f%% of device 0", args.vram_fraction * 100)
    load(args.model, args.device, args.dtype)
    uvicorn.run(app, host=args.host, port=args.port, log_level="warning")


if __name__ == "__main__":
    main()