app.py

"""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())