Coverage for transformer_lens/tools/model_registry/registry_io.py: 88%
168 statements
« prev ^ index » next coverage.py v7.10.1, created at 2026-09-21 19:27 +0000
1"""Shared I/O functions for reading and writing model registry data files.
3Consolidates the load-modify-save pattern used by verify_models.py and
4main_benchmark.py into a single module that properly uses the
5VerificationRecord/VerificationHistory dataclasses.
6"""
8import json
9import logging
10from datetime import date
11from functools import lru_cache
12from pathlib import Path
13from typing import Callable, Optional
15from .verification import VerificationHistory, VerificationRecord
17logger = logging.getLogger(__name__)
19_DATA_DIR = Path(__file__).parent / "data"
20_SUPPORTED_MODELS_PATH = _DATA_DIR / "supported_models.json"
21_VERIFICATION_HISTORY_PATH = _DATA_DIR / "verification_history.json"
22_MODEL_ALIASES_PATH = _DATA_DIR / "model_aliases.json"
24# Status codes
25STATUS_UNVERIFIED = 0
26STATUS_VERIFIED = 1
27STATUS_SKIPPED = 2
28STATUS_FAILED = 3
29# Structural-only pass (--no-hf-reference): Phase 1 ran without an HF reference,
30# so the forward was never numerically compared to HuggingFace. Recorded as a
31# real result but deliberately NOT counted as verified.
32STATUS_PROVISIONAL = 4
34# Human-readable labels for docs display. STATUS_SKIPPED deliberately renders as
35# "Unverified": a skip (memory/tooling) is not a verification outcome.
36STATUS_LABELS: dict[int, str] = {
37 STATUS_UNVERIFIED: "Unverified",
38 STATUS_VERIFIED: "Verified",
39 STATUS_SKIPPED: "Unverified",
40 STATUS_FAILED: "Failed",
41 STATUS_PROVISIONAL: "Provisional",
42}
44# Registry phase-score columns: text 1-4, multimodal 7, audio 8, vision 9.
45# PHASES is derived so the full column set can never drift from the two groups.
46TEXT_PHASES: tuple[int, ...] = (1, 2, 3, 4)
47MODALITY_PHASES: tuple[int, ...] = (7, 8, 9)
48PHASES: tuple[int, ...] = TEXT_PHASES + MODALITY_PHASES
50# HF-loadable quantization formats. Admitted to the registry; verification gates
51# on `required_quant_library_for_model()` at run time.
52_HF_LOADABLE_QUANT_PATTERNS = [
53 "-awq",
54 "_awq",
55 "-AWQ",
56 "_AWQ",
57 "-gptq",
58 "_gptq",
59 "-GPTQ",
60 "_GPTQ",
61 "GPTQ",
62 "-bnb-",
63 "_bnb_",
64 "bnb-4bit",
65 "bnb-8bit",
66 "-4bit",
67 "_4bit",
68 "-8bit",
69 "_8bit",
70 "-int4",
71 "_int4",
72 "-int8",
73 "_int8",
74 "-w4a16",
75 "-w8a8",
76 "-W4A16",
77 "-W8A8",
78 ".w4a16",
79 ".W4A16",
80 "-hqq",
81 "_hqq",
82 "-HQQ",
83 "_HQQ",
84 "-3bit",
85 "_3bit",
86 "-2bit",
87 "_2bit",
88 "-5bit",
89 "-6bit",
90 "-oQ",
91 "_oQ",
92 "-quantized.",
93 "_Quantized",
94 "-Quantized",
95]
97# Formats that need a non-HF loader (GGUF→llama.cpp, MLX→Apple, FP4/FP8→NVIDIA).
98_INCOMPATIBLE_QUANT_PATTERNS = [
99 "-gguf",
100 "_gguf",
101 "-GGUF",
102 "_GGUF",
103 "mlx-community/",
104 "-mlx",
105 "-MLX",
106 "_mlx",
107 "_MLX",
108 ".mlx",
109 ".MLX",
110 "-fp8",
111 "_fp8",
112 "-FP8",
113 "_FP8",
114 "-nvfp4",
115 "_nvfp4",
116 "-NVFP4",
117 "_NVFP4",
118 "-mxfp4",
119 "_mxfp4",
120 "-MXFP4",
121 "_MXFP4",
122]
124# Values are Python import names, not PyPI package names. Order matters: explicit
125# format markers must precede generic bit-width markers (HQQ-4bit IDs match both).
126_QUANT_LIBRARY_BY_PATTERN: list[tuple[tuple[str, ...], str]] = [
127 (("-hqq", "_hqq", "-HQQ", "_HQQ"), "hqq"),
128 (("-gptq", "_gptq", "-GPTQ", "_GPTQ", "GPTQ"), "auto_gptq"),
129 (("-awq", "_awq", "-AWQ", "_AWQ"), "awq"),
130 (("-w4a16", "-w8a8", "-W4A16", "-W8A8", ".w4a16", ".W4A16"), "auto_gptq"),
131 (("-bnb-", "_bnb_", "bnb-4bit", "bnb-8bit"), "bitsandbytes"),
132 (("-4bit", "_4bit", "-8bit", "_8bit", "-int4", "_int4", "-int8", "_int8"), "bitsandbytes"),
133]
135QUANTIZED_NOTE = "Quantized format not loadable by HF transformers"
138def is_incompatible_quantized(model_id: str) -> bool:
139 """True for quantization formats the bridge can't ingest (GGUF, MLX, FP4/FP8)."""
140 return any(pat in model_id for pat in _INCOMPATIBLE_QUANT_PATTERNS)
143def is_hf_loadable_quantized(model_id: str) -> bool:
144 """True for quantizations loadable by HF transformers + a quant library."""
145 return any(pat in model_id for pat in _HF_LOADABLE_QUANT_PATTERNS)
148def required_quant_library_for_model(model_id: str) -> Optional[str]:
149 """Return the Python import name needed to load this model, or None if unquantized."""
150 for patterns, library in _QUANT_LIBRARY_BY_PATTERN:
151 if any(pat in model_id for pat in patterns):
152 return library
153 return None
156def is_quantized_model(model_id: str) -> bool:
157 """Alias for ``is_incompatible_quantized`` — kept for back-compat with existing call sites."""
158 return is_incompatible_quantized(model_id)
161@lru_cache(maxsize=1)
162def load_model_aliases() -> dict[str, list[str]]:
163 """Load the canonical alias table: official HF model name -> deprecated short aliases."""
164 with open(_MODEL_ALIASES_PATH) as f:
165 return json.load(f)["aliases"]
168def resolve_model_alias(model_name: str) -> Optional[str]:
169 """Return the official HF name if ``model_name`` is a deprecated alias, else None."""
170 for official_name, aliases in load_model_aliases().items():
171 if model_name in aliases:
172 return official_name
173 return None
176def load_supported_models_raw() -> dict:
177 """Load supported_models.json as a raw dict."""
178 with open(_SUPPORTED_MODELS_PATH) as f:
179 return json.load(f)
182def save_supported_models_raw(data: dict) -> None:
183 """Save raw dict back to supported_models.json."""
184 with open(_SUPPORTED_MODELS_PATH, "w") as f:
185 json.dump(data, f, indent=2)
186 f.write("\n")
189def load_verification_history() -> VerificationHistory:
190 """Load verification_history.json into a VerificationHistory dataclass."""
191 if _VERIFICATION_HISTORY_PATH.exists():
192 with open(_VERIFICATION_HISTORY_PATH) as f:
193 data = json.load(f)
194 return VerificationHistory.from_dict(data)
195 return VerificationHistory()
198def save_verification_history(history: VerificationHistory) -> None:
199 """Save VerificationHistory dataclass to verification_history.json."""
200 with open(_VERIFICATION_HISTORY_PATH, "w") as f:
201 json.dump(history.to_dict(), f, indent=2)
202 f.write("\n")
205def _get_tl_version() -> Optional[str]:
206 """Get the current TransformerLens version, or None."""
207 try:
208 import transformer_lens
210 version = getattr(transformer_lens, "__version__", None)
211 if version: 211 ↛ 212line 211 didn't jump to line 212 because the condition on line 211 was never true
212 return str(version)
213 # The package exports no __version__; installed-distribution
214 # metadata is the fallback (dev installs record 0.0.0).
215 from importlib.metadata import version as dist_version
217 return dist_version("transformer-lens")
218 except Exception:
219 return None
222def pass_status(use_hf_reference: bool) -> int:
223 """Status for a passing run: VERIFIED with an HF reference, else PROVISIONAL
224 (a --no-hf-reference structural-only pass is recorded but not counted verified)."""
225 return STATUS_VERIFIED if use_hf_reference else STATUS_PROVISIONAL
228def extract_phase_scores(results: list) -> dict[int, Optional[float]]:
229 """Extract phase scores from benchmark results.
231 Shared home for both registry-writing paths (verify_models and
232 main_benchmark.update_model_registry) so they cannot drift.
234 Args:
235 results: List of BenchmarkResult objects
237 Returns:
238 Dict mapping phase number to score (0-100) or None
239 """
240 from transformer_lens.benchmarks.utils import BenchmarkSeverity
242 phase_results: dict[int, list[bool]] = {phase: [] for phase in PHASES}
243 for result in results:
244 if result.phase in phase_results and result.severity != BenchmarkSeverity.SKIPPED:
245 phase_results[result.phase].append(result.passed)
247 scores: dict[int, Optional[float]] = {}
248 for phase, passed_list in phase_results.items():
249 if passed_list:
250 scores[phase] = round(sum(passed_list) / len(passed_list) * 100, 1)
251 # Omit phases with no results — they weren't run, so their
252 # existing registry scores should be preserved.
254 # Phase 4 (text quality): store the actual 0-100 quality score from the
255 # benchmark details instead of a binary pass/fail percentage.
256 if 4 in scores:
257 for result in results: 257 ↛ 262line 257 didn't jump to line 262 because the loop on line 257 didn't complete
258 if result.phase == 4 and result.details and "score" in result.details:
259 scores[4] = round(result.details["score"], 1)
260 break
262 return scores
265def recompute_registry_totals(models: list[dict]) -> dict:
266 """Header totals for supported_models.json, recomputed from the models list.
268 Shared by both writers (``update_model_status`` here and hf_scraper's report
269 builder) so the counting rules cannot drift.
270 """
271 return {
272 "total_architectures": len({m["architecture_id"] for m in models}),
273 "total_models": len(models),
274 "total_verified": sum(1 for m in models if m.get("status", 0) == STATUS_VERIFIED),
275 "total_provisional": sum(1 for m in models if m.get("status", 0) == STATUS_PROVISIONAL),
276 }
279def update_model_status(
280 model_id: str,
281 arch_id: str,
282 status: Optional[int] = None,
283 note: Optional[str] = None,
284 phase_scores: Optional[dict[int, Optional[float]]] = None,
285 sanitize_fn: Optional[Callable[[Optional[str]], Optional[str]]] = None,
286 prompt_profile: Optional[str] = None,
287) -> bool:
288 """Update a single model entry in supported_models.json.
290 If the model is not found in the registry and status is STATUS_VERIFIED or
291 STATUS_PROVISIONAL, a new entry is appended.
293 When status is None (partial-phase update), only the provided phase_scores
294 are updated — status, note, and other scores are preserved.
296 Args:
297 model_id: The model to update
298 arch_id: Architecture of the model
299 status: New status code (0-4), or None for score-only updates
300 note: Optional note for skip/fail reason
301 phase_scores: Phase score dict {1: float, 2: float, 3: float, 4: float}
302 sanitize_fn: Optional callable to sanitize note strings
303 prompt_profile: Phase-4 prompt profile actually used (e.g.
304 "task:translation@en-de"). Sparse: the default "continuation"
305 removes the key (clearing a stale non-default value), None (no
306 Phase-4 result) leaves it untouched.
308 Returns:
309 True if entry was found/created and updated
310 """
311 # Deferred: benchmarks imports model_bridge, and this module loads during
312 # `import transformer_lens` (via supported_models) — importing it at module
313 # scope closes an import cycle back through the bridge's HF source.
314 from transformer_lens.benchmarks.text_quality_profiles import (
315 P4_SCORING_VERSION,
316 is_default_profile,
317 )
319 if phase_scores is None:
320 phase_scores = {}
322 if sanitize_fn and note:
323 note = sanitize_fn(note)
325 data = load_supported_models_raw()
326 updated = False
328 for entry in data.get("models", []):
329 if entry["model_id"] == model_id and entry["architecture_id"] == arch_id:
330 if status is not None: 330 ↛ 336line 330 didn't jump to line 336 because the condition on line 330 was always true
331 entry["status"] = status
332 entry["verified_date"] = (
333 date.today().isoformat() if status != STATUS_UNVERIFIED else None
334 )
335 entry["note"] = note
336 elif note is not None:
337 # Score-only update with an explicit note — overwrite stale notes
338 entry["note"] = note
339 elif phase_scores and "exceeds" in (entry.get("note") or "").lower():
340 # Writing real scores clears a stale memory-skip note
341 entry["note"] = None
342 for phase_num in PHASES:
343 key = f"phase{phase_num}_score"
344 if phase_num in phase_scores:
345 entry[key] = phase_scores[phase_num]
346 elif key not in entry:
347 entry[key] = None
348 if prompt_profile is not None and is_default_profile(prompt_profile):
349 entry.pop("prompt_profile", None)
350 elif prompt_profile is not None:
351 entry["prompt_profile"] = prompt_profile
352 if 4 in phase_scores:
353 entry["p4_scoring_version"] = P4_SCORING_VERSION
354 # Reorder keys so phase scores are always in numerical order
355 _KEY_ORDER = [
356 "architecture_id",
357 "model_id",
358 "status",
359 "verified_date",
360 "metadata",
361 "note",
362 "prompt_profile",
363 "p4_scoring_version",
364 *[f"phase{p}_score" for p in PHASES],
365 ]
366 reordered = {k: entry[k] for k in _KEY_ORDER if k in entry}
367 for k in entry:
368 if k not in reordered: 368 ↛ 369line 368 didn't jump to line 369 because the condition on line 368 was never true
369 reordered[k] = entry[k]
370 entry.clear()
371 entry.update(reordered)
372 updated = True
373 break
375 if not updated and status in (STATUS_VERIFIED, STATUS_PROVISIONAL):
376 # Model not in registry -- add it. A structural-only (provisional) pass
377 # is a real result worth recording; skipped/failed on a missing model
378 # are not, so they still fall through.
379 data.get("models", []).append(
380 {
381 "model_id": model_id,
382 "architecture_id": arch_id,
383 "status": status,
384 "verified_date": date.today().isoformat(),
385 "metadata": None,
386 "note": note,
387 **{f"phase{p}_score": phase_scores.get(p) for p in PHASES},
388 }
389 )
390 new_entry = data["models"][-1]
391 extras: list[tuple[str, object]] = []
392 if prompt_profile is not None and not is_default_profile(prompt_profile):
393 extras.append(("prompt_profile", prompt_profile))
394 if phase_scores.get(4) is not None:
395 extras.append(("p4_scoring_version", P4_SCORING_VERSION))
396 if extras:
397 # Keep key position consistent with _KEY_ORDER (after "note").
398 items = list(new_entry.items())
399 idx = [k for k, _ in items].index("note") + 1
400 for offset, pair in enumerate(extras):
401 items.insert(idx + offset, pair)
402 new_entry.clear()
403 new_entry.update(items)
404 updated = True
406 if updated:
407 data.update(recompute_registry_totals(data.get("models", [])))
408 save_supported_models_raw(data)
410 return updated
413def registry_prompt_profile(model_id: str) -> Optional[str]:
414 """Stored prompt_profile for a model, or None. Uncached read: the sweep
415 rewrites the registry between models."""
416 try:
417 data = load_supported_models_raw()
418 except Exception:
419 return None
420 for entry in data.get("models", []):
421 if entry.get("model_id") == model_id:
422 profile = entry.get("prompt_profile")
423 return profile if isinstance(profile, str) else None
424 return None
427def add_verification_record(
428 model_id: str,
429 arch_id: str,
430 notes: Optional[str] = None,
431 verified_by: str = "verify_models",
432 sanitize_fn: Optional[Callable[[Optional[str]], Optional[str]]] = None,
433 prompt_profile: Optional[str] = None,
434 p4_scoring_version: Optional[int] = None,
435) -> None:
436 """Append a VerificationRecord to verification_history.json.
438 Uses the VerificationRecord dataclass properly instead of raw dict
439 manipulation.
441 Args:
442 model_id: The verified model
443 arch_id: Architecture type
444 notes: Optional verification notes
445 verified_by: Who/what performed the verification
446 sanitize_fn: Optional callable to sanitize note strings
447 """
448 if sanitize_fn and notes:
449 notes = sanitize_fn(notes)
451 record = VerificationRecord(
452 model_id=model_id,
453 architecture_id=arch_id,
454 verified_date=date.today(),
455 verified_by=verified_by,
456 transformerlens_version=_get_tl_version(),
457 notes=notes,
458 prompt_profile=prompt_profile,
459 p4_scoring_version=p4_scoring_version,
460 )
462 history = load_verification_history()
463 history.add_record(record)
464 save_verification_history(history)