needle / scripts /token_accounting_study.py
Leon4gr45's picture
Upload folder using huggingface_hub (part 2)
25a2c1c verified
Raw History Blame Contribute Delete
5.47 kB
"""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()