Spaces:
Paused
Paused
| """Prompt layout, sampling, and streaming decode for booper-chat. | |
| booper-chat was SFT'd on the *pair* layout -- `<bos> prompt <sep> reply <eos>`, | |
| loss on the reply only -- so that is what it has to be served with. Fed the | |
| plain continuation layout it keeps writing more of the user's message instead | |
| of answering it, which is the single easiest way to make this model look | |
| broken. See the model card and babble's `BABBLE_SERVE_LAYOUT=pair` note. | |
| Multi-turn is the same layout with the earlier turns newline-joined into the | |
| prompt, tail-truncated to whatever is left of the 1024-token window after the | |
| reply's reservation. | |
| """ | |
| from __future__ import annotations | |
| import os | |
| import time | |
| from dataclasses import dataclass, field | |
| from pathlib import Path | |
| import torch | |
| from .bpe import BPETokenizer | |
| from .model import Babbler, ModelConfig | |
| BLOCK_SIZE = 1024 | |
| class SamplingParams: | |
| """Defaults follow babble's shipped serving config: temperature 0.5 was | |
| measured to be the single biggest quality lever over 1.0, and the two | |
| repetition knobs are both on because they fail differently. A flat | |
| `repetition_penalty` cannot escape a loop once one is running (live booper | |
| emitted ~85 consecutive "boop" tokens with it nominally active), so the | |
| count-scaled `frequency_penalty` carries the actual loop protection and | |
| keeps growing every repeat; `no_repeat_ngram_size` is the hard backstop | |
| that no penalty can be talked around. | |
| """ | |
| max_new_tokens: int = 96 | |
| temperature: float = 0.7 | |
| top_k: int = 40 | |
| top_p: float = 0.9 | |
| repetition_penalty: float = 1.15 | |
| frequency_penalty: float = 0.12 | |
| presence_penalty: float = 0.0 | |
| no_repeat_ngram_size: int = 3 | |
| seed: int | None = None | |
| def clamped(self) -> "SamplingParams": | |
| """Clamp every knob to a sane band -- these arrive from a public web | |
| form, so nothing here may be trusted to be a number in range.""" | |
| def num(value, lo, hi, fallback): | |
| try: | |
| value = float(value) | |
| except (TypeError, ValueError): | |
| return fallback | |
| if value != value: # NaN | |
| return fallback | |
| return min(hi, max(lo, value)) | |
| return SamplingParams( | |
| max_new_tokens=int(num(self.max_new_tokens, 1, 512, 96)), | |
| temperature=num(self.temperature, 0.0, 2.0, 0.7), | |
| top_k=int(num(self.top_k, 0, 16384, 40)), | |
| top_p=num(self.top_p, 0.0, 1.0, 0.9), | |
| repetition_penalty=num(self.repetition_penalty, 1.0, 2.0, 1.15), | |
| frequency_penalty=num(self.frequency_penalty, 0.0, 2.0, 0.12), | |
| presence_penalty=num(self.presence_penalty, 0.0, 2.0, 0.0), | |
| no_repeat_ngram_size=int(num(self.no_repeat_ngram_size, 0, 8, 3)), | |
| seed=self.seed, | |
| ) | |
| class Turn: | |
| role: str # "user" | "assistant" | |
| content: str | |
| class Booper: | |
| def __init__(self, checkpoint: str | Path, tokenizer_json: str | Path, *, | |
| threads: int | None = None, quantize: bool | str | None = None) -> None: | |
| threads = threads if threads is not None else int(os.environ.get("BOOPER_THREADS", "0")) or None | |
| if threads: | |
| torch.set_num_threads(max(1, threads)) | |
| try: | |
| torch.set_num_interop_threads(1) | |
| except RuntimeError: | |
| pass # already set by an earlier init in this process | |
| if hasattr(torch.backends, "mkldnn"): | |
| torch.backends.mkldnn.enabled = True | |
| try: | |
| torch.set_flush_denormal(True) | |
| except Exception: | |
| pass | |
| self.tok = BPETokenizer.from_json(tokenizer_json) | |
| blob = torch.load(checkpoint, map_location="cpu", weights_only=False) | |
| config = ModelConfig.from_dict(blob["config"]) | |
| if config.vocab_size != self.tok.vocab_size: | |
| raise ValueError( | |
| f"checkpoint vocab_size {config.vocab_size} != tokenizer " | |
| f"vocab_size {self.tok.vocab_size} -- wrong tokenizer.json" | |
| ) | |
| model = Babbler(config) | |
| model.load_state_dict(blob["model"]) | |
| model.eval() | |
| self.n_params = sum({id(p): p.numel() for p in model.parameters()}.values()) | |
| if quantize is None: | |
| quantize = os.environ.get("BOOPER_QUANTIZE", "auto").strip().lower() | |
| if isinstance(quantize, bool): | |
| quantize = "1" if quantize else "0" | |
| self.quantized = False | |
| self.backend = "fp32" | |
| if quantize != "0": | |
| candidate = _quantize_dynamic(model, config, blob) | |
| if candidate is None: | |
| pass # no quantized engine in this wheel; fp32 it is | |
| elif quantize == "auto": | |
| # int8 is a big win on x86 (fbgemm/VNNI) and a ~2x *loss* on | |
| # Apple qnnpack, so the answer depends entirely on what CPU | |
| # the Space actually landed on. Measuring costs a few hundred | |
| # ms once at startup and is the only way to be right on both. | |
| fp32_s = _decode_bench(model, config) | |
| int8_s = _decode_bench(candidate, config) | |
| if int8_s < fp32_s: | |
| model, self.quantized = candidate, True | |
| self.backend = f"int8 ({int8_s * 1e3:.1f}ms/tok vs fp32 {fp32_s * 1e3:.1f})" if self.quantized \ | |
| else f"fp32 ({fp32_s * 1e3:.1f}ms/tok vs int8 {int8_s * 1e3:.1f})" | |
| else: | |
| model, self.quantized, self.backend = candidate, True, "int8 (forced)" | |
| self.model = model | |
| self.config = config | |
| self.block_size = config.block_size | |
| del blob | |
| # --- prompt layout ---------------------------------------------------- | |
| def build_context(self, turns: list[Turn], reserve: int) -> tuple[list[int], int]: | |
| """`<bos> …history… <sep>`, filling the window from the newest turn back. | |
| Turns are added newest-first until the next one would not fit, so the | |
| model always sees the whole of the current message and as much prior | |
| conversation as the 1024-token window still has room for. Returns the | |
| context and how many turns of history survived the cut. | |
| """ | |
| budget = max(1, self.block_size - 2 - reserve) # bos + sep cost a slot each | |
| newline = self.tok.encode("\n") | |
| pieces: list[list[int]] = [] | |
| used = 0 | |
| kept = 0 | |
| for index, turn in enumerate(reversed(turns)): | |
| ids = self.tok.encode(turn.content.strip()) | |
| cost = len(ids) + (0 if index == 0 else len(newline)) | |
| if used + cost > budget: | |
| if index == 0: | |
| # The current message alone overflows: keep its tail, the | |
| # part being replied to, same rule training truncated with. | |
| pieces.append(ids[-budget:]) | |
| used = budget | |
| kept = 1 | |
| break | |
| pieces.append(ids if index == 0 else ids + newline) | |
| used += cost | |
| kept += 1 | |
| body: list[int] = [] | |
| for ids in reversed(pieces): | |
| body.extend(ids) | |
| return [self.tok.specials.bos, *body, self.tok.specials.sep], kept | |
| # --- sampling --------------------------------------------------------- | |
| def _next_token(self, logits: torch.Tensor, generated: torch.Tensor, p: SamplingParams, | |
| generator: torch.Generator) -> int: | |
| logits = logits.clone() | |
| logits[self._banned] = float("-inf") | |
| if generated.numel(): | |
| counts = torch.bincount(generated, minlength=logits.numel()).to(logits.dtype) | |
| present = counts > 0 | |
| if p.repetition_penalty != 1.0: | |
| # HF-style: positive logits divided, negative multiplied, so a | |
| # seen token gets less likely without flipping sign. | |
| scaled = torch.where(logits < 0, logits * p.repetition_penalty, | |
| logits / p.repetition_penalty) | |
| logits = torch.where(present, scaled, logits) | |
| if p.frequency_penalty or p.presence_penalty: | |
| # Count-scaled, so the discount keeps growing every repeat -- | |
| # this is the one that actually breaks a running loop. | |
| logits = logits - counts * p.frequency_penalty - present.to(logits.dtype) * p.presence_penalty | |
| if p.no_repeat_ngram_size > 1: | |
| n = p.no_repeat_ngram_size | |
| seq = generated.tolist() | |
| if len(seq) >= n - 1: | |
| prefix = tuple(seq[-(n - 1):]) | |
| banned = { | |
| seq[j + n - 1] | |
| for j in range(len(seq) - n + 1) | |
| if tuple(seq[j:j + n - 1]) == prefix | |
| } | |
| if banned: | |
| logits[list(banned)] = float("-inf") | |
| if p.temperature <= 0: | |
| return int(logits.argmax()) | |
| logits = logits / p.temperature | |
| if p.top_k and 0 < p.top_k < logits.numel(): | |
| cutoff = torch.topk(logits, p.top_k).values[-1] | |
| logits = logits.masked_fill(logits < cutoff, float("-inf")) | |
| if 0.0 < p.top_p < 1.0: | |
| ordered, index = torch.sort(logits, descending=True) | |
| mass = torch.cumsum(torch.softmax(ordered, dim=-1), dim=-1) | |
| drop = mass > p.top_p | |
| drop[1:] = drop[:-1].clone() # always keep the top token | |
| drop[0] = False | |
| logits = torch.full_like(logits, float("-inf")).scatter( | |
| 0, index, ordered.masked_fill(drop, float("-inf")) | |
| ) | |
| probs = torch.softmax(logits, dim=-1) | |
| return int(torch.multinomial(probs, 1, generator=generator)) | |
| def _banned(self) -> list[int]: | |
| # Structural ids are inputs, never outputs -- a stray <sep> or <bos> | |
| # in the middle of a reply is always a decode bug, never a message. | |
| s = self.tok.specials | |
| return [s.pad, s.bos, s.sep] | |
| def stream(self, turns: list[Turn], params: SamplingParams | None = None): | |
| """Yield `("token", text)` chunks, then one `("done", stats)`. | |
| Bytes are emitted as soon as they form valid UTF-8 -- a multi-byte | |
| character (an emoji, most of what a Discord model likes to say) spans | |
| several BPE ids, and flushing half of one would render as a replacement | |
| glyph that never un-renders. | |
| """ | |
| p = (params or SamplingParams()).clamped() | |
| generator = torch.Generator() | |
| generator.manual_seed(p.seed if p.seed is not None else int.from_bytes(os.urandom(8), "big") % (2**63)) | |
| context, kept = self.build_context(turns, reserve=p.max_new_tokens) | |
| room = self.block_size - len(context) | |
| budget = max(1, min(p.max_new_tokens, room)) | |
| started = time.perf_counter() | |
| cache = self.model.new_cache(1, max_len=len(context) + budget) | |
| logits = self.model(torch.tensor([context]), cache=cache)[0, -1] | |
| prefill_s = time.perf_counter() - started | |
| eos = self.tok.specials.eos | |
| produced: list[int] = [] | |
| generated = torch.empty(0, dtype=torch.long) | |
| pending = bytearray() | |
| step = torch.empty((1, 1), dtype=torch.long) | |
| stop = "eos" | |
| for i in range(budget): | |
| nxt = self._next_token(logits, generated, p, generator) | |
| if nxt == eos: | |
| break | |
| produced.append(nxt) | |
| generated = torch.cat([generated, torch.tensor([nxt])]) | |
| pending.extend(self.tok.vocab.get(nxt, b"")) | |
| try: | |
| text = pending.decode("utf-8") | |
| except UnicodeDecodeError: | |
| pass # mid-character; hold the bytes until it completes | |
| else: | |
| pending.clear() | |
| if text: | |
| yield "token", text | |
| if i == budget - 1 or cache.length >= cache.max_len: | |
| stop = "length" | |
| break | |
| step[0, 0] = nxt | |
| logits = self.model(step, cache=cache)[0, -1] | |
| if pending: | |
| yield "token", pending.decode("utf-8", errors="replace") | |
| elapsed = time.perf_counter() - started | |
| yield "done", { | |
| "stop": stop, | |
| "tokens": len(produced), | |
| "prompt_tokens": len(context), | |
| "context_used": len(context) + len(produced), | |
| "context_limit": self.block_size, | |
| "turns_kept": kept, | |
| "prefill_ms": round(prefill_s * 1000, 1), | |
| "elapsed_ms": round(elapsed * 1000, 1), | |
| "tok_per_s": round(len(produced) / elapsed, 1) if elapsed > 0 and produced else 0.0, | |
| } | |
| def reply(self, turns: list[Turn], params: SamplingParams | None = None) -> str: | |
| parts = [text for kind, text in self.stream(turns, params) if kind == "token"] | |
| return "".join(parts).strip() | |
| def _quantize_dynamic(model: Babbler, config: ModelConfig, blob: dict) -> Babbler | None: | |
| """A dynamic-int8 copy of `model`, or None if this build has no engine. | |
| `quantize_dynamic` needs `lm_head` to own its weight, but booper ties it to | |
| the token embedding, so the copy unties first. That untying is decode-only | |
| and never reaches a checkpoint -- which is why this builds a fresh model | |
| rather than mutating the one that was passed in. | |
| """ | |
| import torch.nn as nn | |
| try: | |
| if torch.backends.quantized.engine == "none": | |
| # macOS/ARM wheels leave this unset and every prepack then fails | |
| # with NoQEngine; x86 Linux picks fbgemm on its own. | |
| engines = [e for e in torch.backends.quantized.supported_engines if e != "none"] | |
| if not engines: | |
| return None | |
| torch.backends.quantized.engine = engines[0] | |
| clone = Babbler(config) | |
| clone.load_state_dict(blob["model"]) | |
| clone.lm_head.weight = nn.Parameter(clone.tok_emb.weight.detach().clone()) | |
| clone.eval() | |
| return torch.ao.quantization.quantize_dynamic(clone, {nn.Linear}, dtype=torch.qint8) | |
| except Exception: | |
| return None | |
| def _decode_bench(model, config: ModelConfig, steps: int = 12) -> float: | |
| """Median seconds per cached single-token forward -- the thing that decides | |
| how fast a reply feels. Median, not mean, so one scheduler hiccup on a | |
| shared vCPU does not pick the wrong backend.""" | |
| cache = model.new_cache(1, max_len=steps + 8) | |
| step = torch.zeros((1, 1), dtype=torch.long) | |
| model(torch.zeros((1, 4), dtype=torch.long), cache=cache) # warm the kernels | |
| times = [] | |
| for _ in range(steps): | |
| started = time.perf_counter() | |
| model(step, cache=cache) | |
| times.append(time.perf_counter() - started) | |
| times.sort() | |
| return times[len(times) // 2] | |