Skills › Scientific Agent Skills › geniml › tokenizer_compatibility.py
tokenizer_compatibility.py
321 lines · skills/geniml/scripts/tokenizer_compatibility.py at 5f2bc03
1#!/usr/bin/env python3
2"""Plan local tokenizer/universe/model compatibility checks without imports."""
4from __future__ import annotations
6import argparse
7import os
8import stat
9import sys
10from collections import Counter
11from pathlib import Path
13from _common import (
14 HARD_MAX_BYTES,
15 HARD_MAX_RECORDS,
16 MAX_COORDINATE,
17 SafetyError,
18 add_path_mode_argument,
19 display_path,
20 fail_json,
21 int_type,
22 iter_text_lines,
23 local_path,
24 print_json,
25 sha256_file,
26 simple_yaml_mapping,
27)
30TOOL = "tokenizer-universe-compatibility"
33def build_parser() -> argparse.ArgumentParser:
34 parser = argparse.ArgumentParser(
35 description=(
36 "Compare a local model bundle and universe using checksums and bounded "
37 "metadata only. No model is deserialized and no package is imported."
38 )
39 )
40 parser.add_argument("--model-dir", required=True, help="Local model bundle directory.")
41 parser.add_argument("--universe", required=True, help="Expected local universe BED.")
42 parser.add_argument("--assembly", required=True, help="Declared assembly/accession.")
43 parser.add_argument("--config-name", default="config.yaml")
44 parser.add_argument("--checkpoint-name", default="checkpoint.pt")
45 parser.add_argument("--bundle-universe-name", default="universe.bed")
46 parser.add_argument(
47 "--token-corpus",
48 help="Optional local Parquet token corpus to fingerprint (not parsed).",
49 )
50 parser.add_argument("--geniml-version", default="0.8.4")
51 parser.add_argument("--gtars-version", default="0.10.0")
52 parser.add_argument(
53 "--special-token-count",
54 type=int_type(minimum=0, maximum=100, label="special-token-count"),
55 default=7,
56 help="Expected Gtars special tokens (0.10.0 default: 7).",
57 )
58 parser.add_argument(
59 "--max-bytes",
60 type=int_type(minimum=1, maximum=HARD_MAX_BYTES, label="max-bytes"),
61 default=2 * 1024**3,
62 )
63 parser.add_argument(
64 "--max-universe-records",
65 type=int_type(
66 minimum=1,
67 maximum=HARD_MAX_RECORDS,
68 label="max-universe-records",
69 ),
70 default=5_000_000,
71 )
72 add_path_mode_argument(parser)
73 return parser
76def _safe_child(directory: Path, name: str) -> Path:
77 if not name or Path(name).name != name or name in {".", ".."}:
78 raise SafetyError("artifact names must be plain filenames")
79 return local_path(str(directory / name), kind="file")
82def _count_universe(
83 path: Path,
84 *,
85 max_bytes: int,
86 max_records: int,
87) -> tuple[int, int, Counter[str]]:
88 records = 0
89 duplicates = 0
90 issues: Counter[str] = Counter()
91 seen: set[tuple[str, int, int]] = set()
92 for _, line in iter_text_lines(
93 path,
94 max_bytes=max_bytes,
95 max_records=max_records,
96 ):
97 if not line or line.startswith("#") or line.startswith("track") or line.startswith("browser"):
98 continue
99 fields = line.split("\t")
100 if len(fields) < 3:
101 issues["fewer_than_three_columns"] += 1
102 continue
103 chrom = fields[0]
104 try:
105 start = int(fields[1])
106 end = int(fields[2])
107 except ValueError:
108 issues["non_integer_coordinate"] += 1
109 continue
110 if (
111 not chrom
112 or start < 0
113 or end <= start
114 or start > MAX_COORDINATE
115 or end > MAX_COORDINATE
116 ):
117 issues["invalid_interval"] += 1
118 continue
119 records += 1
120 key = (chrom, start, end)
121 if key in seen:
122 duplicates += 1
123 else:
124 seen.add(key)
125 return records, duplicates, issues
128def _parquet_magic(path: Path) -> bool:
129 flags = os.O_RDONLY
130 if hasattr(os, "O_NOFOLLOW"):
131 flags |= os.O_NOFOLLOW
132 descriptor = os.open(path, flags)
133 try:
134 mode = os.fstat(descriptor).st_mode
135 size = os.fstat(descriptor).st_size
136 if not stat.S_ISREG(mode) or size < 8:
137 return False
138 first = os.read(descriptor, 4)
139 os.lseek(descriptor, -4, os.SEEK_END)
140 last = os.read(descriptor, 4)
141 return first == b"PAR1" and last == b"PAR1"
142 finally:
143 os.close(descriptor)
146def inspect(args: argparse.Namespace) -> tuple[dict, int]:
147 if not args.assembly.strip() or len(args.assembly) > 200:
148 raise SafetyError("assembly must be a nonempty value of at most 200 characters")
149 model_dir = local_path(args.model_dir, kind="dir")
150 expected_universe = local_path(args.universe, kind="file")
151 config_path = _safe_child(model_dir, args.config_name)
152 checkpoint_path = _safe_child(model_dir, args.checkpoint_name)
153 bundle_universe = _safe_child(model_dir, args.bundle_universe_name)
155 errors: Counter[str] = Counter()
156 warnings: Counter[str] = Counter()
157 checks: list[dict[str, str | int | bool | None]] = []
159 expected_digest, expected_size = sha256_file(
160 expected_universe,
161 max_bytes=args.max_bytes,
162 )
163 bundle_digest, bundle_size = sha256_file(
164 bundle_universe,
165 max_bytes=args.max_bytes,
166 )
167 universe_rows, duplicate_rows, universe_issues = _count_universe(
168 expected_universe,
169 max_bytes=args.max_bytes,
170 max_records=args.max_universe_records,
171 )
172 for code, count in universe_issues.items():
173 errors[f"universe:{code}"] = count
174 if duplicate_rows:
175 errors["universe:duplicate_intervals"] = duplicate_rows
176 if universe_rows == 0:
177 errors["universe:no_valid_records"] += 1
179 universe_match = expected_digest == bundle_digest
180 if not universe_match:
181 errors["bundle_universe_checksum_mismatch"] += 1
182 checks.append(
183 {
184 "check": "exact_universe_bytes_and_order",
185 "passed": universe_match,
186 "expected_sha256": expected_digest,
187 "bundle_sha256": bundle_digest,
188 }
189 )
191 config = simple_yaml_mapping(
192 config_path,
193 max_bytes=min(args.max_bytes, 16 * 1024**2),
194 )
195 vocab_size = config.get("vocab_size")
196 embedding_dim = config.get("embedding_dim")
197 old_embedding_dim = config.get("embedding_size")
198 if embedding_dim is None and old_embedding_dim is not None:
199 embedding_dim = old_embedding_dim
200 warnings["deprecated_config_key:embedding_size"] += 1
201 if type(vocab_size) is not int or vocab_size <= 0:
202 errors["config:invalid_vocab_size"] += 1
203 if type(embedding_dim) is not int or embedding_dim <= 0:
204 errors["config:invalid_embedding_dim"] += 1
205 pooling = config.get("pooling_method")
206 if pooling is not None and pooling not in {"mean", "max"}:
207 errors["config:invalid_pooling_method"] += 1
209 expected_vocab_size = universe_rows + args.special_token_count
210 vocab_match = type(vocab_size) is int and vocab_size == expected_vocab_size
211 if not vocab_match:
212 errors["model_tokenizer_vocab_size_mismatch"] += 1
213 checks.append(
214 {
215 "check": "vocab_size_equals_universe_plus_special_tokens",
216 "passed": vocab_match,
217 "model_vocab_size": vocab_size,
218 "universe_records": universe_rows,
219 "special_token_count": args.special_token_count,
220 "expected_vocab_size": expected_vocab_size,
221 }
222 )
224 config_digest, config_size = sha256_file(
225 config_path,
226 max_bytes=min(args.max_bytes, 16 * 1024**2),
227 )
228 checkpoint_digest, checkpoint_size = sha256_file(
229 checkpoint_path,
230 max_bytes=args.max_bytes,
231 )
232 if checkpoint_size == 0:
233 errors["checkpoint:empty"] += 1
235 token_corpus_summary = None
236 if args.token_corpus:
237 token_path = local_path(args.token_corpus, kind="file")
238 token_digest, token_size = sha256_file(token_path, max_bytes=args.max_bytes)
239 magic_ok = _parquet_magic(token_path)
240 if not magic_ok:
241 errors["token_corpus:invalid_parquet_magic"] += 1
242 token_corpus_summary = {
243 "path": display_path(token_path, 4, args.path_mode),
244 "sha256": token_digest,
245 "size_bytes": token_size,
246 "parquet_magic_valid": magic_ok,
247 "schema_inspected": False,
248 }
250 warnings["assembly_provenance_not_embedded_in_standard_config"] += 1
251 report = {
252 "ok": not errors,
253 "tool": TOOL,
254 "contract": {
255 "assembly": args.assembly,
256 "coordinate_system": "0-based-half-open",
257 "geniml_version": args.geniml_version,
258 "gtars_version": args.gtars_version,
259 "network_used": False,
260 "model_deserialized": False,
261 },
262 "paths": {
263 "model_dir": display_path(model_dir, 1, args.path_mode),
264 "expected_universe": display_path(expected_universe, 2, args.path_mode),
265 "bundle_universe": display_path(bundle_universe, 3, args.path_mode),
266 },
267 "artifacts": {
268 "checkpoint": {
269 "name": checkpoint_path.name,
270 "sha256": checkpoint_digest,
271 "size_bytes": checkpoint_size,
272 "deserialized": False,
273 },
274 "config": {
275 "name": config_path.name,
276 "sha256": config_digest,
277 "size_bytes": config_size,
278 "metadata": {
279 "vocab_size": vocab_size,
280 "embedding_dim": embedding_dim,
281 "pooling_method": pooling,
282 },
283 },
284 "expected_universe": {
285 "sha256": expected_digest,
286 "size_bytes": expected_size,
287 "records": universe_rows,
288 "duplicate_intervals": duplicate_rows,
289 },
290 "bundle_universe": {
291 "sha256": bundle_digest,
292 "size_bytes": bundle_size,
293 },
294 "token_corpus": token_corpus_summary,
295 },
296 "checks": checks,
297 "errors": dict(sorted(errors.items())),
298 "warnings": dict(sorted(warnings.items())),
299 "next_checks": [
300 "compare assembly accession and chromosome-sizes checksum",
301 "instantiate the pinned tokenizer only after this plan passes",
302 "verify every special-token name and ID in a synthetic smoke test",
303 "verify all token IDs are in range and Parquet has one list-valued tokens column",
304 "load the checkpoint only in an isolated, resource-bounded environment",
305 ],
306 }
307 return report, 0 if not errors else 2
310def main() -> int:
311 args = build_parser().parse_args()
312 try:
313 report, exit_code = inspect(args)
314 print_json(report)
315 return exit_code
316 except (OSError, SafetyError, UnicodeError) as exc:
317 return fail_json(TOOL, exc)
320if __name__ == "__main__":
321 sys.exit(main())