| """Standalone tokenizer for Needle 3 ONNX models. |
| |
| Decodes the SentencePiece BPE vocabulary embedded in tokenizer.json |
| without external sentencepiece or c++ binary dependencies. |
| """ |
|
|
| import json |
| import struct |
| from typing import Any, Dict, List, Optional, Tuple, Union |
|
|
| _SP_META_SPACE = "\u2581" |
| TK_NORMAL, TK_UNKNOWN, TK_CONTROL, TK_USER_DEFINED, TK_BYTE = 0, 1, 2, 3, 4 |
| _TK_HDR = "<IIIIIBBH" |
| _TK_REC = "<fBH" |
|
|
| |
| PAD_ID = 0 |
| EOS_ID = 1 |
| BOS_ID = 2 |
| UNK_ID = 3 |
| IM_START_ID = 4 |
| IM_END_ID = 5 |
| THINK_START_ID = 6 |
| THINK_END_ID = 7 |
| TOOLS_START_ID = 8 |
| TOOLS_END_ID = 9 |
| TOOL_CALL_START_ID = 10 |
| TOOL_CALL_END_ID = 11 |
|
|
| IM_START = "<|im_start|>" |
| IM_END = "<|im_end|>" |
| THINK_START = "<think>" |
| THINK_END = "</think>" |
| TOOLS_START = "<tools>" |
| TOOLS_END = "</tools>" |
| TOOL_CALL_START = "<tool_call>" |
| TOOL_CALL_END = "</tool_call>" |
|
|
|
|
| def parse_tokenizer_blob(blob: bytes) -> Dict[str, Any]: |
| """Parse raw SentencePiece vocabulary blob from needle3.cact if provided.""" |
| off = struct.calcsize(_TK_HDR) |
| n, pad, eos, bos, unk, add_dummy, byte_fb, _ = struct.unpack_from(_TK_HDR, blob, 0) |
| rec = struct.calcsize(_TK_REC) |
| pieces, scores, types = [], [], [] |
| for _ in range(n): |
| score, t, ln = struct.unpack_from(_TK_REC, blob, off) |
| off += rec |
| pieces.append(blob[off : off + ln].decode("utf-8")) |
| scores.append(score) |
| types.append(t) |
| off += ln |
| return { |
| "pieces": pieces, |
| "scores": scores, |
| "types": types, |
| "pad_id": pad, |
| "eos_id": eos, |
| "bos_id": bos, |
| "unk_id": unk, |
| "add_dummy_prefix": bool(add_dummy), |
| "byte_fallback": bool(byte_fb), |
| } |
|
|
|
|
| class NeedleTokenizer: |
| """Tokenizer for Needle 3 models.""" |
|
|
| def __init__(self, meta: Dict[str, Any]): |
| self.pieces = meta["pieces"] |
| self.scores = meta["scores"] |
| self.types = meta["types"] |
| self.add_dummy = meta["add_dummy_prefix"] |
| self.byte_fallback = meta["byte_fallback"] |
| self.unk_id = meta["unk_id"] |
| self.bos_id = meta.get("bos_id", BOS_ID) |
| self.eos_id = meta.get("eos_id", EOS_ID) |
| self.pad_id = meta.get("pad_id", PAD_ID) |
|
|
| self.p2id = {p: i for i, p in enumerate(self.pieces)} |
| self.byte_id = { |
| int(p[3:5], 16): i |
| for i, (p, t) in enumerate(zip(self.pieces, self.types)) |
| if t == TK_BYTE |
| } |
| self.markers = sorted( |
| (p for p, t in zip(self.pieces, self.types) if t == TK_USER_DEFINED), |
| key=len, |
| reverse=True, |
| ) |
|
|
| @classmethod |
| def from_json(cls, json_path: str) -> "NeedleTokenizer": |
| """Load tokenizer vocabulary from JSON file.""" |
| with open(json_path, "r", encoding="utf-8") as f: |
| meta = json.load(f) |
| return cls(meta) |
|
|
| @classmethod |
| def from_cact_blob(cls, blob: bytes) -> "NeedleTokenizer": |
| meta = parse_tokenizer_blob(blob) |
| return cls(meta) |
|
|
| def _bpe(self, seg: str) -> List[int]: |
| syms = list(seg) |
| while len(syms) > 1: |
| best_score, best_j = None, -1 |
| for j in range(len(syms) - 1): |
| idx = self.p2id.get(syms[j] + syms[j + 1]) |
| if idx is not None and (best_score is None or self.scores[idx] > best_score): |
| best_score, best_j = self.scores[idx], j |
| if best_j < 0: |
| break |
| syms[best_j : best_j + 2] = [syms[best_j] + syms[best_j + 1]] |
| ids = [] |
| for s in syms: |
| idx = self.p2id.get(s) |
| if idx is not None: |
| ids.append(idx) |
| elif self.byte_fallback: |
| ids.extend(self.byte_id[b] for b in s.encode("utf-8")) |
| else: |
| ids.append(self.unk_id) |
| return ids |
|
|
| def encode(self, text: str) -> List[int]: |
| if not text: |
| return [] |
| esc = text.replace(" ", _SP_META_SPACE) |
| if self.add_dummy: |
| esc = _SP_META_SPACE + esc |
| ids, buf, i, n = [], [], 0, len(esc) |
| while i < n: |
| marker = next((m for m in self.markers if esc.startswith(m, i)), None) |
| if marker is not None: |
| ids += self._bpe("".join(buf)) |
| buf = [] |
| ids.append(self.p2id[marker]) |
| i += len(marker) |
| else: |
| buf.append(esc[i]) |
| i += 1 |
| ids += self._bpe("".join(buf)) |
| return ids |
|
|
| def decode(self, ids: List[int]) -> str: |
| buf = bytearray() |
| for i in ids: |
| t = self.types[i] |
| if t == TK_BYTE: |
| buf.append(int(self.pieces[i][3:5], 16)) |
| elif t in (TK_CONTROL, TK_UNKNOWN): |
| continue |
| else: |
| buf += self.pieces[i].encode("utf-8") |
| text = buf.decode("utf-8", "replace").replace(_SP_META_SPACE, " ") |
| if self.add_dummy and text.startswith(" "): |
| text = text[1:] |
| return text |
|
|
| def build_prompt( |
| self, |
| query: str, |
| tools: Optional[List[Dict[str, Any]]] = None, |
| system: Optional[str] = None, |
| ) -> Tuple[str, List[int]]: |
| """Construct prompt and tokenize with BOS token.""" |
| prefix = "" |
| if system: |
| prefix = f"{IM_START}system\n{system}{IM_END}\n" |
| tools_str = "" |
| if tools is not None: |
| tools_json = json.dumps(tools, separators=(",", ":"), ensure_ascii=False) |
| tools_str = f"{TOOLS_START}{tools_json}{TOOLS_END}\n" |
| prompt = f"{prefix}{IM_START}user\n{tools_str}{query}{IM_END}\n{IM_START}assistant\n" |
| ids = [self.bos_id] + self.encode(prompt) |
| return prompt, ids |
|
|
| def parse_tool_calls(self, output_ids: List[int]) -> Tuple[str, List[Dict[str, Any]]]: |
| """Extract reasoning and tool call JSON from decoded tokens.""" |
| pieces = [self.pieces[i] for i in output_ids] |
| raw = "".join(pieces).replace(_SP_META_SPACE, " ") |
|
|
| reasoning = "" |
| if THINK_START in raw and THINK_END in raw: |
| reasoning = raw.split(THINK_START, 1)[1].split(THINK_END, 1)[0].strip() |
|
|
| calls = [] |
| if TOOL_CALL_START in raw: |
| body = raw.split(TOOL_CALL_START, 1)[1].split(TOOL_CALL_END, 1)[0].strip() |
| if body.startswith("[]name") or body.startswith('[]"name'): |
| body = '[{"' + body[2:] |
| try: |
| parsed = json.loads(body) |
| if isinstance(parsed, list): |
| calls = parsed |
| elif isinstance(parsed, dict): |
| calls = [parsed] |
| except Exception: |
| calls = [] |
| return reasoning, calls |
|
|
