transformer_lens.tools.model_registry.registry_io module

Shared I/O functions for reading and writing model registry data files.

Consolidates the load-modify-save pattern used by verify_models.py and main_benchmark.py into a single module that properly uses the VerificationRecord/VerificationHistory dataclasses.

transformer_lens.tools.model_registry.registry_io.add_verification_record(model_id: str, arch_id: str, notes: str | None = None, verified_by: str = 'verify_models', sanitize_fn: Callable[[str | None], str | None] | None = None) → None

Append a VerificationRecord to verification_history.json.

Uses the VerificationRecord dataclass properly instead of raw dict manipulation.

Parameters:

  • model_id – The verified model

  • arch_id – Architecture type

  • notes – Optional verification notes

  • verified_by – Who/what performed the verification

  • sanitize_fn – Optional callable to sanitize note strings

transformer_lens.tools.model_registry.registry_io.is_quantized_model(model_id: str) → bool

Check if a model ID indicates a quantized model variant.

Detects AWQ, GPTQ, GGUF, BitsAndBytes (bnb), FP8, INT4/INT8, MLX quantized, and other common quantization suffixes.

transformer_lens.tools.model_registry.registry_io.load_supported_models_raw() → dict

Load supported_models.json as a raw dict.

transformer_lens.tools.model_registry.registry_io.load_verification_history() → VerificationHistory

Load verification_history.json into a VerificationHistory dataclass.

transformer_lens.tools.model_registry.registry_io.save_supported_models_raw(data: dict) → None

Save raw dict back to supported_models.json.

transformer_lens.tools.model_registry.registry_io.save_verification_history(history: VerificationHistory) → None

Save VerificationHistory dataclass to verification_history.json.

transformer_lens.tools.model_registry.registry_io.update_model_status(model_id: str, arch_id: str, status: int | None = None, note: str | None = None, phase_scores: dict[int, float | None] | None = None, sanitize_fn: Callable[[str | None], str | None] | None = None) → bool

Update a single model entry in supported_models.json.

If the model is not found in the registry and status == STATUS_VERIFIED, a new entry is appended.

When status is None (partial-phase update), only the provided phase_scores are updated — status, note, and other scores are preserved.

Parameters:

  • model_id – The model to update

  • arch_id – Architecture of the model

  • status – New status code (0-3), or None for score-only updates

  • note – Optional note for skip/fail reason

  • phase_scores – Phase score dict {1: float, 2: float, 3: float, 4: float}

  • sanitize_fn – Optional callable to sanitize note strings

Returns:

True if entry was found/created and updated