#!/usr/bin/env python3 """Audit Qwen235B trace token lengths and source prefix-hash identities.""" from __future__ import annotations import argparse import hashlib import json import math import platform import socket import time from pathlib import Path from typing import Any def sha256_file(path: Path) -> str: digest = hashlib.sha256() with path.open("rb") as source: for chunk in iter(lambda: source.read(1 << 20), b""): digest.update(chunk) return digest.hexdigest() def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser() parser.add_argument("--trace", type=Path, required=True) parser.add_argument("--model", type=Path, required=True) parser.add_argument("--output", type=Path, required=True) parser.add_argument("--max-model-len", type=int, default=40960) parser.add_argument("--source-block-size", type=int, default=64) parser.add_argument("--batch-size", type=int, default=16) return parser.parse_args() def main() -> None: args = parse_args() if args.max_model_len <= 0 or args.source_block_size <= 0 or args.batch_size <= 0: raise ValueError("length and batch-size arguments must be positive") import transformers from transformers import AutoTokenizer rows = [json.loads(line) for line in args.trace.open() if line.strip()] context_exceeded = [ row for row in rows if int(row["input_length"]) + int(row["output_length"]) > args.max_model_len ] zero_output = [row for row in rows if int(row["output_length"]) == 0] eligible = [ row for row in rows if int(row["input_length"]) + int(row["output_length"]) <= args.max_model_len and int(row["output_length"]) > 0 ] tokenizer = AutoTokenizer.from_pretrained(args.model, trust_remote_code=True) length_mismatch_count = 0 hash_count_mismatch_count = 0 hash_to_key: dict[str, bytes] = {} key_to_hash: dict[bytes, str] = {} hash_to_key_conflict_count = 0 key_to_hash_conflict_count = 0 total_tokens = 0 full_blocks = 0 partial_blocks = 0 length_digest = hashlib.sha256() token_digest = hashlib.sha256() started = time.time() for start in range(0, len(eligible), args.batch_size): batch = eligible[start : start + args.batch_size] encoded = tokenizer( [row["prompt"] for row in batch], add_special_tokens=False, padding=False, truncation=False, )["input_ids"] for offset, (row, token_ids) in enumerate(zip(batch, encoded, strict=True)): row_index = start + offset actual_length = len(token_ids) expected_length = int(row["input_length"]) total_tokens += actual_length length_digest.update(f"{row_index}:{actual_length}\n".encode()) if actual_length != expected_length: length_mismatch_count += 1 source_hashes = row["hash_ids"] expected_hashes = math.ceil(actual_length / args.source_block_size) if len(source_hashes) != expected_hashes: hash_count_mismatch_count += 1 continue request_token_digest = hashlib.sha256() parent = b"ROOT" for block_index, source_hash in enumerate(source_hashes): begin = block_index * args.source_block_size chunk = token_ids[begin : begin + args.source_block_size] token_payload = b"".join( int(token_id).to_bytes(4, "little", signed=False) for token_id in chunk ) request_token_digest.update(token_payload) chunk_digest = hashlib.blake2b(token_payload, digest_size=16).digest() key_digest = hashlib.blake2b( parent + b"\0" + chunk_digest, digest_size=16 ).digest() if len(chunk) == args.source_block_size: full_blocks += 1 else: partial_blocks += 1 hash_id = str(source_hash) previous_key = hash_to_key.setdefault(hash_id, key_digest) if previous_key != key_digest: hash_to_key_conflict_count += 1 previous_hash = key_to_hash.setdefault(key_digest, hash_id) if previous_hash != hash_id: key_to_hash_conflict_count += 1 parent = hash_id.encode() token_digest.update(row_index.to_bytes(4, "little")) token_digest.update(request_token_digest.digest()) payload: dict[str, Any] = { "schema": "qwen235b-trace-contract-audit-v1", "status": "pass_offline_source_contract" if not any( ( length_mismatch_count, hash_count_mismatch_count, hash_to_key_conflict_count, key_to_hash_conflict_count, ) ) else "fail", "execution": { "host": socket.gethostname(), "device": "cpu_only", "elapsed_seconds": round(time.time() - started, 3), "python_version": platform.python_version(), "tokenizer_class": type(tokenizer).__name__, "transformers_version": transformers.__version__, "model_path": str(args.model.resolve()), }, "trace": { "path": str(args.trace.resolve()), "sha256": sha256_file(args.trace), "source_request_count": len(rows), "context_exceeded_count": len(context_exceeded), "zero_output_count": len(zero_output), "exclusion_overlap_count": sum(row in zero_output for row in context_exceeded), "eligible_request_count": len(eligible), }, "tokenization": { "total_token_count": total_tokens, "input_length_mismatch_count": length_mismatch_count, "length_order_sha256": length_digest.hexdigest(), "per_request_token_digest_sha256": token_digest.hexdigest(), }, "source_hash_contract": { "source_block_size_tokens": args.source_block_size, "hash_count_mismatch_count": hash_count_mismatch_count, "full_block_count": full_blocks, "partial_block_count": partial_blocks, "unique_hash_id_count": len(hash_to_key), "unique_parent_chunk_key_count": len(key_to_hash), "hash_id_to_parent_chunk_conflict_count": hash_to_key_conflict_count, "parent_chunk_to_hash_id_conflict_count": key_to_hash_conflict_count, "key_definition": ( "(parent source hash id, BLAKE2b-128 of the tokenizer token-id chunk)" ), }, } args.output.parent.mkdir(parents=True, exist_ok=True) args.output.write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n") if payload["status"] != "pass_offline_source_contract": raise SystemExit(1) if __name__ == "__main__": main()