184 lines
7.0 KiB
Python
184 lines
7.0 KiB
Python
#!/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()
|