"""Side chat — a Gradio port of the browser steganacrostics app. Completely normal text assistant, with a secret talking on the side: every line of the answer starts with successive letters of a hidden "secret" word (an acrostic), produced by grammar-constrained decoding. A list-vs-prose classifier auto-picks the render mode, and an optional local-crossing search spends extra attention at each constraint cliff so the forced letters read as the natural next word. Runs the model locally on CPU with PyTorch transformers (the remote Inference API can't do custom logits processing, which is the whole point here). """ from __future__ import annotations import os import re import threading import queue import time import torch from transformers import ( AutoModelForCausalLM, AutoTokenizer, LogitsProcessorList, TextIteratorStreamer, ) import gradio as gr from grammar import compile_acrostic, union_grammars from logits import GrammarLogitsProcessor, build_token_text_table from tokinfo import build_tok_info from classifier import classify, DEFAULT_VARIANT from crossing_search import generate_crossing_search # Default to MiniCPM5-1B (OpenBMB); override with SIDECHAT_MODEL, e.g. # SIDECHAT_MODEL=LiquidAI/LFM2.5-350M for the smaller, faster original. MODEL_ID = os.environ.get("SIDECHAT_MODEL", "openbmb/MiniCPM5-1B") DEVICE = "cpu" # pure CPU by request LIST_SYSTEM = ( "You are a helpful assistant. Answer as a plain bulleted list — one short " "item per line. Do not use markdown, bold text, headings, code, or numbered " "lists." ) PROSE_SYSTEM = ( "You are a helpful assistant. Answer in plain prose. Do not use markdown, " "bold text, headings, code, or bulleted/numbered lists." ) class Context: """Everything the generation + classifier code needs, built once at startup.""" def __init__(self): print(f"loading {MODEL_ID} on {DEVICE}…", flush=True) self.tokenizer = AutoTokenizer.from_pretrained(MODEL_ID) self.model = ( AutoModelForCausalLM.from_pretrained(MODEL_ID, dtype=torch.float32) .to(DEVICE) .eval() ) self.model.device # noqa: B018 (touch to confirm) vocab = self.model.config.vocab_size t0 = time.perf_counter() self.token_text = build_token_text_table(self.tokenizer, vocab) print(f"token table built in {time.perf_counter() - t0:.1f}s ({vocab} tokens)", flush=True) eos = set() def add_eos(x): if x is None: return if isinstance(x, (list, tuple)): for y in x: add_eos(y) else: eos.add(int(x)) add_eos(self.tokenizer.eos_token_id) add_eos(getattr(self.model.config, "eos_token_id", None)) add_eos(getattr(self.model.generation_config, "eos_token_id", None)) self.eos_token_ids = sorted(eos) pad = self.tokenizer.pad_token_id if pad is None: pad = getattr(self.model.generation_config, "pad_token_id", None) if pad is None: pad = self.eos_token_ids[0] self.pad_token_id = int(pad) self.tok_info = build_tok_info(self.token_text, self.eos_token_ids) print("context ready.", flush=True) CTX = Context() # Case-insensitive acrostic; in list mode the very first ` * ` prefix is optional # (some models start with a preamble-free letter, others don't). def build_grammar(secret, list_mode, max_line): if not list_mode: return compile_acrostic(secret, list_prefix="", max_line=max_line, case_insensitive=True) with_prefix = compile_acrostic( secret, list_prefix=" * ", max_line=max_line, case_insensitive=True, first_line_prefix=True ) without_prefix = compile_acrostic( secret, list_prefix=" * ", max_line=max_line, case_insensitive=True, first_line_prefix=False ) return union_grammars([with_prefix, without_prefix]) def check_acrostic(output, secret): """Find a window of line-initial letters matching the secret (case-insensitive, bullets stripped). Returns (ok, firsts).""" lines = [l.strip() for l in output.split("\n")] lines = [l for l in lines if l] def strip(l): return re.sub(r"^\*?\s*", "", l) firsts = [(strip(l)[:1] or "") for l in lines] n = len(secret) for i in range(0, len(firsts) - n + 1): if all(firsts[i + j].lower() == secret[j].lower() for j in range(n)): return True, "".join(firsts[i:i + n]) return False, "".join(firsts) # --- Classifier: drive the list/prose checkbox ------------------------------ def classify_fn(prompt): if not (prompt or "").strip(): return gr.update(), "enter a prompt to detect list vs. prose" pred, raw = classify(CTX, prompt, DEFAULT_VARIANT) label = "list" if pred else "prose" return pred, f"detected **{label}** (classifier raw: {raw!r})" def maybe_detect(prompt, list_mode, auto_detect): """Runs before Generate: when auto-detect is on, classify the prompt and set the list/prose checkbox from it. Otherwise leave the manual choice alone.""" if auto_detect and (prompt or "").strip(): pred, raw = classify(CTX, prompt, DEFAULT_VARIANT) return pred, f"detected **{'list' if pred else 'prose'}** (raw {raw!r}) — generating…" return list_mode, gr.update() # --- Generation ------------------------------------------------------------- def _run_in_thread(target): """Run target() in a daemon thread; return a queue it pushes to. target receives the queue and must push a None sentinel when finished.""" q = queue.Queue() threading.Thread(target=target, args=(q,), daemon=True).start() return q def generate_fn(prompt, secret, list_mode, max_line, crossing, k, j, R, min_line): # Strip spaces: a multi-word secret spells its letters across lines; spaces # would force odd punctuation-prefixed "word-break" lines. The field still # shows the spaced version; the acrostic uses only the letters. secret = re.sub(r"\s+", "", (secret or "").strip()) if not secret: yield "(secret is empty — open ⚙️ Settings and set one)", "", "" return list_mode = bool(list_mode) max_line = max(1, int(max_line or 80)) system_prompt = LIST_SYSTEM if list_mode else PROSE_SYSTEM try: grammar = build_grammar(secret, list_mode, max_line) except Exception as e: # noqa: BLE001 yield f"grammar build error: {e}", "", "" return # --- Local-crossing search (prose only) --------------------------------- if crossing and not list_mode: k = max(0, int(k or 4)) j = max(0, int(j or 3)) R = max(0, int(R or 4)) min_line = min(max_line, max(0, int(min_line or 30))) status = f"generating (local-crossing search · k={k}, j={j}, R={R}, minLine={min_line})…" committed = [""] t0 = time.perf_counter() def worker(q): def on_line(line_text, info): committed[0] += line_text q.put(committed[0]) try: res = generate_crossing_search( CTX, grammar, secret, max_line, prompt, system_prompt, k=k, j=j, R=R, min_line=min_line, on_line=on_line, ) q.put(("done", res)) except Exception as e: # noqa: BLE001 q.put(("error", str(e))) q.put(None) q = _run_in_thread(worker) result = None yield "", "", status while True: item = q.get() if item is None: break if isinstance(item, tuple) and item[0] == "done": result = item[1] elif isinstance(item, tuple) and item[0] == "error": yield committed[0], f"error: {item[1]}", "error" else: yield item, "", status elapsed = time.perf_counter() - t0 text = result["text"] if result else committed[0] per_line = result["per_line"] if result else [] n_moved = sum(1 for p in per_line if p.get("r", 0) > 0) ok, firsts = check_acrostic(text, secret) metrics = ( f"local-crossing · {elapsed:.2f}s · {len(per_line)} lines · " f"{n_moved} breaks moved · acrostic {'OK' if ok else 'MISS'} ({firsts})" ) yield text, metrics, "done (local-crossing search)." return # --- Plain grammar-constrained greedy (token-streamed) ------------------ proc = GrammarLogitsProcessor(grammar, CTX.tokenizer, CTX.token_text, CTX.eos_token_ids) messages = [ {"role": "system", "content": system_prompt}, {"role": "user", "content": prompt}, ] enc = CTX.tokenizer.apply_chat_template( messages, add_generation_prompt=True, return_tensors="pt", return_dict=True ).to(DEVICE) streamer = TextIteratorStreamer(CTX.tokenizer, skip_prompt=True, skip_special_tokens=True) gen_kwargs = dict( **enc, max_new_tokens=400, do_sample=False, logits_processor=LogitsProcessorList([proc]), streamer=streamer, pad_token_id=CTX.pad_token_id, ) t0 = time.perf_counter() thread = threading.Thread(target=CTX.model.generate, kwargs=gen_kwargs, daemon=True) thread.start() acc = "" t_first = None tokens = 0 chars = 0 yield "", "", "generating (grammar-constrained)…" for chunk in streamer: if not chunk: continue if t_first is None: t_first = time.perf_counter() tokens += 1 chars += len(chunk) acc += chunk gen_s = max(0.001, time.perf_counter() - t_first) tps = tokens / gen_s ttft = (t_first - t0) yield acc, f"TTFT {ttft:.2f}s · ~{tps:.1f} tok/s · {tokens} tokens · {chars} chars", "generating…" thread.join() wall = time.perf_counter() - t0 s = proc.stats proc_ms = s["total_ms"] ttft = (t_first - t0) if t_first else 0.0 ok, firsts = check_acrostic(acc, secret) metrics = ( f"TTFT {ttft:.2f}s · {tokens} tokens · {chars} chars · wall {wall:.2f}s · " f"mask {proc_ms:.0f}ms ({(proc_ms/1000)/wall*100:.0f}%) · " f"acrostic {'OK' if ok else 'MISS'} ({firsts})" ) yield acc, metrics, "done. edit the secret and/or prompt and click Generate again." # --- UI --------------------------------------------------------------------- with gr.Blocks(title="Side chat") as demo: gr.Markdown("# Side chat") gr.Markdown( "Completely normal text assistant, with talking on the side. Each line " "of the answer secretly starts with the next letter of your **secret** " f"word — grammar-constrained decoding on `{MODEL_ID}`, running locally on CPU." ) prompt = gr.Textbox(label="Prompt", value="what are some easy-to-make home recipes?", lines=2) gr.Examples( examples=[ ["what are some easy-to-make home recipes?"], ["please write a few sentences about regular expressions"], ], inputs=prompt, label="Demo prompts (one detects as a list, one as prose)", ) run = gr.Button("Generate", variant="primary") output = gr.Textbox(label="Output", lines=10, interactive=False) metrics = gr.Markdown("") with gr.Accordion("⚙️ Settings", open=False): secret = gr.Textbox( label="Secret (each line will start with these letters)", value="subtle" ) auto_detect = gr.Checkbox( label="auto-detect list vs. prose on Generate (LLM classifier)", value=True, ) list_mode = gr.Checkbox( label="render as bulleted list (each line prefixed with ` * `) — " "set by auto-detect; uncheck auto-detect to set it manually", value=True, ) # Manual preview: run the classifier without generating (debug aid). detect = gr.Button("🔎 Detect list / prose (preview only)", size="sm") max_line = gr.Number(label="Max chars per line (after the prefix + letter)", value=80, precision=0) gr.Markdown("**Local-crossing search** (prose only) — extra attention at each constraint cliff") crossing = gr.Checkbox( label="enable local-crossing search (greedy line, then pick the break " "that makes the crossing read best; list mode stays greedy)", value=False, ) win_k = gr.Number(label="↳ window before the break (k content tokens)", value=4, precision=0) win_j = gr.Number(label="↳ window after the forced letter (j content tokens)", value=3, precision=0) max_rewind = gr.Number(label="↳ max tokens to trim the break earlier (R; 0 = greedy)", value=4, precision=0) min_line = gr.Number(label="↳ min chars per line (avoid stubby lines; 0 = off)", value=30, precision=0) status = gr.Markdown("ready.") # Manual preview: detect list vs. prose without generating. detect.click(classify_fn, [prompt], [list_mode, status]) prompt.submit(classify_fn, [prompt], [list_mode, status]) # Generate: auto-detect first (updates the checkbox), then generate using it. run.click( maybe_detect, [prompt, list_mode, auto_detect], [list_mode, status] ).then( generate_fn, [prompt, secret, list_mode, max_line, crossing, win_k, win_j, max_rewind, min_line], [output, metrics, status], ) if __name__ == "__main__": demo.queue().launch(theme=gr.themes.Soft())