432 lines
18 KiB
Python
432 lines
18 KiB
Python
"""
|
|
tokenmax.py — Deterministic token-maxing post-processor for compressed think blocks.
|
|
|
|
Applies ONLY substitutions that are verified to save tokens on the Qwen 248K tokenizer
|
|
(OmniCoder-9B / Qwen3.5). Every substitution was tested in-context (not isolation) to
|
|
confirm real token savings without boundary interference.
|
|
|
|
Design:
|
|
- LLM does semantic compression (what to keep vs drop)
|
|
- This code enforces consistent notation deterministically
|
|
- GUARD: only returns the processed version if ntok(result) < ntok(original)
|
|
- Idempotent: safe to run multiple times
|
|
|
|
Usage:
|
|
from caveman.compress.tokenmax import tokenmax, tokenmax_with_stats
|
|
compressed = tokenmax(think_text, tokenizer)
|
|
compressed, stats = tokenmax_with_stats(think_text, tokenizer)
|
|
|
|
Verified: 2026-05-31 on Qwen 248K vocab. 27/28 substitutions save in-context.
|
|
Zero false positives. One zero-effect (贪心 for "greedy" — boundary-dependent).
|
|
"""
|
|
|
|
import re
|
|
from typing import Optional
|
|
|
|
# ── Phase 1: Filler drops ──────────────────────────────────────────────
|
|
# Phrases that carry zero information in compressed reasoning.
|
|
# Only patterns that are NEVER load-bearing in a think block.
|
|
_FILLER_PATTERNS = [
|
|
# Metacognition (the model narrating its own process)
|
|
r"\bI need to\b",
|
|
r"\bwe need to\b",
|
|
r"\bI will\b",
|
|
r"\bI'll\b",
|
|
r"\blet me\b",
|
|
r"\blet's\b",
|
|
r"\bI want to\b",
|
|
r"\bI should\b",
|
|
r"\bwe should\b",
|
|
# Hedging
|
|
r"\bprobably\b",
|
|
r"\bbasically\b",
|
|
r"\bessentially\b",
|
|
r"\bit seems like\b",
|
|
# Filler transitions
|
|
r"\bin order to\b",
|
|
r"\bfirst of all\b",
|
|
r"\bin other words\b",
|
|
r"\bon the other hand\b",
|
|
r"\bmore specifically\b",
|
|
r"\bto be more precise\b",
|
|
r"\band so on\b",
|
|
# Conversational padding (require word boundary at end to avoid "Greatest", "Perfectly")
|
|
r"\bGreat\b[,!.]?\s*",
|
|
r"\bPerfect\b[,!.]?\s*",
|
|
# Obvious statements
|
|
r"\bAs (?:we|you) can see\b",
|
|
r"\bAs mentioned (?:above|earlier|before)\b",
|
|
]
|
|
|
|
# ── Phase 2: Phrase → cheapest token ───────────────────────────────────
|
|
# Ordered LONGEST FIRST to prevent partial matches.
|
|
# Each entry: (regex_pattern, replacement, category)
|
|
# Categories: 'cjk', 'symbol', 'abbrev' — for stats tracking.
|
|
_SUBSTITUTIONS = [
|
|
# ── COMPOUND PATTERNS FIRST (must fire before their components) ──
|
|
|
|
# Verbose comparison phrases (+5t savings)
|
|
(r'\bis\s+greater\s+than\s+or\s+equal\s+to\b', '≥', 'symbol'), # +5t
|
|
(r'\bis\s+less\s+than\s+or\s+equal\s+to\b', '≤', 'symbol'), # +5t
|
|
|
|
# Verbose discourse phrases (+3t savings)
|
|
(r'\bwe\s+can\s+see\s+that\b', '可知', 'cjk'), # +3t
|
|
(r'\bat\s+the\s+same\s+time\b', '同时', 'cjk'), # +3t
|
|
(r'\bthat\s+is\s+to\s+say\b', '即', 'cjk'), # +3t
|
|
|
|
# Multi-word phrases (+2t savings)
|
|
(r'\bin\s+this\s+case\b', '此时', 'cjk'), # +2t
|
|
(r'\bthe\s+number\s+of\b', '个数', 'cjk'), # +2t
|
|
(r'\bis\s+equal\s+to\b', '等于', 'cjk'), # +2t
|
|
|
|
# Complexity boilerplate (biggest per-occurrence savings)
|
|
(r'[Oo]\(n\)\s*time[,;]?\s*[Oo]\(n\)\s*space\.?', 'O(n|n).', 'abbrev'),
|
|
(r'[Oo]\(n\)\s*time[,;]?\s*[Oo]\(1\)\s*space\.?', 'O(n|1).', 'abbrev'),
|
|
(r'[Oo]\(n\s*log\s*n\)\s*time[,;]?\s*[Oo]\(n\)\s*space', 'O(n㏒n|n)', 'abbrev'),
|
|
(r'[Oo]\(n\s*log\s*n\)\s*time[,;]?\s*[Oo]\(1\)\s*space', 'O(n㏒n|1)', 'abbrev'),
|
|
(r'[Tt]ime\s*complexity[:\s]+', 'T=', 'abbrev'),
|
|
(r'[Ss]pace\s*complexity[:\s]+', 'S=', 'abbrev'),
|
|
|
|
# Multi-word compounds (BEFORE their single-word components)
|
|
(r'\bassume without loss of generality\b', '设 不妨', 'cjk'), # before "assume"
|
|
(r'\bproof by contradiction\b', '反证', 'cjk'), # before "proof", "contradiction"
|
|
(r'\bnecessary and sufficient\b', '充要', 'cjk'), # before "sufficient"
|
|
(r'\bnot equal(?:\s+to)?\b', '≠', 'symbol'), # before "is not", "does not"
|
|
(r'\bif and only if\b', 'iff', 'abbrev'), # before "for all"
|
|
(r'\bmuch greater than\b', '≫', 'symbol'), # before "greater than"
|
|
(r'\bkeep track(?:\s+of)?\b', '记录', 'cjk'), # before article strip
|
|
(r'\bin ascending order\b', 'asc', 'abbrev'),
|
|
(r'\bin descending order\b', 'desc', 'abbrev'),
|
|
(r'\bmaximum value\b', '最大值', 'cjk'),
|
|
(r'\bminimum value\b', '最小值', 'cjk'),
|
|
(r'\breturn value\b', '返回値', 'cjk'),
|
|
(r'\brather than\b', '而非', 'cjk'),
|
|
(r'\baccording to\b', '按照', 'cjk'),
|
|
|
|
# DS compounds (before components)
|
|
(r'\bdoubly linked list\b', 'DLL', 'abbrev'), # BEFORE "linked list"
|
|
(r'\bbinary indexed tree\b', 'BIT', 'abbrev'), # BEFORE "binary"
|
|
(r'\bminimum spanning tree\b', 'MST', 'abbrev'),
|
|
(r'\bdepth[- ]first search\b', 'DFS', 'abbrev'),
|
|
(r'\bbreadth[- ]first search\b', 'BFS', 'abbrev'),
|
|
(r'\bdynamic programming\b', 'DP', 'abbrev'),
|
|
(r'\bdivide and conquer\b', '分治', 'cjk'),
|
|
(r'\bmonot(?:onic|one)\s*stack\b', '单调栈', 'cjk'),
|
|
(r'\btime limit exceeded\b', '超时', 'cjk'),
|
|
(r'\bout of bounds\b', '越界', 'cjk'),
|
|
(r'\bremove duplicates?\b', '去重', 'cjk'),
|
|
(r'\benumerate all\b', '穷举', 'cjk'),
|
|
(r'\bbinary search\b', '二分', 'cjk'),
|
|
(r'\bsliding window\b', 'sw', 'abbrev'),
|
|
(r'\bunion[- ]find\b', 'UF', 'abbrev'),
|
|
(r'\btopological sort\b', '拓扑序', 'cjk'),
|
|
(r'\bshortest path\b', 'sp', 'abbrev'),
|
|
(r'\blinked list\b', 'LL', 'abbrev'),
|
|
(r'\bpriority queue\b', 'heap', 'abbrev'),
|
|
(r'\bprefix sum\b', 'ps', 'abbrev'),
|
|
(r'\bbrute force\b', '暴力', 'cjk'),
|
|
(r'\bno solution\b', '无解', 'cjk'),
|
|
(r'\bedge cases?\b', '边界', 'cjk'),
|
|
(r'\bbase case\b', 'bc', 'abbrev'),
|
|
(r'\bworst case\b', 'wc', 'abbrev'),
|
|
|
|
# ── SINGLE-WORD SUBSTITUTIONS (safe after compounds consumed) ──
|
|
|
|
# +4t savings
|
|
(r'\bobviously\b', '显然', 'cjk'),
|
|
# +3t savings
|
|
(r'\bredundant\b', '冗余', 'cjk'),
|
|
(r'\bsatisf(?:y|ies|ied)\b', '满足', 'cjk'), # +2t, 333x in data
|
|
(r'\bunsorted\b', '无序', 'cjk'),
|
|
(r'\bdue to\b', '由于', 'cjk'),
|
|
(r'\bhence\b', '故', 'cjk'),
|
|
(r'\bnamely\b', '即', 'cjk'),
|
|
(r'\bassume\b', '设', 'cjk'),
|
|
(r'\bsuppose\b', '设', 'cjk'),
|
|
(r'\bderive\b', '推导', 'cjk'),
|
|
# +2t savings
|
|
(r'\bmonotone\b', '单调', 'cjk'),
|
|
(r'\bconvergent\b', '收敛', 'cjk'),
|
|
(r'\bdivergent\b', '发散', 'cjk'),
|
|
(r'\bcommutative\b', '交换', 'cjk'),
|
|
(r'\bdeterministic\b', '确定', 'cjk'),
|
|
(r'\bprobabilistic\b', '概率', 'cjk'),
|
|
(r'\bprove\b', '证明', 'cjk'),
|
|
(r'\bproof\b', '证明', 'cjk'),
|
|
(r'\bflip\b', '翻转', 'cjk'),
|
|
(r'\bsorted\b(?!\s*[=(\[])', '有序', 'cjk'), # not before = ( [ (assignment/call)
|
|
# +1t savings
|
|
(r'\bbacktrack(?:ing)?\b', '回溯', 'cjk'),
|
|
(r'\btravers(?:e|al|ing)\b', '遍历', 'cjk'),
|
|
(r'\brecursi(?:on|ve|vely)\b', '递归', 'cjk'),
|
|
(r'\bcontradiction\b', '矛盾', 'cjk'),
|
|
(r'\bsufficient\b', '充分', 'cjk'),
|
|
(r'\bequivalent\b', '等价', 'cjk'),
|
|
(r'\bsymmetric\b', '对称', 'cjk'),
|
|
(r'\binvariant\b', '不变', 'cjk'),
|
|
(r'\bexponential\b', '指数', 'cjk'),
|
|
(r'\bpermutation\b', '排列', 'cjk'),
|
|
(r'\badjacent\b', '相邻', 'cjk'),
|
|
(r'\boptimal\b', '最优', 'cjk'),
|
|
(r'\bfeasible\b', '可行', 'cjk'),
|
|
(r'\binduction\b', '归纳', 'cjk'),
|
|
(r'\bmaintain\b', '维护', 'cjk'),
|
|
(r'\bswap\b(?!\s*[=(\[])', '交换', 'cjk'), # not before = ( [ (assignment/call)
|
|
(r'\bcumulative\b', '累积', 'cjk'),
|
|
(r'\bquotient\b', '商', 'cjk'),
|
|
(r'\bmemoiz(?:ation|e)\b', 'memo', 'abbrev'),
|
|
# +2t savings (mined from v19 data)
|
|
(r'\bmathematical\b', '数学', 'cjk'),
|
|
(r'\bcorresponding(?:ly)?\b', '对应', 'cjk'),
|
|
(r'\brequirement\b', '需求', 'cjk'),
|
|
# +1t savings (mined from v19 data)
|
|
(r'\bcomplexity\b', '复杂度', 'cjk'),
|
|
(r'\bsimilarly\b', '同理', 'cjk'),
|
|
(r'\bsubstitut(?:e|ion)\b', '代入', 'cjk'),
|
|
(r'\bincreasing(?:ly)?\b', '递增', 'cjk'),
|
|
(r'\bdecreasing(?:ly)?\b', '递减', 'cjk'),
|
|
(r'\brespectively\b', '分别', 'cjk'),
|
|
(r'\bnecessarily\b', '必然', 'cjk'),
|
|
(r'\btransformation\b', '变换', 'cjk'),
|
|
(r'\bprerequisite\b', '前提', 'cjk'),
|
|
(r'\bconsequently\b', '从而', 'cjk'),
|
|
(r'\boverlapping\b', '重叠', 'cjk'),
|
|
(r'\bcontribut(?:e|ion)\b', '贡献', 'cjk'),
|
|
(r'\bindependent(?:ly)?\b', '独立', 'cjk'),
|
|
(r'\bimpossible\b', '不可能', 'cjk'),
|
|
(r'\biterat(?:e|ion|ing)\b', '迭代', 'cjk'),
|
|
(r'\benumerat(?:e|ion|ing)\b', '枚举', 'cjk'),
|
|
|
|
# ── LOGIC SYMBOLS ──
|
|
(r'\btherefore\b', '⇒', 'symbol'),
|
|
(r'\bthus\b', '⇒', 'symbol'),
|
|
(r'\bsuch that\b', 'st', 'abbrev'),
|
|
(r'\bthere exists?\b', '∃', 'symbol'),
|
|
(r'\bfor each\b', '∀', 'symbol'),
|
|
(r'\bfor every\b', '∀', 'symbol'),
|
|
(r'\bfor all\b', '∀', 'symbol'),
|
|
(r'\bdoes not\b', '¬', 'symbol'),
|
|
(r"\bdoesn't\b", '¬', 'symbol'),
|
|
(r'\bis not\b(?!\s+(?:None|null|undefined|empty|zero|0))', '非', 'cjk'), # protect "is not None" etc
|
|
(r'\bat least\b', '≥', 'symbol'),
|
|
(r'\bat most\b', '≤', 'symbol'),
|
|
(r'\bgreater than\b', '>', 'symbol'),
|
|
(r'\bless than\b', '<', 'symbol'),
|
|
|
|
# ── ARTICLE STRIPPING (last — lowest priority) ──
|
|
(r'\bthe\b\s+(?!(?:same|only|first|last|next|other)\b)', '', 'filler'), # protect "the same", "the only" etc
|
|
(r'\ba\b\s+(?=[bcdfghjklmnpqrstvwxyz])', '', 'filler'),
|
|
(r'\ban\b\s+', '', 'filler'),
|
|
]
|
|
|
|
# ── Compile once ───────────────────────────────────────────────────────
|
|
_FILLER_COMPILED = [(re.compile(p, re.IGNORECASE), '') for p in _FILLER_PATTERNS]
|
|
_SUBS_COMPILED = [(re.compile(p, re.IGNORECASE), r, cat) for p, r, cat in _SUBSTITUTIONS]
|
|
|
|
|
|
def _ntok(text: str, tokenizer) -> int:
|
|
"""Token count using the provided tokenizer."""
|
|
return len(tokenizer.encode(text, add_special_tokens=False))
|
|
|
|
|
|
def _protect_code_fences(text: str) -> tuple[str, list]:
|
|
"""Extract code-fenced blocks, replace with placeholders.
|
|
Returns (text_with_placeholders, list_of_extracted_blocks)."""
|
|
blocks = []
|
|
def _replace(m):
|
|
blocks.append(m.group(0))
|
|
return f'\x00CODEFENCE{len(blocks)-1}\x00'
|
|
# Match ```...``` and inline `...` (non-greedy)
|
|
protected = re.sub(r'```.*?```|`[^`\n]+`', _replace, text, flags=re.DOTALL)
|
|
return protected, blocks
|
|
|
|
|
|
def _restore_code_fences(text: str, blocks: list) -> str:
|
|
"""Restore code-fenced blocks from placeholders."""
|
|
for i, block in enumerate(blocks):
|
|
text = text.replace(f'\x00CODEFENCE{i}\x00', block)
|
|
return text
|
|
|
|
|
|
def _apply_subs(text: str) -> tuple[str, dict]:
|
|
"""Apply all substitutions, return (result, stats).
|
|
Code fences (``` and inline `) are protected from substitution."""
|
|
stats = {'filler_drops': 0, 'cjk': 0, 'symbol': 0, 'abbrev': 0, 'total_subs': 0}
|
|
|
|
# Phase 0: protect code fences from substitution
|
|
text, code_blocks = _protect_code_fences(text)
|
|
|
|
# Phase 1: filler drops
|
|
for pat, repl in _FILLER_COMPILED:
|
|
text, n = pat.subn(repl, text)
|
|
if n:
|
|
stats['filler_drops'] += n
|
|
stats['total_subs'] += n
|
|
|
|
# Phase 2: substitutions
|
|
for pat, repl, cat in _SUBS_COMPILED:
|
|
text, n = pat.subn(repl, text)
|
|
if n:
|
|
stats[cat] = stats.get(cat, 0) + n
|
|
stats['total_subs'] += n
|
|
|
|
# Phase 3: restore code fences
|
|
text = _restore_code_fences(text, code_blocks)
|
|
|
|
# Phase 4: whitespace normalization
|
|
text = re.sub(r'[ \t]+', ' ', text)
|
|
text = re.sub(r'\n{3,}', '\n\n', text)
|
|
text = re.sub(r' *\n *', '\n', text)
|
|
text = text.strip()
|
|
|
|
return text, stats
|
|
|
|
|
|
def tokenmax(text: str, tokenizer, force_cjk: bool = False) -> str:
|
|
"""Apply token-maxing. Returns original if no savings achieved.
|
|
|
|
Args:
|
|
text: The think block content (without <think> tags).
|
|
tokenizer: A HuggingFace tokenizer with .encode() method.
|
|
force_cjk: If True, always return the processed version when CJK
|
|
substitutions were applied, even if total token count increased.
|
|
Use this to maximize CJK adoption in training data.
|
|
|
|
Returns:
|
|
The token-maxed text, or the original if processing didn't save tokens
|
|
(unless force_cjk=True and CJK subs were applied).
|
|
"""
|
|
if not text or not text.strip():
|
|
return text
|
|
|
|
original_tokens = _ntok(text, tokenizer)
|
|
result, stats = _apply_subs(text)
|
|
result_tokens = _ntok(result, tokenizer)
|
|
|
|
# GUARD: only return processed version if it actually saves tokens
|
|
# OVERRIDE: force_cjk bypasses the guard when CJK substitutions were made
|
|
if result_tokens < original_tokens:
|
|
return result
|
|
if force_cjk and stats.get('cjk', 0) > 0:
|
|
return result
|
|
return text
|
|
|
|
|
|
def tokenmax_with_stats(text: str, tokenizer, force_cjk: bool = False) -> tuple[str, dict]:
|
|
"""Like tokenmax() but also returns substitution statistics.
|
|
|
|
Args:
|
|
force_cjk: If True, always apply when CJK substitutions were made,
|
|
even if total token count increased. Prioritizes CJK adoption
|
|
over token savings.
|
|
|
|
Returns:
|
|
(processed_text, stats_dict) where stats_dict contains:
|
|
- original_tokens: token count before processing
|
|
- result_tokens: token count after processing
|
|
- saved: tokens saved (negative = token increase; check forced_cjk)
|
|
- applied: whether the processed version was used
|
|
- forced_cjk: True when force_cjk override caused acceptance despite no savings
|
|
- filler_drops, cjk, symbol, abbrev: substitution counts by category
|
|
- total_subs: total substitutions applied
|
|
"""
|
|
if not text or not text.strip():
|
|
return text, {'original_tokens': 0, 'result_tokens': 0, 'saved': 0,
|
|
'applied': False, 'forced_cjk': False, 'total_subs': 0}
|
|
|
|
original_tokens = _ntok(text, tokenizer)
|
|
result, stats = _apply_subs(text)
|
|
result_tokens = _ntok(result, tokenizer)
|
|
saved = original_tokens - result_tokens
|
|
|
|
stats['original_tokens'] = original_tokens
|
|
stats['result_tokens'] = result_tokens
|
|
stats['saved'] = saved
|
|
stats['forced_cjk'] = False
|
|
|
|
if saved > 0:
|
|
stats['applied'] = True
|
|
return result, stats
|
|
if force_cjk and stats.get('cjk', 0) > 0:
|
|
stats['applied'] = True
|
|
stats['forced_cjk'] = True
|
|
return result, stats
|
|
stats['applied'] = False
|
|
return text, stats
|
|
|
|
|
|
# ── CLI: batch process a JSONL file ────────────────────────────────────
|
|
if __name__ == '__main__':
|
|
import json, sys, argparse
|
|
from transformers import AutoTokenizer
|
|
|
|
parser = argparse.ArgumentParser(description='Token-max post-processor for think blocks')
|
|
parser.add_argument('--input', required=True, help='Input JSONL (messages format)')
|
|
parser.add_argument('--output', help='Output JSONL (default: dry run, stats only)')
|
|
parser.add_argument('--tokenizer', default='ZelligeAI/tessera-compressor',
|
|
help='HF repo id or local path of the tokenizer to count savings under')
|
|
parser.add_argument('--force-cjk', action='store_true',
|
|
help='Force CJK substitutions even if total tokens increase. '
|
|
'Prioritizes CJK adoption over token savings.')
|
|
args = parser.parse_args()
|
|
|
|
tok = AutoTokenizer.from_pretrained(args.tokenizer, trust_remote_code=True)
|
|
|
|
total_before = total_after = applied = skipped = forced = 0
|
|
|
|
out_lines = []
|
|
with open(args.input) as f:
|
|
for line in f:
|
|
line = line.strip()
|
|
if not line:
|
|
continue
|
|
rec = json.loads(line)
|
|
for m in rec.get('messages', rec.get('conversations', [])):
|
|
role = m.get('role', m.get('from', ''))
|
|
if role not in ('assistant', 'gpt'):
|
|
continue
|
|
content_key = 'content' if 'content' in m else 'value'
|
|
c = m.get(content_key, '') or ''
|
|
if '<think>' not in c or '</think>' not in c:
|
|
continue
|
|
|
|
# Extract think content, preserving prefix before <think> and suffix after </think>
|
|
think_start = c.index('<think>') + len('<think>')
|
|
think_end = c.index('</think>')
|
|
prefix = c[:think_start - len('<think>')]
|
|
think = c[think_start:think_end]
|
|
suffix = c[think_end + len('</think>'):]
|
|
|
|
maxed, stats = tokenmax_with_stats(think, tok, force_cjk=args.force_cjk)
|
|
|
|
total_before += stats['original_tokens']
|
|
if stats['applied']:
|
|
applied += 1
|
|
total_after += stats['result_tokens']
|
|
m[content_key] = f'{prefix}<think>{maxed}</think>{suffix}'
|
|
if stats.get('forced_cjk'):
|
|
forced += 1
|
|
else:
|
|
skipped += 1
|
|
total_after += stats['original_tokens']
|
|
|
|
out_lines.append(json.dumps(rec, ensure_ascii=False))
|
|
|
|
if args.output:
|
|
with open(args.output, 'w') as f:
|
|
for line in out_lines:
|
|
f.write(line + '\n')
|
|
|
|
total = applied + skipped
|
|
saved = total_before - total_after
|
|
if total > 0:
|
|
print(f'Processed {total} think blocks')
|
|
print(f' Applied: {applied} ({100*applied/total:.0f}%)')
|
|
if forced:
|
|
print(f' Forced CJK: {forced} (applied despite no token savings)')
|
|
print(f' Skipped (no savings): {skipped}')
|
|
pct = f'{100*saved/total_before:.1f}' if total_before > 0 else '0.0'
|
|
print(f' Tokens: {total_before} → {total_after} = {saved:+d} ({pct}%)')
|
|
else:
|
|
print(f'No think blocks found in {args.input}')
|