liautaud/Next_token_probability
0
1import math2import os3import time4 5import gradio as gr6import torch7from transformers import AutoModelForCausalLM, AutoTokenizer8 9# ---- Model config ----10MODEL_NAME = "microsoft/biogpt" # e.g. "distilgpt2", "HuggingFaceTB/SmolLM2-135M"11DEVICE = "cuda" if torch.cuda.is_available() else "cpu"12FULL_BATCH_CONTEXT_THRESHOLD = 64 # short contexts are often faster as one batched full forward13EPS = 1e-914 15# Set TORCH_COMPILE=1 to compile.16ENABLE_TORCH_COMPILE = os.environ.get("TORCH_COMPILE", "0") != "0"17COMPILE_STATUS = "not attempted"18 19# T4-friendly defaults: fp16 on CUDA, no gradients, eval mode.20torch.set_grad_enabled(False)21if DEVICE == "cuda":22 torch.backends.cudnn.benchmark = True23 torch.backends.cuda.matmul.allow_tf32 = True # harmless on T4; useful on newer GPUs24 try:25 torch.set_float32_matmul_precision("high")26 except Exception:27 pass28 29model_kwargs = {}30if DEVICE == "cuda":31 model_kwargs["torch_dtype"] = torch.float1632 33tok = AutoTokenizer.from_pretrained(MODEL_NAME)34model = AutoModelForCausalLM.from_pretrained(MODEL_NAME, **model_kwargs).to(DEVICE)35model.eval()36model.config.use_cache = True37 38# Causal LMs such as GPT-2 often have no pad token. Right padding is safe here39# because we score only real positions and pass an attention mask.40if tok.pad_token_id is None:41 if tok.eos_token_id is not None:42 tok.pad_token = tok.eos_token43 else:44 tok.add_special_tokens({"pad_token": "<|pad|>"})45 model.resize_token_embeddings(len(tok))46 47model.config.pad_token_id = tok.pad_token_id48PAD_ID = tok.pad_token_id49 50# Pre-tokenized hot candidates. Exact string match only.51# Add more dictation commands/ambiguities here as they become common.52PRETOKENIZED_CANDIDATES = {53 text: tok.encode(text, add_special_tokens=False)54 for text in ("column", "colon", ":")55}56 57# Optional torch.compile. This can reduce Python/dispatch overhead after warmup, but58# may not help all GPU/model/shape combinations, so it is deliberately best-effort.59if ENABLE_TORCH_COMPILE and hasattr(torch, "compile"):60 try:61 model = torch.compile(model, mode="reduce-overhead", fullgraph=False)62 COMPILE_STATUS = "torch.compile attempted: enabled"63 except Exception as exc:64 COMPILE_STATUS = f"torch.compile attempted: failed ({type(exc).__name__})"65elif not ENABLE_TORCH_COMPILE:66 COMPILE_STATUS = "torch.compile disabled by TORCH_COMPILE=0"67else:68 COMPILE_STATUS = "torch.compile unavailable in this PyTorch"69 70 71def cuda_sync() -> None:72 if DEVICE == "cuda":73 torch.cuda.synchronize()74 75 76def now_ms() -> float:77 return time.perf_counter() * 1000.078 79 80def safe_exp(x: float) -> str:81 try:82 return f"{math.exp(x):.6e}"83 except OverflowError:84 return "inf (overflow)"85 except Exception:86 return "-"87 88 89def is_finite(x: float) -> bool:90 return x is not None and math.isfinite(x)91 92 93def encode_candidate(candidate: str) -> list[int]:94 cached = PRETOKENIZED_CANDIDATES.get(candidate)95 if cached is not None:96 # Return a copy so downstream code can treat all candidate id lists normally.97 return list(cached)98 return tok.encode(candidate, add_special_tokens=False)99 100 101def encode_inputs(context: str, candidates: list[str]):102 ctx_ids_cpu = tok.encode(context, return_tensors="pt").squeeze(0)103 cand_ids_list = [encode_candidate(c) for c in candidates]104 return ctx_ids_cpu, cand_ids_list105 106 107def token_strings(cand_ids: list[int]) -> list[str]:108 return [tok.decode([token_id]) for token_id in cand_ids]109 110 111def empty_result(cand_ids: list[int], message: str):112 return {113 "total_score": None,114 "score_kind": "not_scored",115 "token_list": token_strings(cand_ids),116 "num_tokens": len(cand_ids),117 "per_token_scores": [],118 "message": message,119 }120 121 122def make_result(cand_ids: list[int], per_token_scores: list[float], score_kind: str):123 return {124 "total_score": float(sum(per_token_scores)),125 "score_kind": score_kind, # "raw_logit" for one-token fast path, "logprob" otherwise126 "token_list": token_strings(cand_ids),127 "num_tokens": len(cand_ids),128 "per_token_scores": [float(x) for x in per_token_scores],129 "message": "",130 }131 132 133def repeat_cache_to_batch(cache, batch_size: int):134 """Repeat a KV cache from batch size 1 to batch_size."""135 if cache is None:136 return None137 138 # Newer Transformers Cache objects may expose batch_repeat_interleave.139 if hasattr(cache, "batch_repeat_interleave"):140 maybe_returned = cache.batch_repeat_interleave(batch_size)141 return cache if maybe_returned is None else maybe_returned142 143 if torch.is_tensor(cache):144 # expand avoids work, contiguous makes it safe for all attention implementations.145 return cache.expand(batch_size, *([-1] * (cache.dim() - 1))).contiguous()146 147 if isinstance(cache, tuple):148 return tuple(repeat_cache_to_batch(x, batch_size) for x in cache)149 150 if isinstance(cache, list):151 return [repeat_cache_to_batch(x, batch_size) for x in cache]152 153 raise TypeError(f"Unsupported cache type: {type(cache)}")154 155 156def score_single_token_fast(ctx_ids_cpu: torch.Tensor, cand_ids_list: list[list[int]]):157 """158 Fastest path: one context forward, then gather raw next-token logits.159 160 No softmax/log_softmax is needed. For candidates scored from the same next-token161 distribution, logprob(A) - logprob(B) == logit(A) - logit(B) because the shared162 normalization denominator cancels.163 """164 ctx_ids = ctx_ids_cpu.unsqueeze(0).to(DEVICE)165 cuda_sync()166 t0 = now_ms()167 with torch.inference_mode():168 outputs = model(input_ids=ctx_ids, use_cache=False)169 next_logits = outputs.logits[:, -1, :].float().squeeze(0)170 token_ids = torch.tensor([ids[0] for ids in cand_ids_list], dtype=torch.long, device=DEVICE)171 raw_logits = next_logits.index_select(0, token_ids)172 cuda_sync()173 model_ms = now_ms() - t0174 175 scores = raw_logits.detach().cpu().tolist()176 results = [make_result(ids, [score], "raw_logit") for ids, score in zip(cand_ids_list, scores)]177 return results, model_ms, "single context forward for 1-token candidates; raw logits, no softmax"178 179 180def score_full_batch(ctx_ids_cpu: torch.Tensor, cand_ids_list: list[list[int]]):181 """Score context+candidate sequences in one right-padded batched forward."""182 ctx_len = int(ctx_ids_cpu.numel())183 seqs = [torch.cat([ctx_ids_cpu, torch.tensor(ids, dtype=torch.long)]) for ids in cand_ids_list]184 max_len = max(int(seq.numel()) for seq in seqs)185 batch_size = len(seqs)186 187 input_ids = torch.full((batch_size, max_len), PAD_ID, dtype=torch.long)188 attention_mask = torch.zeros((batch_size, max_len), dtype=torch.long)189 for row, seq in enumerate(seqs):190 seq_len = int(seq.numel())191 input_ids[row, :seq_len] = seq192 attention_mask[row, :seq_len] = 1193 194 input_ids = input_ids.to(DEVICE)195 attention_mask = attention_mask.to(DEVICE)196 197 cuda_sync()198 t0 = now_ms()199 with torch.inference_mode():200 outputs = model(input_ids=input_ids, attention_mask=attention_mask, use_cache=False)201 logits = outputs.logits.float()202 cuda_sync()203 model_ms = now_ms() - t0204 205 results = []206 for row, cand_ids in enumerate(cand_ids_list):207 per_token_lps = []208 for i, token_id in enumerate(cand_ids):209 pred_pos = ctx_len + i - 1210 step_logits = logits[row, pred_pos, :]211 lp = float((step_logits[token_id] - torch.logsumexp(step_logits, dim=-1)).item())212 per_token_lps.append(lp)213 results.append(make_result(cand_ids, per_token_lps, "logprob"))214 215 return results, model_ms, "one full batched forward pass"216 217 218def score_with_prefix_cache(ctx_ids_cpu: torch.Tensor, cand_ids_list: list[list[int]]):219 """Score long-context multi-token candidates using one context pass + one batched continuation pass."""220 ctx_ids = ctx_ids_cpu.unsqueeze(0).to(DEVICE)221 ctx_len = int(ctx_ids.shape[1])222 batch_size = len(cand_ids_list)223 max_cand_len = max(len(ids) for ids in cand_ids_list)224 225 cuda_sync()226 t0 = now_ms()227 with torch.inference_mode():228 prefix_outputs = model(input_ids=ctx_ids, use_cache=True)229 prefix_logits = prefix_outputs.logits[:, -1, :].float()230 suffix_logits = None231 232 if max_cand_len > 1:233 past = repeat_cache_to_batch(prefix_outputs.past_key_values, batch_size)234 suffix_len = max_cand_len - 1235 suffix_input_ids = torch.full((batch_size, suffix_len), PAD_ID, dtype=torch.long, device=DEVICE)236 suffix_attention = torch.zeros((batch_size, suffix_len), dtype=torch.long, device=DEVICE)237 238 for row, ids in enumerate(cand_ids_list):239 prefix_ids = ids[:-1]240 if prefix_ids:241 n = len(prefix_ids)242 suffix_input_ids[row, :n] = torch.tensor(prefix_ids, dtype=torch.long, device=DEVICE)243 suffix_attention[row, :n] = 1244 245 full_attention = torch.cat(246 [torch.ones((batch_size, ctx_len), dtype=torch.long, device=DEVICE), suffix_attention],247 dim=1,248 )249 suffix_outputs = model(250 input_ids=suffix_input_ids,251 attention_mask=full_attention,252 past_key_values=past,253 use_cache=False,254 )255 suffix_logits = suffix_outputs.logits.float()256 cuda_sync()257 model_ms = now_ms() - t0258 259 results = []260 prefix_log_denom = torch.logsumexp(prefix_logits[0], dim=-1)261 for row, cand_ids in enumerate(cand_ids_list):262 per_token_lps = []263 first_token_id = cand_ids[0]264 first_lp = float((prefix_logits[0, first_token_id] - prefix_log_denom).item())265 per_token_lps.append(first_lp)266 267 for i in range(1, len(cand_ids)):268 step_logits = suffix_logits[row, i - 1, :]269 token_id = cand_ids[i]270 lp = float((step_logits[token_id] - torch.logsumexp(step_logits, dim=-1)).item())271 per_token_lps.append(lp)272 273 results.append(make_result(cand_ids, per_token_lps, "logprob"))274 275 return results, model_ms, "shared-prefix cache + batched candidate pass"276 277 278def score_candidates(context: str, candidates: list[str]):279 """280 Compute scores for all candidates.281 282 Candidates are scored exactly as typed. No leading space is added.283 One-token candidates use raw logits without softmax/log_softmax; multi-token284 candidates use summed log probabilities.285 """286 total_t0 = now_ms()287 ctx_ids_cpu, cand_ids_list = encode_inputs(context, candidates)288 289 if any(len(ids) == 0 for ids in cand_ids_list):290 results = [291 empty_result(ids, "Candidate tokenized to an empty sequence. Type the candidate exactly as you want it scored.")292 if len(ids) == 0293 else empty_result(ids, "Not scored because another candidate tokenized to an empty sequence.")294 for ids in cand_ids_list295 ]296 return results, 0.0, now_ms() - total_t0, "not scored"297 298 max_cand_len = max(len(ids) for ids in cand_ids_list)299 ctx_len = int(ctx_ids_cpu.numel())300 301 try:302 if max_cand_len == 1:303 results, model_ms, mode = score_single_token_fast(ctx_ids_cpu, cand_ids_list)304 elif ctx_len <= FULL_BATCH_CONTEXT_THRESHOLD:305 results, model_ms, mode = score_full_batch(ctx_ids_cpu, cand_ids_list)306 else:307 results, model_ms, mode = score_with_prefix_cache(ctx_ids_cpu, cand_ids_list)308 except Exception:309 # Robust fallback for any model/Transformers cache/compile incompatibility.310 results, model_ms, mode = score_full_batch(ctx_ids_cpu, cand_ids_list)311 mode = f"fallback: {mode}"312 313 total_ms = now_ms() - total_t0314 return results, model_ms, total_ms, mode315 316 317def compare_candidates(context, candA, candB, use_len_norm):318 request_t0 = now_ms()319 320 errors = []321 if not context.strip():322 errors.append("Please enter a context.")323 if not candA.strip():324 errors.append("Please enter Candidate A.")325 if not candB.strip():326 errors.append("Please enter Candidate B.")327 if errors:328 msg = " ".join(errors)329 return f"<div style='color:#b00020;font-weight:600'>{msg}</div>", "", ""330 331 scored, model_ms, scoring_ms, inference_mode = score_candidates(context, [candA, candB])332 resA, resB = scored[0], scored[1]333 rawA = resA["total_score"]334 rawB = resB["total_score"]335 nA = resA["num_tokens"]336 nB = resB["num_tokens"]337 338 if not (is_finite(rawA) and is_finite(rawB)):339 return (340 "<div style='color:#b00020;font-weight:600'>Numerical or tokenization issue. "341 "Try shorter context, a smaller model, or check the candidate text.</div>",342 summarize_candidate("Candidate A", candA, resA),343 summarize_candidate("Candidate B", candB, resB),344 )345 346 if use_len_norm:347 scoreA = rawA / nA348 scoreB = rawB / nB349 label_suffix = " (per-token)"350 else:351 scoreA = rawA352 scoreB = rawB353 label_suffix = ""354 355 diff = scoreA - scoreB356 if abs(diff) <= EPS:357 winner = "Tie"358 win_color = "#92400e"359 elif diff > 0:360 winner = "Candidate A"361 win_color = "#166534"362 else:363 winner = "Candidate B"364 win_color = "#1d4ed8"365 366 request_ms = now_ms() - request_t0367 ratio_str = safe_exp(diff)368 score_kind = resA["score_kind"] if resA["score_kind"] == resB["score_kind"] else "mixed"369 ratio_label = "exp(raw-logit difference)" if score_kind == "raw_logit" else "odds A/B"370 371 headline = (372 f"<div style='padding:14px;border-radius:12px;background:#f8fafc;"373 f"border:1px solid #e2e8f0;margin-bottom:10px'>"374 f"<div style='font-size:20px;font-weight:800;color:{win_color};'>Winner: {winner}{label_suffix}</div>"375 f"<div style='margin-top:6px;font-size:16px;'>"376 f"{ratio_label}{label_suffix} = <b>{ratio_str}</b> | "377 f"score diff A-B{label_suffix} = <b>{diff:.6f}</b>"378 f"</div>"379 f"<div style='margin-top:6px;font-size:16px;'>"380 f"Model inference = <b>{model_ms:.2f} ms</b> | "381 f"Scoring total = <b>{scoring_ms:.2f} ms</b> | "382 f"Request function = <b>{request_ms:.2f} ms</b>"383 f"</div>"384 f"<div style='margin-top:6px;color:#475569'>"385 f"Mode: {inference_mode} on {DEVICE}. Compile: {COMPILE_STATUS}. "386 f"Pre-tokenized exact candidates: {', '.join(repr(k) for k in PRETOKENIZED_CANDIDATES.keys())}. "387 f"Candidates are scored exactly as typed; no leading space is added. "388 f"{'Per-token uses average score.' if use_len_norm else 'Whole-sequence comparison.'}"389 f"</div></div>"390 )391 392 return headline, summarize_candidate("Candidate A", candA, resA), summarize_candidate("Candidate B", candB, resB)393 394 395def summarize_candidate(label: str, cand: str, res: dict) -> str:396 if res["total_score"] is None:397 return (398 f"**{label}**: {repr(cand)}\n\n"399 f"Tokenization: {res['token_list']}\n"400 f"Tokens: {res['num_tokens']}\n"401 f"{res['message']}"402 )403 404 per_token = ", ".join(f"{x:.4f}" for x in res["per_token_scores"])405 score_kind = res["score_kind"]406 if score_kind == "raw_logit":407 score_lines = (408 f"Raw logit score: {res['total_score']:.6f}\n"409 f"Per-token raw logits: [{per_token}]\n"410 "Sequence probability: not computed in one-token fast path\n"411 )412 else:413 score_lines = (414 f"Total logprob: {res['total_score']:.6f}\n"415 f"Sequence probability: {math.exp(res['total_score']):.6e}\n"416 f"Per-token logprobs: [{per_token}]\n"417 )418 419 return (420 f"**{label}**: {repr(cand)}\n\n"421 f"Tokenization: {res['token_list']}\n"422 f"{score_lines}"423 f"Tokens: {res['num_tokens']}"424 )425 426 427def swap(a, b):428 return b, a429 430 431with gr.Blocks(title="Ultra-Fast Two-Candidate Next-Token Comparator") as demo:432 gr.Markdown(433 "# Ultra-Fast Two-Candidate Next-Word/Token Comparator\n"434 "Compare candidate continuations from a pretrained causal LM.\n"435 "- One-token candidates use raw logits only: no softmax/log_softmax.\n"436 "- Exact candidates `column`, `colon`, and `:` are pre-tokenized at startup.\n"437 "- The app doesn't attempt `torch.compile(..., mode='reduce-overhead')` unless `TORCH_COMPILE=1`.\n"438 "- Candidates are scored exactly as typed; no leading space is automatically added.\n"439 "- Multi-token candidates still use summed log probabilities."440 )441 with gr.Row():442 context = gr.Textbox(label="Context (prompt)", lines=6, placeholder="Paste prior text here...")443 with gr.Row():444 candA = gr.Textbox(label="Candidate A", value="colon")445 candB = gr.Textbox(label="Candidate B", value=":")446 with gr.Row():447 use_len_norm = gr.Checkbox(value=False, label="Use length normalization (average score per token)")448 with gr.Row():449 btn_compare = gr.Button("Compare", variant="primary")450 btn_swap = gr.Button("Swap A <-> B")451 452 winner_html = gr.HTML()453 with gr.Row():454 summaryA = gr.Markdown()455 summaryB = gr.Markdown()456 457 btn_compare.click(458 fn=compare_candidates,459 inputs=[context, candA, candB, use_len_norm],460 outputs=[winner_html, summaryA, summaryB],461 )462 463 btn_swap.click(fn=swap, inputs=[candA, candB], outputs=[candA, candB])464 465 466demo.launch()467