File size: 2,528 Bytes
31596cd | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 | #!/usr/bin/env python3
"""Compare GGUF token IDs with the pinned Hugging Face tokenizer outputs."""
from __future__ import annotations
import argparse
import ast
import json
import subprocess
from pathlib import Path
def main() -> None:
parser = argparse.ArgumentParser()
parser.add_argument("--llama-tokenize", type=Path, required=True)
parser.add_argument("--model", type=Path, required=True)
parser.add_argument("--hf-reference", type=Path, required=True)
parser.add_argument("--output", type=Path, required=True)
args = parser.parse_args()
reference = json.loads(args.hf_reference.read_text(encoding="utf-8"))
cases = reference.get("tokenizer_cases")
if not isinstance(cases, list) or len(cases) != 6:
raise SystemExit("HF reference must contain exactly six tokenizer cases")
if any(
not isinstance(case, dict)
or not isinstance(case.get("text"), str)
or not case["text"]
or not isinstance(case.get("token_ids"), list)
or not case["token_ids"]
for case in cases
):
raise SystemExit("HF reference contains an empty or malformed tokenizer case")
comparisons = []
for case in cases:
completed = subprocess.run(
[
str(args.llama_tokenize),
"--model",
str(args.model),
"--prompt",
case["text"],
"--ids",
"--no-bos",
"--offline",
"--log-disable",
],
check=True,
text=True,
capture_output=True,
)
gguf_ids = ast.literal_eval(completed.stdout.strip())
match = gguf_ids == case["token_ids"]
comparisons.append(
{
"text": case["text"],
"hf_token_ids": case["token_ids"],
"gguf_token_ids": gguf_ids,
"exact_match": match,
}
)
result = {
"cases": comparisons,
"case_count": len(comparisons),
"all_exact": all(case["exact_match"] for case in comparisons),
}
args.output.parent.mkdir(parents=True, exist_ok=True)
args.output.write_text(
json.dumps(result, ensure_ascii=False, indent=2, sort_keys=True) + "\n",
encoding="utf-8",
)
print(json.dumps(result, ensure_ascii=False, sort_keys=True))
if not result["all_exact"]:
raise SystemExit("tokenizer mismatch")
if __name__ == "__main__":
main()
|