Spaces:
Running
Running
Download scripts/token_accounting_study.py from Leon4gr45/needle: direct link, hf CLI and curl.
- Browser
- Download file 5.47 kB
-
https://huggingface.co/spaces/Leon4gr45/needle/resolve/main/scripts/token_accounting_study.py
- Command line
-
hf download hf://spaces/Leon4gr45/needle/scripts/token_accounting_study.py
-
curl -L -o token_accounting_study.py https://huggingface.co/spaces/Leon4gr45/needle/resolve/main/scripts/token_accounting_study.py
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() | |