scripts/token_accounting_study.py

"""Offline reference-tokenizer investigation; never imported by the gateway.

Uses reviewed declarations from a checksum-pinned official Apache-2.0 SDK wheel,
without installing it (upstream license: Apache-2.0, retained in that wheel).
The existing SDK worker's own ready response is observed;
no new native symbol/signature is called. Reference counts are NOT native counts.
"""
import argparse
import ast
import hashlib
import heapq
import json
import random
import time
import zipfile
from pathlib import Path

WHEEL_SHA256 = '8b77639b0676588523bca82f1b63f98775379b35d2c252feaf1b483eb4c15473'


def reference(wheel):
    raw = Path(wheel).read_bytes()
    if hashlib.sha256(raw).hexdigest() != WHEEL_SHA256:
        raise ValueError('Official wheel checksum mismatch')
    with zipfile.ZipFile(wheel) as z:
        source = z.read('needle/model/export.py').decode()
    names = {'TAG', 'RAW', '_HDR_FMT', '_REC_FMT', 'REC_SIZE', '_TK_HDR', '_TK_REC',
             '_SP_META_SPACE', 'TK_NORMAL', 'TK_UNKNOWN', 'TK_CONTROL', 'TK_USER_DEFINED',
             'TK_BYTE', 'read_tokenizer_blob', 'parse_tokenizer_blob', 'RefTokenizer'}
    nodes = []
    for node in ast.parse(source).body:
        if isinstance(node, (ast.FunctionDef, ast.ClassDef)) and node.name in names:
            nodes.append(node)
        elif isinstance(node, ast.Assign) and any(isinstance(n, ast.Name) and n.id in names for t in node.targets for n in ast.walk(t)):
            nodes.append(node)
    ns = {'struct': __import__('struct')}
    exec(compile(ast.Module(body=nodes, type_ignores=[]), '<reviewed-official-reference>', 'exec'), ns)
    return ns


def fast_bpe(tokenizer, segment):
    """Same highest-score/leftmost merge rule as official quadratic reference."""
    symbols = list(segment)
    length = len(symbols)
    prev = [i-1 for i in range(length)]
    nxt = [i+1 for i in range(length)]
    alive = bytearray([1])*length
    versions = [0]*length
    heap = []
    def offer(i):
        if i < 0 or i >= length or not alive[i]: return
        j = nxt[i]
        if j >= length: return
        idx = tokenizer.p2id.get(symbols[i]+symbols[j])
        if idx is not None:
            heapq.heappush(heap, (-tokenizer.scores[idx], i, j, versions[i], versions[j]))
    for i in range(length-1): offer(i)
    while heap:
        _, i, j, vi, vj = heapq.heappop(heap)
        if not alive[i] or not alive[j] or nxt[i] != j or versions[i] != vi or versions[j] != vj: continue
        symbols[i] += symbols[j]; versions[i] += 1; alive[j] = 0
        nxt[i] = nxt[j]
        if nxt[j] < length: prev[nxt[j]] = i
        offer(prev[i]); offer(i)
    ids = []
    for i, symbol in enumerate(symbols):
        if not alive[i]: continue
        idx = tokenizer.p2id.get(symbol)
        if idx is not None: ids.append(idx)
        elif tokenizer.byte_fallback: ids.extend(tokenizer.byte_id[b] for b in symbol.encode())
        else: ids.append(tokenizer.unk_id)
    return ids


def native_prefix(library, weights, system, tools):
    from needle._worker import FineTuneWorker
    class Observer(FineTuneWorker):
        def _receive(self, timeout=None):
            message = super()._receive(timeout)
            if message.get('status') == 'ready': self.observed_prefix = message['prefix_tokens']
            return message
    worker = Observer(library, weights, system, json.dumps(tools), None, 65536, generation=3)
    try: return worker.observed_prefix
    finally: worker.close()


def main():
    p = argparse.ArgumentParser(); p.add_argument('--wheel', required=True); p.add_argument('--assets', required=True)
    args = p.parse_args(); ns = reference(args.wheel)
    blob = ns['read_tokenizer_blob'](str(Path(args.assets)/'needle3.cact'))
    meta = ns['parse_tokenizer_blob'](blob)
    ref = ns['RefTokenizer'](meta); fast = ns['RefTokenizer'](meta)
    fast._bpe = lambda text: fast_bpe(fast, text)
    corpus = ['', 'Read /repo/auth.py', '中文 日本語 한국어', 'مرحبا بالعالم', 'café e\u0301 😀', '\t\n\r\0', '<|im_start|>system\n<tools>{}</tools>', json.dumps({'type':'object','properties':{'path':{'type':'string'}}})]
    rng=random.Random(91231)
    corpus += [''.join(rng.choice('abcXYZ /._{}[]:\n\té中😀') for _ in range(rng.randrange(1,240))) for _ in range(100)]
    assert all(ref.encode(s)==fast.encode(s) for s in corpus)
    tools=[{'name':'read_file','description':'Read file','parameters':{'type':'object','properties':{'path':{'type':'string'}},'required':['path']}}]
    samples=[]
    for system in corpus[:8]:
        prefix=native_prefix(str(Path(args.assets)/'libneedle.so'),str(Path(args.assets)/'needle3.cact'),system,tools)
        samples.append({'system_utf8_bytes':len(system.encode()),'reference_system_tokens':len(ref.encode(system)), 'reference_tools_tokens':len(ref.encode(json.dumps(tools))),'native_static_prefix_tokens':prefix})
    print(json.dumps({'official_wheel_sha256':WHEEL_SHA256,'vocabulary_pieces':len(meta['pieces']),'byte_fallback':meta['byte_fallback'],'dummy_prefix':meta['add_dummy_prefix'],'reference_conformance_cases':len(corpus),'native_prefix_observations':samples,'verified_exact_runtime_accounting':False,'limitation':'Native prefix is observable through the existing SDK worker, but query token IDs/counts and actual prefix-template bytes are not exposed. Reference conformance is not native conformance. Keep production byte proxy and reserve.'},indent=2))

if __name__=='__main__': main()