Coverage for transformer_lens/tools/model_registry/schemas.py: 89%
142 statements
« prev ^ index » next coverage.py v7.10.1, created at 2026-09-21 19:27 +0000
1"""Data schemas for the model registry.
3This module defines the dataclasses used throughout the model registry for
4representing supported models, architecture gaps, and related metadata.
5"""
7from dataclasses import dataclass, field
8from datetime import date, datetime
9from typing import Optional
12@dataclass
13class ModelMetadata:
14 """Metadata for a model from HuggingFace.
16 Attributes:
17 downloads: Total download count for the model
18 likes: Number of likes/stars on HuggingFace
19 last_modified: When the model was last updated
20 tags: List of tags associated with the model
21 parameter_count: Estimated number of parameters (if available)
22 """
24 downloads: int = 0
25 likes: int = 0
26 last_modified: Optional[datetime] = None
27 tags: list[str] = field(default_factory=list)
28 parameter_count: Optional[int] = None
30 def to_dict(self) -> dict:
31 """Convert to a JSON-serializable dictionary."""
32 return {
33 "downloads": self.downloads,
34 "likes": self.likes,
35 "last_modified": self.last_modified.isoformat() if self.last_modified else None,
36 "tags": self.tags,
37 "parameter_count": self.parameter_count,
38 }
40 @classmethod
41 def from_dict(cls, data: dict) -> "ModelMetadata":
42 """Create from a dictionary."""
43 last_modified = None
44 if data.get("last_modified"):
45 last_modified = datetime.fromisoformat(data["last_modified"])
46 return cls(
47 downloads=data.get("downloads", 0),
48 likes=data.get("likes", 0),
49 last_modified=last_modified,
50 tags=data.get("tags", []),
51 parameter_count=data.get("parameter_count"),
52 )
55@dataclass
56class ModelEntry:
57 """A single model entry in the supported models list.
59 Attributes:
60 architecture_id: The architecture type (e.g., "GPT2LMHeadModel")
61 model_id: The HuggingFace model ID (e.g., "gpt2", "openai-community/gpt2")
62 status: Verification status (0=unverified, 1=verified, 2=skipped, 3=failed,
63 4=provisional — structural-only pass via --no-hf-reference, not counted as verified)
64 verified_date: Date when verification was performed
65 metadata: Optional metadata from HuggingFace
66 note: Optional note (skip/fail reason, e.g. "Estimated 48 GB exceeds 16 GB limit")
67 prompt_profile: Phase-4 prompt profile used (e.g. "task:translation@en-de");
68 omitted from JSON for the default continuation profile
69 phase1_score: Benchmark Phase 1 score (HF vs Bridge), 0-100 or None
70 phase2_score: Benchmark Phase 2 score (unprocessed runtime self-checks +
71 HF logits/loss equivalence), 0-100 or None
72 phase3_score: Benchmark Phase 3 score (compatibility mode + HF logits/loss
73 equivalence), 0-100 or None
74 phase4_score: Benchmark Phase 4 score (generation / text quality), 0-100 or None
75 phase7_score: Benchmark Phase 7 score (multimodal), 0-100 or None
76 phase8_score: Benchmark Phase 8 score (audio), 0-100 or None
77 phase9_score: Benchmark Phase 9 score (vision), 0-100 or None
78 """
80 architecture_id: str
81 model_id: str
82 status: int = 0
83 verified_date: Optional[date] = None
84 metadata: Optional[ModelMetadata] = None
85 note: Optional[str] = None
86 prompt_profile: Optional[str] = None
87 p4_scoring_version: Optional[int] = None
88 phase1_score: Optional[float] = None
89 phase2_score: Optional[float] = None
90 phase3_score: Optional[float] = None
91 phase4_score: Optional[float] = None
92 phase7_score: Optional[float] = None
93 phase8_score: Optional[float] = None
94 phase9_score: Optional[float] = None
96 def to_dict(self) -> dict:
97 """Convert to a JSON-serializable dictionary. prompt_profile is sparse:
98 omitted when None so default-profile entries carry no key."""
99 result = {
100 "architecture_id": self.architecture_id,
101 "model_id": self.model_id,
102 "status": self.status,
103 "verified_date": self.verified_date.isoformat() if self.verified_date else None,
104 "metadata": self.metadata.to_dict() if self.metadata else None,
105 "note": self.note,
106 "phase1_score": self.phase1_score,
107 "phase2_score": self.phase2_score,
108 "phase3_score": self.phase3_score,
109 "phase4_score": self.phase4_score,
110 "phase7_score": self.phase7_score,
111 "phase8_score": self.phase8_score,
112 "phase9_score": self.phase9_score,
113 }
114 extras: list[tuple[str, object]] = []
115 if self.prompt_profile is not None: 115 ↛ 116line 115 didn't jump to line 116 because the condition on line 115 was never true
116 extras.append(("prompt_profile", self.prompt_profile))
117 if self.p4_scoring_version is not None: 117 ↛ 118line 117 didn't jump to line 118 because the condition on line 117 was never true
118 extras.append(("p4_scoring_version", self.p4_scoring_version))
119 if extras: 119 ↛ 120line 119 didn't jump to line 120 because the condition on line 119 was never true
120 note_index = list(result).index("note") + 1
121 items = list(result.items())
122 for offset, pair in enumerate(extras):
123 items.insert(note_index + offset, pair)
124 result = dict(items)
125 return result
127 @classmethod
128 def from_dict(cls, data: dict) -> "ModelEntry":
129 """Create from a dictionary."""
130 verified_date = None
131 if data.get("verified_date"):
132 verified_date = date.fromisoformat(data["verified_date"])
133 metadata = None
134 if data.get("metadata"):
135 metadata = ModelMetadata.from_dict(data["metadata"])
136 # Backwards compat: convert old "verified" bool to new "status" int
137 if "status" in data:
138 status = data["status"]
139 elif data.get("verified", False):
140 status = 1
141 else:
142 status = 0
143 return cls(
144 architecture_id=data["architecture_id"],
145 model_id=data["model_id"],
146 status=status,
147 verified_date=verified_date,
148 metadata=metadata,
149 note=data.get("note"),
150 prompt_profile=data.get("prompt_profile"),
151 p4_scoring_version=data.get("p4_scoring_version"),
152 phase1_score=data.get("phase1_score"),
153 phase2_score=data.get("phase2_score"),
154 phase3_score=data.get("phase3_score"),
155 phase4_score=data.get("phase4_score"),
156 phase7_score=data.get("phase7_score"),
157 phase8_score=data.get("phase8_score"),
158 phase9_score=data.get("phase9_score"),
159 )
162@dataclass
163class ArchitectureGap:
164 """An unsupported architecture with model count and relevancy metrics.
166 Attributes:
167 architecture_id: The architecture type not supported by TransformerLens
168 total_models: Number of models on HuggingFace using this architecture
169 sample_models: Top models by downloads for this architecture (up to 10)
170 total_downloads: Aggregate download count across all models of this architecture
171 min_param_count: Parameter count of the smallest model (None if unknown)
172 relevancy_score: Composite relevancy score (0-100), or None if not computed
173 """
175 architecture_id: str
176 total_models: int
177 sample_models: list[str] = field(default_factory=list)
178 total_downloads: int = 0
179 min_param_count: Optional[int] = None
180 relevancy_score: Optional[float] = None
182 def to_dict(self) -> dict:
183 """Convert to a JSON-serializable dictionary."""
184 d: dict = {
185 "architecture_id": self.architecture_id,
186 "total_models": self.total_models,
187 "total_downloads": self.total_downloads,
188 "min_param_count": self.min_param_count,
189 "relevancy_score": self.relevancy_score,
190 "sample_models": self.sample_models,
191 }
192 return d
194 @classmethod
195 def from_dict(cls, data: dict) -> "ArchitectureGap":
196 """Create from a dictionary."""
197 return cls(
198 architecture_id=data["architecture_id"],
199 total_models=data["total_models"],
200 sample_models=data.get("sample_models", []),
201 total_downloads=data.get("total_downloads", 0),
202 min_param_count=data.get("min_param_count"),
203 relevancy_score=data.get("relevancy_score"),
204 )
207@dataclass
208class ScanInfo:
209 """Metadata about a scraping run.
211 Attributes:
212 total_scanned: Total number of models scanned in this run
213 task_filter: HuggingFace task filter used (e.g., "text-generation")
214 scan_duration_seconds: How long the scan took in seconds (if available)
215 """
217 total_scanned: int
218 task_filter: str
219 scan_duration_seconds: Optional[float] = None
221 def to_dict(self) -> dict:
222 """Convert to a JSON-serializable dictionary."""
223 d: dict = {
224 "total_scanned": self.total_scanned,
225 "task_filter": self.task_filter,
226 }
227 if self.scan_duration_seconds is not None: 227 ↛ 228line 227 didn't jump to line 228 because the condition on line 227 was never true
228 d["scan_duration_seconds"] = self.scan_duration_seconds
229 return d
231 @classmethod
232 def from_dict(cls, data: dict) -> "ScanInfo":
233 """Create from a dictionary."""
234 return cls(
235 total_scanned=data["total_scanned"],
236 task_filter=data["task_filter"],
237 scan_duration_seconds=data.get("scan_duration_seconds"),
238 )
241@dataclass
242class SupportedModelsReport:
243 """Report containing all supported models.
245 Attributes:
246 generated_at: Date when this report was generated
247 scan_info: Metadata about the scraping run
248 total_architectures: Number of unique supported architectures
249 total_models: Total number of supported models
250 total_verified: Number of models that have been verified
251 models: List of all model entries
252 """
254 generated_at: date
255 total_models: int
256 models: list[ModelEntry]
257 scan_info: Optional[ScanInfo] = None
258 total_architectures: int = 0
259 total_verified: int = 0
260 # Structural-only (--no-hf-reference) passes; not counted as verified.
261 total_provisional: int = 0
263 def to_dict(self) -> dict:
264 """Convert to a JSON-serializable dictionary."""
265 d: dict = {
266 "generated_at": self.generated_at.isoformat(),
267 "scan_info": self.scan_info.to_dict() if self.scan_info else None,
268 "total_architectures": self.total_architectures,
269 "total_models": self.total_models,
270 "total_verified": self.total_verified,
271 "total_provisional": self.total_provisional,
272 "models": [m.to_dict() for m in self.models],
273 }
274 return d
276 @classmethod
277 def from_dict(cls, data: dict) -> "SupportedModelsReport":
278 """Create from a dictionary."""
279 scan_info = None
280 if data.get("scan_info"): 280 ↛ 282line 280 didn't jump to line 282 because the condition on line 280 was always true
281 scan_info = ScanInfo.from_dict(data["scan_info"])
282 return cls(
283 generated_at=date.fromisoformat(data["generated_at"]),
284 scan_info=scan_info,
285 total_architectures=data.get("total_architectures", 0),
286 total_models=data.get("total_models", len(data.get("models", []))),
287 total_verified=data.get("total_verified", 0),
288 total_provisional=data.get("total_provisional", 0),
289 models=[ModelEntry.from_dict(m) for m in data["models"]],
290 )
293@dataclass
294class ArchitectureGapsReport:
295 """Report containing unsupported architectures.
297 Attributes:
298 generated_at: Date when this report was generated
299 scan_info: Metadata about the scraping run
300 total_unsupported_architectures: Number of unsupported architectures
301 total_unsupported_models: Total models across all unsupported architectures
302 gaps: List of architecture gaps sorted by model count
303 """
305 generated_at: date
306 gaps: list[ArchitectureGap]
307 scan_info: Optional[ScanInfo] = None
308 total_unsupported_architectures: int = 0
309 total_unsupported_models: int = 0
311 def to_dict(self) -> dict:
312 """Convert to a JSON-serializable dictionary."""
313 return {
314 "generated_at": self.generated_at.isoformat(),
315 "scan_info": self.scan_info.to_dict() if self.scan_info else None,
316 "total_unsupported_architectures": self.total_unsupported_architectures,
317 "total_unsupported_models": self.total_unsupported_models,
318 "gaps": [g.to_dict() for g in self.gaps],
319 }
321 @classmethod
322 def from_dict(cls, data: dict) -> "ArchitectureGapsReport":
323 """Create from a dictionary."""
324 scan_info = None
325 if data.get("scan_info"): 325 ↛ 327line 325 didn't jump to line 327 because the condition on line 325 was always true
326 scan_info = ScanInfo.from_dict(data["scan_info"])
327 gaps = [ArchitectureGap.from_dict(g) for g in data["gaps"]]
328 return cls(
329 generated_at=date.fromisoformat(data["generated_at"]),
330 scan_info=scan_info,
331 total_unsupported_architectures=data.get(
332 "total_unsupported_architectures",
333 data.get("total_unsupported", len(gaps)),
334 ),
335 total_unsupported_models=data.get(
336 "total_unsupported_models",
337 sum(g.total_models for g in gaps),
338 ),
339 gaps=gaps,
340 )
343@dataclass
344class ArchitectureStats:
345 """Statistics about an architecture including supported and gap info.
347 Attributes:
348 architecture_id: The architecture identifier
349 is_supported: Whether TransformerLens supports this architecture
350 model_count: Number of models using this architecture
351 verified_count: Number of verified models (if supported)
352 example_models: Sample model IDs for this architecture
353 """
355 architecture_id: str
356 is_supported: bool
357 model_count: int
358 verified_count: int = 0
359 example_models: list[str] = field(default_factory=list)
361 def to_dict(self) -> dict:
362 """Convert to a JSON-serializable dictionary."""
363 return {
364 "architecture_id": self.architecture_id,
365 "is_supported": self.is_supported,
366 "model_count": self.model_count,
367 "verified_count": self.verified_count,
368 "example_models": self.example_models,
369 }
372@dataclass
373class ArchitectureAnalysis:
374 """Analysis result for prioritizing architecture support.
376 Attributes:
377 architecture_id: The architecture identifier
378 total_models: Total models using this architecture
379 total_downloads: Sum of downloads across all models
380 priority_score: Computed priority score for implementation
381 top_models: Most popular models for this architecture
382 """
384 architecture_id: str
385 total_models: int
386 total_downloads: int
387 priority_score: float
388 top_models: list[str] = field(default_factory=list)
390 def to_dict(self) -> dict:
391 """Convert to a JSON-serializable dictionary."""
392 return {
393 "architecture_id": self.architecture_id,
394 "total_models": self.total_models,
395 "total_downloads": self.total_downloads,
396 "priority_score": self.priority_score,
397 "top_models": self.top_models,
398 }