File size: 4,246 Bytes
36b8f8e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 | """Stable delimiter-based encoding used by the independent model."""
from __future__ import annotations
import re
QWEN_SPECIAL = ["<|fim_prefix|>", "<|fim_middle|>", "<|box_start|>", "<|box_end|>", "<|fim_suffix|>"]
# Reuse five reserved Gemma 4 control tokens instead of resizing its 262K
# vocabulary and training new embeddings. They are only structural markers in
# this prefill-only model; user text containing control-token syntax is escaped.
GEMMA4_SPECIAL = ["<|turn>", "<|channel>", "<|tool>", "<tool|>", "<turn|>"]
SPECIAL = QWEN_SPECIAL
SPECIAL_RE = re.compile(r"<\|([A-Za-z0-9_]+)\|>")
def user_tokens(tokenizer, text: str) -> list[int]:
safe = SPECIAL_RE.sub(r"<¦\1¦>", str(text)).replace("<turn|>", "<turn¦>").replace("<tool|>", "<tool¦>")
return tokenizer(safe, add_special_tokens=False).input_ids
def delimiter_tokens(tokenizer) -> list[str]:
"""Select five distinct, existing delimiter tokens for the backbone."""
unknown = tokenizer.unk_token_id
for profile in (QWEN_SPECIAL, GEMMA4_SPECIAL):
ids = [tokenizer.convert_tokens_to_ids(token) for token in profile]
if None not in ids and unknown not in ids and len(set(ids)) == len(ids):
return profile
raise ValueError("tokenizer has no supported five-token delimiter profile")
def _text_tokens(tokenizer, text: str) -> list[int]:
return tokenizer(text, add_special_tokens=False).input_ids
def encode_question(tokenizer, record: dict, question: dict, max_state: int = 384, max_branch: int = 1024,
profile: str = "legacy") -> dict:
if profile == "gemma_chat":
state = user_tokens(tokenizer, record["state"])
if len(state) + 1 > max_state:
raise ValueError(f"state exceeds the configured {max_state}-token limit: {len(state) + 1}")
bos = [tokenizer.bos_token_id] if tokenizer.bos_token_id is not None else []
prefix = bos + _text_tokens(tokenizer, "<|turn>user\nContext:\n") + state
instr = _text_tokens(tokenizer, "\nQuestion:\n") + user_tokens(tokenizer, question["instr"]) + _text_tokens(tokenizer, "\nOptions:\n")
spans = [_text_tokens(tokenizer, "\nOption: ") + user_tokens(tokenizer, option) + _text_tokens(tokenizer, "\n")
for option in question["options"]]
suffix = _text_tokens(tokenizer, "<turn|>\n<|turn>model\nDecision:")
branch = instr + [token for span in spans for token in span] + suffix
if len(branch) + len(prefix) > max_branch:
raise ValueError("question branch exceeds the configured context")
opt_ends, cursor = [], len(instr)
for span in spans:
cursor += len(span)
opt_ends.append(cursor - 2 if len(span) > 1 else cursor - 1)
return {"ids": prefix + branch, "state_len": len(prefix), "decide": len(prefix) + len(branch) - 1,
"opts": [len(prefix) + index for index in opt_ends], "label": question.get("label", 0),
"qtype": question.get("qtype", "choice"), "keys": question.get("keys", [])}
if profile != "legacy":
raise ValueError(f"unsupported encoding profile: {profile}")
state = user_tokens(tokenizer, record["state"])
if len(state) + 1 > max_state:
raise ValueError(f"state exceeds the configured {max_state}-token limit: {len(state) + 1}")
state_id, q_id, opt_id, end_id, decide_id = [tokenizer.convert_tokens_to_ids(x) for x in delimiter_tokens(tokenizer)]
ids = [state_id] + state
instr = [q_id] + user_tokens(tokenizer, question["instr"])
spans = [[opt_id] + user_tokens(tokenizer, option) + [end_id] for option in question["options"]]
branch = instr + [t for span in spans for t in span] + [decide_id]
if len(branch) + len(ids) > max_branch:
raise ValueError("question branch exceeds the configured context")
opt_ends, cursor = [], len(instr)
for span in spans:
cursor += len(span)
opt_ends.append(cursor - 1)
return {"ids": ids + branch, "state_len": len(ids), "decide": len(ids) + len(branch) - 1,
"opts": [len(ids) + i for i in opt_ends], "label": question.get("label", 0),
"qtype": question.get("qtype", "choice"), "keys": question.get("keys", [])}
|