| """Typed think decode scaffold (Loop 2 P0/P1). No weights. |
| |
| Gateway-safe: this module must not import train/eval stacks. |
| |
| P0: prefill ``Answer: <think>\\n``. |
| P1: allow a real working budget (192 tokens, same as GRPO ``max_new``). |
| Force-close if still open. Cut only on a triple 20-char repeat (loop XOR), |
| not on the first equation. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import re |
|
|
| THINK_CUE = "Show brief working." |
| THINK_PREFILL = "Answer: <think>\n" |
| THINK_CLOSE = "</think>" |
| MAX_THINK_BODY_TOKENS = 192 |
|
|
| _LOOP_RE = re.compile(r"(.{20,}?)\1\1", re.I | re.DOTALL) |
|
|
|
|
| def whitespace_tokens(text: str) -> list[str]: |
| return [t for t in (text or "").split() if t] |
|
|
|
|
| def trim_think_body(body: str, *, max_tokens: int = MAX_THINK_BODY_TOKENS) -> str: |
| """Keep multi-step work. Cut a triple-repeat loop, then cap at max_tokens.""" |
| raw = (body or "").replace("<|endoftext|>", "").strip() |
| if not raw: |
| return "" |
| m = _LOOP_RE.search(raw) |
| if m: |
| raw = raw[: m.start() + len(m.group(1))].strip() |
| toks = whitespace_tokens(raw) |
| if len(toks) > max_tokens: |
| return " ".join(toks[:max_tokens]) |
| return raw |
|
|
