"""Next Token Explorer: live next-token probabilities from a small language model. Gradio Space, runs on ZeroGPU or on a free CPU. The whole interface is one gr.HTML component (ui/template.html, ui/style.css, ui/app.js): - tokenizing while you type is a server function (CPU only, no queue, no GPU), - every model step goes through a normal queued Gradio event, so ZeroGPU can attribute GPU time to the visitor. One GPU call can take several steps (used by "Write on its own"). """ import spaces # import before torch touches CUDA (ZeroGPU) import json import os import time from pathlib import Path import gradio as gr import torch from transformers import AutoModelForCausalLM, AutoTokenizer MODEL_ID = os.environ.get("MODEL_ID", "Qwen/Qwen2.5-1.5B-Instruct") DTYPE_NAME = os.environ.get("DTYPE", "float32") DTYPE = {"float32": torch.float32, "bfloat16": torch.bfloat16, "float16": torch.float16}[DTYPE_NAME] ZERO_GPU = spaces.config.Config.zero_gpu DEVICE = "cuda" if (ZERO_GPU or torch.cuda.is_available()) else "cpu" TOP_K = 20 MAX_INPUT_TOKENS = int(os.environ.get("MAX_INPUT_TOKENS", "1024")) MAX_STEPS = 40 GRID = [round(0.05 * i, 2) for i in range(1, 41)] # temperatures 0.05 ... 2.0 UI = Path(__file__).parent / "ui" if DEVICE == "cpu": torch.set_num_threads(max(1, os.cpu_count() or 1)) torch.set_grad_enabled(False) tokenizer = AutoTokenizer.from_pretrained(MODEL_ID) model = AutoModelForCausalLM.from_pretrained(MODEL_ID, dtype=DTYPE) model.to(DEVICE) # on ZeroGPU this must happen at module level model.eval() EOS_IDS = sorted({i for i in [tokenizer.eos_token_id, tokenizer.convert_tokens_to_ids("<|im_end|>"), tokenizer.convert_tokens_to_ids("<|endoftext|>")] if i is not None}) VOCAB = len(tokenizer) _probe = tokenizer.apply_chat_template([{"role": "user", "content": "x"}], tokenize=False) DEFAULT_SYSTEM = _probe.split("<|im_start|>system\n", 1)[1].split("<|im_end|>", 1)[0] if "<|im_start|>system\n" in _probe else "" IM_START, IM_END = tokenizer.convert_tokens_to_ids("<|im_start|>"), tokenizer.convert_tokens_to_ids("<|im_end|>") SYS_HEAD = len(tokenizer("<|im_start|>system\n", add_special_tokens=False).input_ids) # <|im_start|> system \n OUTPUT_SIZE = model.config.vocab_size # ---------------------------------------------------------------- tokens as readable pieces def _bytes_to_unicode(): bs = list(range(ord("!"), ord("~") + 1)) + list(range(ord("¡"), ord("¬") + 1)) + list(range(ord("®"), ord("ÿ") + 1)) cs = bs[:] n = 0 for b in range(256): if b not in bs: bs.append(b) cs.append(256 + n) n += 1 return {chr(c): b for b, c in zip(bs, cs)} U2B = _bytes_to_unicode() SPECIAL_IDS = set(tokenizer.all_special_ids) _pieces = {} def piece(token_id: int) -> dict: """Readable form of one token: its text, or its raw bytes if it holds only part of a character.""" if token_id in _pieces: return _pieces[token_id] raw = tokenizer.convert_ids_to_tokens(token_id) if token_id in SPECIAL_IDS or (raw.startswith("<|") and raw.endswith("|>")): out = {"id": token_id, "s": raw, "special": True} else: try: b = bytes(U2B[ch] for ch in raw) except KeyError: b = raw.encode("utf-8") try: out = {"id": token_id, "s": b.decode("utf-8")} except UnicodeDecodeError: out = {"id": token_id, "s": " ".join(f"{x:02x}" for x in b), "partial": True} _pieces[token_id] = out return out def chat_messages(text: str, system: str | None) -> list[dict]: """An empty system prompt means no system message: Qwen's chat template then inserts its own default.""" msgs = [{"role": "system", "content": system}] if (system or "").strip() else [] return msgs + [{"role": "user", "content": text}] def build_ids(text: str, mode: str, cont: list, system: str | None = None) -> list[int]: text = (text or "")[:4000] system = (system or "")[:4000] if isinstance(system, str) else "" if mode == "chat": prompt = tokenizer.apply_chat_template(chat_messages(text, system), tokenize=False, add_generation_prompt=True) else: prompt = text cont = [int(i) for i in (cont or [])][:400] if any(i < 0 or i >= VOCAB for i in cont): raise ValueError("Unknown token in the continuation.") ids = tokenizer(prompt, add_special_tokens=False).input_ids + cont if len(ids) > MAX_INPUT_TOKENS: raise ValueError(f"The input is too long for this demo ({len(ids)} tokens, limit {MAX_INPUT_TOKENS}). Shorten the text or the system prompt.") return ids # ---------------------------------------------------------------- the model (GPU on ZeroGPU) _cpu_cache = {"ids": [], "kv": None} # CPU only: reuse earlier tokens between calls def _forward(seq, new, past): out = model(input_ids=torch.tensor([new], device=DEVICE), past_key_values=past, use_cache=True) return torch.log_softmax(out.logits[0, -1].float(), dim=-1), out.past_key_values def _summary(lp): top = torch.topk(lp, TOP_K) temps = torch.tensor(GRID, device=lp.device).unsqueeze(1) lD = torch.logsumexp(lp.unsqueeze(0) / temps, dim=1) # log sum_j p_j^(1/T), exact over the whole output return {"ids": top.indices.tolist(), "lp": [round(v, 5) for v in top.values.tolist()], "lD": [round(v, 5) for v in lD.tolist()]} def _duration(ids, actions, temperature): # seconds of GPU time to reserve (on ZeroGPU), the actual compute is a small fraction of this return min(40, 10 + 0.25 * len(actions)) @spaces.GPU(duration=_duration) def advance(ids: list[int], actions: list[dict], temperature: float) -> dict: """Run the model on ids, then apply each action: {"id": token} adds that token, {"r": x} samples the token at position x of the cumulative distribution at this temperature (exact, over the whole output layer). Returns one summary per state and the chosen tokens.""" with torch.inference_mode(): return _advance(ids, actions, temperature) def _advance(ids, actions, temperature): t0 = time.time() seq = list(ids) if not ZERO_GPU and _cpu_cache["kv"] is not None and 0 < len(_cpu_cache["ids"]) < len(seq) and seq[:len(_cpu_cache["ids"])] == _cpu_cache["ids"]: lp, past = _forward(seq, seq[len(_cpu_cache["ids"]):], _cpu_cache["kv"]) else: lp, past = _forward(seq, seq, None) states, choices = [_summary(lp)], [] for a in actions[:MAX_STEPS]: if "id" in a: tid = int(a["id"]) if not 0 <= tid < VOCAB: break p_t = torch.softmax(lp / temperature, dim=-1) if temperature > 0 else torch.exp(lp) rank = int((p_t > p_t[tid]).sum()) + 1 p = float(p_t[tid]) elif temperature == 0: tid, rank, p = int(torch.argmax(lp)), 1, 1.0 else: probs = torch.softmax(lp / temperature, dim=-1) vals, order = torch.sort(probs, descending=True) k = int(torch.searchsorted(torch.cumsum(vals, 0), torch.tensor(float(a["r"]), device=vals.device)).clamp(max=len(vals) - 1)) tid, rank, p = int(order[k]), k + 1, float(vals[k]) if tid >= VOCAB: # unused padding slot of the output layer tid, rank, p = int(order[0]), 1, float(vals[0]) choices.append({"id": tid, "rank": rank, "p": p}) seq.append(tid) lp, past = _forward(seq, [tid], past) states.append(_summary(lp)) if tid in EOS_IDS: break if not ZERO_GPU: _cpu_cache["ids"], _cpu_cache["kv"] = seq, past return {"states": states, "choices": choices, "ms": round((time.time() - t0) * 1000)} # ---------------------------------------------------------------- handlers def system_span(ids: list[int], mode: str): """[first, end) token positions of the system prompt text, so the page can colour it. None outside chat mode.""" if mode != "chat" or not ids or ids[0] != IM_START or IM_END not in ids: return None end = ids.index(IM_END) return [min(SYS_HEAD, end), end] def tokenize(payload: dict) -> dict: """Server function for live tokenization while typing (CPU only, not queued, no GPU). Called from the page as server.tokenize({text, mode, cont, system}).""" payload = payload if isinstance(payload, dict) else {} cont = payload.get("cont") or [] system = payload.get("system") if isinstance(payload.get("system"), str) else "" try: ids = build_ids(payload.get("text", ""), payload.get("mode", "raw"), cont, system) except ValueError as e: return {"error": str(e)} return {"tokens": [piece(i) for i in ids], "n": len(ids), "n_cont": len(cont), "sys": system_span(ids, payload.get("mode", "raw")), "system_default": not system.strip()} def info(_=None) -> dict: return {"model": MODEL_ID, "dtype": DTYPE_NAME, "device": "ZeroGPU" if ZERO_GPU else DEVICE.upper(), "default_system": DEFAULT_SYSTEM, "vocab": VOCAB, "output_size": OUTPUT_SIZE, "grid": GRID, "eos": EOS_IDS} def handle(value: str): """Queued event: the page writes a request into the component's value and fires `submit`, the reply goes back into the value, where the page picks it up.""" try: req = json.loads(value or "{}") except json.JSONDecodeError: return gr.skip() if req.get("kind") != "req": return gr.skip() rid = req.get("rid") try: ids = build_ids(req.get("text", ""), req.get("mode", "raw"), req.get("cont", []), req.get("system")) if not ids: raise ValueError("Type a start text first.") temperature = float(req.get("temperature", 1.0)) if not 0 <= temperature <= 5: raise ValueError("Temperature out of range.") actions = [a for a in req.get("actions", []) if isinstance(a, dict)][:MAX_STEPS] res = advance(ids, actions, temperature) states = [{"top": [dict(piece(i), lp=l) for i, l in zip(s["ids"], s["lp"])], "lD": s["lD"]} for s in res["states"]] choices = [dict(piece(c["id"]), rank=c["rank"], p=c["p"]) for c in res["choices"]] final_ids = ids + [c["id"] for c in res["choices"]] reply = {"kind": "res", "rid": rid, "states": states, "choices": choices, "ms": res["ms"], "tokens": [piece(i) for i in final_ids], "n": len(final_ids), "n_cont": len(req.get("cont", [])) + len(choices), "sys": system_span(ids, req.get("mode", "raw"))} except Exception as e: # shown in the page, including ZeroGPU quota messages msg = getattr(e, "message", None) or str(e) or type(e).__name__ reply = {"kind": "res", "rid": rid, "error": msg} return json.dumps(reply) # ---------------------------------------------------------------- page with gr.Blocks(title="Next Token Explorer") as demo: ui = gr.HTML( value="", html_template=(UI / "template.html").read_text(encoding="utf-8"), js_on_load=(UI / "app.js").read_text(encoding="utf-8"), apply_default_css=False, elem_id="nte", server_functions=[tokenize, info], ) # on a CPU one request at a time (shared key/value cache), on ZeroGPU several in parallel ui.submit(handle, inputs=ui, outputs=ui, show_progress="hidden", queue=True, concurrency_limit=4 if ZERO_GPU else 1) if __name__ == "__main__": demo.launch( css=(UI / "style.css").read_text(encoding="utf-8"), head='' '', footer_links=[], ssr_mode=False, )