Coverage for transformer_lens/tools/model_registry/registry_io.py: 87%
137 statements
« prev ^ index » next coverage.py v7.10.1, created at 2026-09-01 16:23 +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 pathlib import Path
12from typing import Callable, Optional
14from transformer_lens.benchmarks.text_quality_profiles import (
15 P4_SCORING_VERSION,
16 is_default_profile,
17)
19from .verification import VerificationHistory, VerificationRecord
21logger = logging.getLogger(__name__)
23_DATA_DIR = Path(__file__).parent / "data"
24_SUPPORTED_MODELS_PATH = _DATA_DIR / "supported_models.json"
25_VERIFICATION_HISTORY_PATH = _DATA_DIR / "verification_history.json"
27# Status codes
28STATUS_UNVERIFIED = 0
29STATUS_VERIFIED = 1
30STATUS_SKIPPED = 2
31STATUS_FAILED = 3
32# Structural-only pass (--no-hf-reference): Phase 1 ran without an HF reference,
33# so the forward was never numerically compared to HuggingFace. Recorded as a
34# real result but deliberately NOT counted as verified.
35STATUS_PROVISIONAL = 4
37# HF-loadable quantization formats. Admitted to the registry; verification gates
38# on `required_quant_library_for_model()` at run time.
39_HF_LOADABLE_QUANT_PATTERNS = [
40 "-awq",
41 "_awq",
42 "-AWQ",
43 "_AWQ",
44 "-gptq",
45 "_gptq",
46 "-GPTQ",
47 "_GPTQ",
48 "GPTQ",
49 "-bnb-",
50 "_bnb_",
51 "bnb-4bit",
52 "bnb-8bit",
53 "-4bit",
54 "_4bit",
55 "-8bit",
56 "_8bit",
57 "-int4",
58 "_int4",
59 "-int8",
60 "_int8",
61 "-w4a16",
62 "-w8a8",
63 "-W4A16",
64 "-W8A8",
65 ".w4a16",
66 ".W4A16",
67 "-hqq",
68 "_hqq",
69 "-HQQ",
70 "_HQQ",
71 "-3bit",
72 "_3bit",
73 "-2bit",
74 "_2bit",
75 "-5bit",
76 "-6bit",
77 "-oQ",
78 "_oQ",
79 "-quantized.",
80 "_Quantized",
81 "-Quantized",
82]
84# Formats that need a non-HF loader (GGUF→llama.cpp, MLX→Apple, FP4/FP8→NVIDIA).
85_INCOMPATIBLE_QUANT_PATTERNS = [
86 "-gguf",
87 "_gguf",
88 "-GGUF",
89 "_GGUF",
90 "mlx-community/",
91 "-mlx",
92 "-MLX",
93 "_mlx",
94 "_MLX",
95 ".mlx",
96 ".MLX",
97 "-fp8",
98 "_fp8",
99 "-FP8",
100 "_FP8",
101 "-nvfp4",
102 "_nvfp4",
103 "-NVFP4",
104 "_NVFP4",
105 "-mxfp4",
106 "_mxfp4",
107 "-MXFP4",
108 "_MXFP4",
109]
111# Values are Python import names, not PyPI package names. Order matters: explicit
112# format markers must precede generic bit-width markers (HQQ-4bit IDs match both).
113_QUANT_LIBRARY_BY_PATTERN: list[tuple[tuple[str, ...], str]] = [
114 (("-hqq", "_hqq", "-HQQ", "_HQQ"), "hqq"),
115 (("-gptq", "_gptq", "-GPTQ", "_GPTQ", "GPTQ"), "auto_gptq"),
116 (("-awq", "_awq", "-AWQ", "_AWQ"), "awq"),
117 (("-w4a16", "-w8a8", "-W4A16", "-W8A8", ".w4a16", ".W4A16"), "auto_gptq"),
118 (("-bnb-", "_bnb_", "bnb-4bit", "bnb-8bit"), "bitsandbytes"),
119 (("-4bit", "_4bit", "-8bit", "_8bit", "-int4", "_int4", "-int8", "_int8"), "bitsandbytes"),
120]
122QUANTIZED_NOTE = "Quantized format not loadable by HF transformers"
125def is_incompatible_quantized(model_id: str) -> bool:
126 """True for quantization formats the bridge can't ingest (GGUF, MLX, FP4/FP8)."""
127 return any(pat in model_id for pat in _INCOMPATIBLE_QUANT_PATTERNS)
130def is_hf_loadable_quantized(model_id: str) -> bool:
131 """True for quantizations loadable by HF transformers + a quant library."""
132 return any(pat in model_id for pat in _HF_LOADABLE_QUANT_PATTERNS)
135def required_quant_library_for_model(model_id: str) -> Optional[str]:
136 """Return the Python import name needed to load this model, or None if unquantized."""
137 for patterns, library in _QUANT_LIBRARY_BY_PATTERN:
138 if any(pat in model_id for pat in patterns):
139 return library
140 return None
143def is_quantized_model(model_id: str) -> bool:
144 """Alias for ``is_incompatible_quantized`` — kept for back-compat with existing call sites."""
145 return is_incompatible_quantized(model_id)
148def load_supported_models_raw() -> dict:
149 """Load supported_models.json as a raw dict."""
150 with open(_SUPPORTED_MODELS_PATH) as f:
151 return json.load(f)
154def save_supported_models_raw(data: dict) -> None:
155 """Save raw dict back to supported_models.json."""
156 with open(_SUPPORTED_MODELS_PATH, "w") as f:
157 json.dump(data, f, indent=2)
158 f.write("\n")
161def load_verification_history() -> VerificationHistory:
162 """Load verification_history.json into a VerificationHistory dataclass."""
163 if _VERIFICATION_HISTORY_PATH.exists():
164 with open(_VERIFICATION_HISTORY_PATH) as f:
165 data = json.load(f)
166 return VerificationHistory.from_dict(data)
167 return VerificationHistory()
170def save_verification_history(history: VerificationHistory) -> None:
171 """Save VerificationHistory dataclass to verification_history.json."""
172 with open(_VERIFICATION_HISTORY_PATH, "w") as f:
173 json.dump(history.to_dict(), f, indent=2)
174 f.write("\n")
177def _get_tl_version() -> Optional[str]:
178 """Get the current TransformerLens version, or None."""
179 try:
180 import transformer_lens
182 version = getattr(transformer_lens, "__version__", None)
183 if version: 183 ↛ 184line 183 didn't jump to line 184 because the condition on line 183 was never true
184 return str(version)
185 # The package exports no __version__; installed-distribution
186 # metadata is the fallback (dev installs record 0.0.0).
187 from importlib.metadata import version as dist_version
189 return dist_version("transformer-lens")
190 except Exception:
191 return None
194def update_model_status(
195 model_id: str,
196 arch_id: str,
197 status: Optional[int] = None,
198 note: Optional[str] = None,
199 phase_scores: Optional[dict[int, Optional[float]]] = None,
200 sanitize_fn: Optional[Callable[[Optional[str]], Optional[str]]] = None,
201 prompt_profile: Optional[str] = None,
202) -> bool:
203 """Update a single model entry in supported_models.json.
205 If the model is not found in the registry and status is STATUS_VERIFIED or
206 STATUS_PROVISIONAL, a new entry is appended.
208 When status is None (partial-phase update), only the provided phase_scores
209 are updated — status, note, and other scores are preserved.
211 Args:
212 model_id: The model to update
213 arch_id: Architecture of the model
214 status: New status code (0-4), or None for score-only updates
215 note: Optional note for skip/fail reason
216 phase_scores: Phase score dict {1: float, 2: float, 3: float, 4: float}
217 sanitize_fn: Optional callable to sanitize note strings
218 prompt_profile: Phase-4 prompt profile actually used (e.g.
219 "task:translation@en-de"). Sparse: the default "continuation"
220 removes the key (clearing a stale non-default value), None (no
221 Phase-4 result) leaves it untouched.
223 Returns:
224 True if entry was found/created and updated
225 """
226 if phase_scores is None:
227 phase_scores = {}
229 if sanitize_fn and note:
230 note = sanitize_fn(note)
232 data = load_supported_models_raw()
233 updated = False
235 for entry in data.get("models", []):
236 if entry["model_id"] == model_id and entry["architecture_id"] == arch_id:
237 if status is not None: 237 ↛ 243line 237 didn't jump to line 243 because the condition on line 237 was always true
238 entry["status"] = status
239 entry["verified_date"] = (
240 date.today().isoformat() if status != STATUS_UNVERIFIED else None
241 )
242 entry["note"] = note
243 elif note is not None:
244 # Score-only update with an explicit note — overwrite stale notes
245 entry["note"] = note
246 elif phase_scores and "exceeds" in (entry.get("note") or "").lower():
247 # Writing real scores clears a stale memory-skip note
248 entry["note"] = None
249 for phase_num in (1, 2, 3, 4, 7, 8, 9):
250 key = f"phase{phase_num}_score"
251 if phase_num in phase_scores:
252 entry[key] = phase_scores[phase_num]
253 elif key not in entry:
254 entry[key] = None
255 if prompt_profile is not None and is_default_profile(prompt_profile):
256 entry.pop("prompt_profile", None)
257 elif prompt_profile is not None:
258 entry["prompt_profile"] = prompt_profile
259 if 4 in phase_scores:
260 entry["p4_scoring_version"] = P4_SCORING_VERSION
261 # Reorder keys so phase scores are always in numerical order
262 _KEY_ORDER = [
263 "architecture_id",
264 "model_id",
265 "status",
266 "verified_date",
267 "metadata",
268 "note",
269 "prompt_profile",
270 "p4_scoring_version",
271 "phase1_score",
272 "phase2_score",
273 "phase3_score",
274 "phase4_score",
275 "phase7_score",
276 "phase8_score",
277 "phase9_score",
278 ]
279 reordered = {k: entry[k] for k in _KEY_ORDER if k in entry}
280 for k in entry:
281 if k not in reordered: 281 ↛ 282line 281 didn't jump to line 282 because the condition on line 281 was never true
282 reordered[k] = entry[k]
283 entry.clear()
284 entry.update(reordered)
285 updated = True
286 break
288 if not updated and status in (STATUS_VERIFIED, STATUS_PROVISIONAL):
289 # Model not in registry -- add it. A structural-only (provisional) pass
290 # is a real result worth recording; skipped/failed on a missing model
291 # are not, so they still fall through.
292 data.get("models", []).append(
293 {
294 "model_id": model_id,
295 "architecture_id": arch_id,
296 "status": status,
297 "verified_date": date.today().isoformat(),
298 "metadata": None,
299 "note": note,
300 "phase1_score": phase_scores.get(1),
301 "phase2_score": phase_scores.get(2),
302 "phase3_score": phase_scores.get(3),
303 "phase4_score": phase_scores.get(4),
304 "phase7_score": phase_scores.get(7),
305 "phase8_score": phase_scores.get(8),
306 "phase9_score": phase_scores.get(9),
307 }
308 )
309 new_entry = data["models"][-1]
310 extras: list[tuple[str, object]] = []
311 if prompt_profile is not None and not is_default_profile(prompt_profile):
312 extras.append(("prompt_profile", prompt_profile))
313 if phase_scores.get(4) is not None:
314 extras.append(("p4_scoring_version", P4_SCORING_VERSION))
315 if extras:
316 # Keep key position consistent with _KEY_ORDER (after "note").
317 items = list(new_entry.items())
318 idx = [k for k, _ in items].index("note") + 1
319 for offset, pair in enumerate(extras):
320 items.insert(idx + offset, pair)
321 new_entry.clear()
322 new_entry.update(items)
323 updated = True
325 if updated:
326 models = data.get("models", [])
327 data["total_verified"] = sum(1 for m in models if m.get("status", 0) == STATUS_VERIFIED)
328 data["total_provisional"] = sum(
329 1 for m in models if m.get("status", 0) == STATUS_PROVISIONAL
330 )
331 data["total_models"] = len(models)
332 data["total_architectures"] = len(set(m["architecture_id"] for m in models))
333 save_supported_models_raw(data)
335 return updated
338def registry_prompt_profile(model_id: str) -> Optional[str]:
339 """Stored prompt_profile for a model, or None. Uncached read: the sweep
340 rewrites the registry between models."""
341 try:
342 data = load_supported_models_raw()
343 except Exception:
344 return None
345 for entry in data.get("models", []):
346 if entry.get("model_id") == model_id:
347 profile = entry.get("prompt_profile")
348 return profile if isinstance(profile, str) else None
349 return None
352def add_verification_record(
353 model_id: str,
354 arch_id: str,
355 notes: Optional[str] = None,
356 verified_by: str = "verify_models",
357 sanitize_fn: Optional[Callable[[Optional[str]], Optional[str]]] = None,
358 prompt_profile: Optional[str] = None,
359 p4_scoring_version: Optional[int] = None,
360) -> None:
361 """Append a VerificationRecord to verification_history.json.
363 Uses the VerificationRecord dataclass properly instead of raw dict
364 manipulation.
366 Args:
367 model_id: The verified model
368 arch_id: Architecture type
369 notes: Optional verification notes
370 verified_by: Who/what performed the verification
371 sanitize_fn: Optional callable to sanitize note strings
372 """
373 if sanitize_fn and notes:
374 notes = sanitize_fn(notes)
376 record = VerificationRecord(
377 model_id=model_id,
378 architecture_id=arch_id,
379 verified_date=date.today(),
380 verified_by=verified_by,
381 transformerlens_version=_get_tl_version(),
382 notes=notes,
383 prompt_profile=prompt_profile,
384 p4_scoring_version=p4_scoring_version,
385 )
387 history = load_verification_history()
388 history.add_record(record)
389 save_verification_history(history)