Skills › Scientific Agent Skills › gtars › tokenizer_manifest.py
tokenizer_manifest.py
238 lines · skills/gtars/scripts/tokenizer_manifest.py at 5f2bc03
1#!/usr/bin/env python3
2"""Check a local Gtars tokenizer manifest against its exact universe."""
4from __future__ import annotations
6import argparse
7import re
8import sys
10from _common import (
11 HARD_MAX_BYTES,
12 HARD_MAX_RECORDS,
13 SafetyError,
14 add_path_mode_argument,
15 display_path,
16 fail_json,
17 inspect_bed,
18 int_type,
19 load_chrom_sizes,
20 load_json,
21 local_path,
22 print_json,
23 sha256_file,
24)
27TOOL = "gtars-tokenizer-manifest"
28DEFAULT_GTARS_PYTHON_VERSION = "0.10.0"
29SPECIAL_TOKEN_NAMES = {"unk", "pad", "mask", "cls", "bos", "eos", "sep"}
30_SHA256 = re.compile(r"^[0-9a-f]{64}$")
33def build_parser() -> argparse.ArgumentParser:
34 parser = argparse.ArgumentParser(
35 description=(
36 "Compare a local JSON tokenizer manifest with a local universe BED. "
37 "No tokenizer is imported and no model repository is contacted."
38 )
39 )
40 parser.add_argument("--manifest", required=True)
41 parser.add_argument("--universe", required=True)
42 parser.add_argument("--assembly", required=True)
43 parser.add_argument("--chrom-sizes")
44 parser.add_argument(
45 "--expected-python-version",
46 default=DEFAULT_GTARS_PYTHON_VERSION,
47 )
48 parser.add_argument(
49 "--max-bytes",
50 type=int_type(minimum=1, maximum=HARD_MAX_BYTES, label="max-bytes"),
51 default=1024**3,
52 )
53 parser.add_argument(
54 "--max-records",
55 type=int_type(
56 minimum=1,
57 maximum=HARD_MAX_RECORDS,
58 label="max-records",
59 ),
60 default=5_000_000,
61 )
62 add_path_mode_argument(parser)
63 return parser
66def _mapping(value, label: str) -> dict:
67 if not isinstance(value, dict):
68 raise SafetyError(f"{label} must be a JSON object")
69 return value
72def check(args: argparse.Namespace) -> tuple[dict, int]:
73 assembly = args.assembly.strip()
74 if not assembly or len(assembly) > 200:
75 raise SafetyError("assembly must contain 1-200 characters")
76 manifest_path = local_path(args.manifest, kind="file")
77 universe_path = local_path(args.universe, kind="file")
78 manifest = _mapping(
79 load_json(manifest_path, max_bytes=min(args.max_bytes, 16 * 1024**2)),
80 "manifest",
81 )
83 chrom_sizes = None
84 chrom_order = None
85 chrom_digest = None
86 chrom_report = None
87 if args.chrom_sizes:
88 chrom_path = local_path(args.chrom_sizes, kind="file")
89 chrom_sizes, chrom_order = load_chrom_sizes(
90 chrom_path,
91 max_bytes=min(args.max_bytes, 128 * 1024**2),
92 max_records=min(args.max_records, 1_000_000),
93 )
94 chrom_digest, chrom_bytes = sha256_file(
95 chrom_path,
96 max_bytes=min(args.max_bytes, 128 * 1024**2),
97 )
98 chrom_report = {
99 "path": display_path(chrom_path, 3, args.path_mode),
100 "sha256": chrom_digest,
101 "size_bytes": chrom_bytes,
102 }
104 universe = inspect_bed(
105 universe_path,
106 max_bytes=args.max_bytes,
107 max_records=args.max_records,
108 chrom_sizes=chrom_sizes,
109 chrom_order=chrom_order,
110 )
111 errors: dict[str, int] = {
112 f"universe:{code}": count for code, count in universe["errors"].items()
113 }
114 warnings: dict[str, int] = {}
115 if universe["records"] == 0:
116 errors["universe:no_data_records"] = 1
117 if universe["duplicate_intervals"]:
118 errors["universe:duplicate_intervals"] = universe["duplicate_intervals"]
119 if universe["out_of_order_records"]:
120 warnings["universe:order_is_part_of_token_ids"] = universe[
121 "out_of_order_records"
122 ]
124 manifest_universe = _mapping(manifest.get("universe"), "manifest.universe")
125 tokenizer = _mapping(manifest.get("tokenizer"), "manifest.tokenizer")
126 expected_sha = manifest_universe.get("sha256")
127 if not isinstance(expected_sha, str) or not _SHA256.fullmatch(expected_sha):
128 errors["manifest:invalid_universe_sha256"] = 1
129 elif expected_sha != universe["sha256"]:
130 errors["manifest:universe_sha256_mismatch"] = 1
132 expected_records = manifest_universe.get("records")
133 if type(expected_records) is not int or expected_records < 1:
134 errors["manifest:invalid_universe_record_count"] = 1
135 elif expected_records != universe["records"]:
136 errors["manifest:universe_record_count_mismatch"] = 1
138 if manifest.get("schema_version") != "1.0":
139 errors["manifest:unsupported_schema_version"] = 1
140 if manifest.get("assembly") != assembly:
141 errors["manifest:assembly_mismatch"] = 1
142 if manifest.get("coordinate_system") != "0-based-half-open":
143 errors["manifest:coordinate_system_mismatch"] = 1
144 if manifest.get("gtars_python_version") != args.expected_python_version:
145 errors["manifest:gtars_python_version_mismatch"] = 1
146 if chrom_digest is not None:
147 if manifest_universe.get("chrom_sizes_sha256") != chrom_digest:
148 errors["manifest:chrom_sizes_sha256_mismatch"] = 1
150 backend = tokenizer.get("backend")
151 if backend not in {"bits", "ailist"}:
152 errors["manifest:invalid_backend"] = 1
153 vocab_size = tokenizer.get("vocab_size")
154 expected_vocab = universe["records"] + len(SPECIAL_TOKEN_NAMES)
155 if type(vocab_size) is not int or vocab_size < 1:
156 errors["manifest:invalid_vocab_size"] = 1
157 elif vocab_size != expected_vocab:
158 errors["manifest:vocab_size_mismatch"] = 1
160 token_ids = tokenizer.get("special_token_ids")
161 if not isinstance(token_ids, dict):
162 errors["manifest:missing_special_token_ids"] = 1
163 token_ids = {}
164 else:
165 names = set(token_ids)
166 if names != SPECIAL_TOKEN_NAMES:
167 errors["manifest:special_token_name_set_mismatch"] = 1
168 values = list(token_ids.values())
169 if any(type(value) is not int or value < 0 for value in values):
170 errors["manifest:invalid_special_token_id"] = 1
171 elif len(values) != len(set(values)):
172 errors["manifest:duplicate_special_token_ids"] = 1
173 elif type(vocab_size) is int and any(value >= vocab_size for value in values):
174 errors["manifest:special_token_id_out_of_range"] = 1
176 manifest_digest, manifest_bytes = sha256_file(
177 manifest_path,
178 max_bytes=min(args.max_bytes, 16 * 1024**2),
179 )
180 report = {
181 "ok": not errors,
182 "tool": TOOL,
183 "contract": {
184 "assembly": assembly,
185 "coordinate_system": "0-based-half-open",
186 "expected_gtars_python_version": args.expected_python_version,
187 "network_used": False,
188 "packages_imported": False,
189 "model_deserialized": False,
190 "files_written": False,
191 "symlinks_allowed": False,
192 },
193 "manifest": {
194 "path": display_path(manifest_path, 1, args.path_mode),
195 "sha256": manifest_digest,
196 "size_bytes": manifest_bytes,
197 "schema_version": manifest.get("schema_version"),
198 },
199 "universe": {
200 "path": display_path(universe_path, 2, args.path_mode),
201 "sha256": universe["sha256"],
202 "size_bytes": universe["size_bytes"],
203 "records": universe["records"],
204 "duplicate_intervals": universe["duplicate_intervals"],
205 "out_of_order_records": universe["out_of_order_records"],
206 },
207 "chromosome_sizes": chrom_report,
208 "compatibility": {
209 "backend": backend,
210 "vocab_size": vocab_size,
211 "expected_vocab_size": expected_vocab,
212 "special_token_count": len(token_ids),
213 "universe_bytes_and_order_match": expected_sha == universe["sha256"],
214 },
215 "errors": dict(sorted(errors.items())),
216 "warnings": dict(sorted(warnings.items())),
217 "next_checks": [
218 "freeze patient and replicate splits before fitting the universe",
219 "instantiate Tokenizer.from_bed only after this local manifest passes",
220 "verify a small fixed set of token IDs with gtars==0.10.0",
221 "do not use from_pretrained until an immutable snapshot is approved locally",
222 ],
223 }
224 return report, 0 if not errors else 2
227def main() -> int:
228 args = build_parser().parse_args()
229 try:
230 report, status = check(args)
231 print_json(report)
232 return status
233 except (OSError, SafetyError, UnicodeError) as exc:
234 return fail_json(TOOL, exc)
237if __name__ == "__main__":
238 sys.exit(main())