Coverage for transformer_lens/tools/model_registry/schemas.py: 96%
128 statements
« prev ^ index » next coverage.py v7.10.1, created at 2026-08-11 18:50 +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 phase1_score: Benchmark Phase 1 score (HF vs Bridge), 0-100 or None
68 phase2_score: Benchmark Phase 2 score (Bridge vs HT unprocessed), 0-100 or None
69 phase3_score: Benchmark Phase 3 score (Bridge vs HT processed), 0-100 or None
70 phase4_score: Benchmark Phase 4 score (generation / text quality), 0-100 or None
71 phase7_score: Benchmark Phase 7 score (multimodal), 0-100 or None
72 phase8_score: Benchmark Phase 8 score (audio), 0-100 or None
73 phase9_score: Benchmark Phase 9 score (vision), 0-100 or None
74 """
76 architecture_id: str
77 model_id: str
78 status: int = 0
79 verified_date: Optional[date] = None
80 metadata: Optional[ModelMetadata] = None
81 note: Optional[str] = None
82 phase1_score: Optional[float] = None
83 phase2_score: Optional[float] = None
84 phase3_score: Optional[float] = None
85 phase4_score: Optional[float] = None
86 phase7_score: Optional[float] = None
87 phase8_score: Optional[float] = None
88 phase9_score: Optional[float] = None
90 def to_dict(self) -> dict:
91 """Convert to a JSON-serializable dictionary."""
92 return {
93 "architecture_id": self.architecture_id,
94 "model_id": self.model_id,
95 "status": self.status,
96 "verified_date": self.verified_date.isoformat() if self.verified_date else None,
97 "metadata": self.metadata.to_dict() if self.metadata else None,
98 "note": self.note,
99 "phase1_score": self.phase1_score,
100 "phase2_score": self.phase2_score,
101 "phase3_score": self.phase3_score,
102 "phase4_score": self.phase4_score,
103 "phase7_score": self.phase7_score,
104 "phase8_score": self.phase8_score,
105 "phase9_score": self.phase9_score,
106 }
108 @classmethod
109 def from_dict(cls, data: dict) -> "ModelEntry":
110 """Create from a dictionary."""
111 verified_date = None
112 if data.get("verified_date"):
113 verified_date = date.fromisoformat(data["verified_date"])
114 metadata = None
115 if data.get("metadata"): 115 ↛ 116line 115 didn't jump to line 116 because the condition on line 115 was never true
116 metadata = ModelMetadata.from_dict(data["metadata"])
117 # Backwards compat: convert old "verified" bool to new "status" int
118 if "status" in data:
119 status = data["status"]
120 elif data.get("verified", False):
121 status = 1
122 else:
123 status = 0
124 return cls(
125 architecture_id=data["architecture_id"],
126 model_id=data["model_id"],
127 status=status,
128 verified_date=verified_date,
129 metadata=metadata,
130 note=data.get("note"),
131 phase1_score=data.get("phase1_score"),
132 phase2_score=data.get("phase2_score"),
133 phase3_score=data.get("phase3_score"),
134 phase4_score=data.get("phase4_score"),
135 phase7_score=data.get("phase7_score"),
136 phase8_score=data.get("phase8_score"),
137 phase9_score=data.get("phase9_score"),
138 )
141@dataclass
142class ArchitectureGap:
143 """An unsupported architecture with model count and relevancy metrics.
145 Attributes:
146 architecture_id: The architecture type not supported by TransformerLens
147 total_models: Number of models on HuggingFace using this architecture
148 sample_models: Top models by downloads for this architecture (up to 10)
149 total_downloads: Aggregate download count across all models of this architecture
150 min_param_count: Parameter count of the smallest model (None if unknown)
151 relevancy_score: Composite relevancy score (0-100), or None if not computed
152 """
154 architecture_id: str
155 total_models: int
156 sample_models: list[str] = field(default_factory=list)
157 total_downloads: int = 0
158 min_param_count: Optional[int] = None
159 relevancy_score: Optional[float] = None
161 def to_dict(self) -> dict:
162 """Convert to a JSON-serializable dictionary."""
163 d: dict = {
164 "architecture_id": self.architecture_id,
165 "total_models": self.total_models,
166 "total_downloads": self.total_downloads,
167 "min_param_count": self.min_param_count,
168 "relevancy_score": self.relevancy_score,
169 "sample_models": self.sample_models,
170 }
171 return d
173 @classmethod
174 def from_dict(cls, data: dict) -> "ArchitectureGap":
175 """Create from a dictionary."""
176 return cls(
177 architecture_id=data["architecture_id"],
178 total_models=data["total_models"],
179 sample_models=data.get("sample_models", []),
180 total_downloads=data.get("total_downloads", 0),
181 min_param_count=data.get("min_param_count"),
182 relevancy_score=data.get("relevancy_score"),
183 )
186@dataclass
187class ScanInfo:
188 """Metadata about a scraping run.
190 Attributes:
191 total_scanned: Total number of models scanned in this run
192 task_filter: HuggingFace task filter used (e.g., "text-generation")
193 scan_duration_seconds: How long the scan took in seconds (if available)
194 """
196 total_scanned: int
197 task_filter: str
198 scan_duration_seconds: Optional[float] = None
200 def to_dict(self) -> dict:
201 """Convert to a JSON-serializable dictionary."""
202 d: dict = {
203 "total_scanned": self.total_scanned,
204 "task_filter": self.task_filter,
205 }
206 if self.scan_duration_seconds is not None: 206 ↛ 207line 206 didn't jump to line 207 because the condition on line 206 was never true
207 d["scan_duration_seconds"] = self.scan_duration_seconds
208 return d
210 @classmethod
211 def from_dict(cls, data: dict) -> "ScanInfo":
212 """Create from a dictionary."""
213 return cls(
214 total_scanned=data["total_scanned"],
215 task_filter=data["task_filter"],
216 scan_duration_seconds=data.get("scan_duration_seconds"),
217 )
220@dataclass
221class SupportedModelsReport:
222 """Report containing all supported models.
224 Attributes:
225 generated_at: Date when this report was generated
226 scan_info: Metadata about the scraping run
227 total_architectures: Number of unique supported architectures
228 total_models: Total number of supported models
229 total_verified: Number of models that have been verified
230 models: List of all model entries
231 """
233 generated_at: date
234 total_models: int
235 models: list[ModelEntry]
236 scan_info: Optional[ScanInfo] = None
237 total_architectures: int = 0
238 total_verified: int = 0
239 # Structural-only (--no-hf-reference) passes; not counted as verified.
240 total_provisional: int = 0
242 def to_dict(self) -> dict:
243 """Convert to a JSON-serializable dictionary."""
244 d: dict = {
245 "generated_at": self.generated_at.isoformat(),
246 "scan_info": self.scan_info.to_dict() if self.scan_info else None,
247 "total_architectures": self.total_architectures,
248 "total_models": self.total_models,
249 "total_verified": self.total_verified,
250 "total_provisional": self.total_provisional,
251 "models": [m.to_dict() for m in self.models],
252 }
253 return d
255 @classmethod
256 def from_dict(cls, data: dict) -> "SupportedModelsReport":
257 """Create from a dictionary."""
258 scan_info = None
259 if data.get("scan_info"): 259 ↛ 261line 259 didn't jump to line 261 because the condition on line 259 was always true
260 scan_info = ScanInfo.from_dict(data["scan_info"])
261 return cls(
262 generated_at=date.fromisoformat(data["generated_at"]),
263 scan_info=scan_info,
264 total_architectures=data.get("total_architectures", 0),
265 total_models=data.get("total_models", len(data.get("models", []))),
266 total_verified=data.get("total_verified", 0),
267 total_provisional=data.get("total_provisional", 0),
268 models=[ModelEntry.from_dict(m) for m in data["models"]],
269 )
272@dataclass
273class ArchitectureGapsReport:
274 """Report containing unsupported architectures.
276 Attributes:
277 generated_at: Date when this report was generated
278 scan_info: Metadata about the scraping run
279 total_unsupported_architectures: Number of unsupported architectures
280 total_unsupported_models: Total models across all unsupported architectures
281 gaps: List of architecture gaps sorted by model count
282 """
284 generated_at: date
285 gaps: list[ArchitectureGap]
286 scan_info: Optional[ScanInfo] = None
287 total_unsupported_architectures: int = 0
288 total_unsupported_models: int = 0
290 def to_dict(self) -> dict:
291 """Convert to a JSON-serializable dictionary."""
292 return {
293 "generated_at": self.generated_at.isoformat(),
294 "scan_info": self.scan_info.to_dict() if self.scan_info else None,
295 "total_unsupported_architectures": self.total_unsupported_architectures,
296 "total_unsupported_models": self.total_unsupported_models,
297 "gaps": [g.to_dict() for g in self.gaps],
298 }
300 @classmethod
301 def from_dict(cls, data: dict) -> "ArchitectureGapsReport":
302 """Create from a dictionary."""
303 scan_info = None
304 if data.get("scan_info"): 304 ↛ 306line 304 didn't jump to line 306 because the condition on line 304 was always true
305 scan_info = ScanInfo.from_dict(data["scan_info"])
306 gaps = [ArchitectureGap.from_dict(g) for g in data["gaps"]]
307 return cls(
308 generated_at=date.fromisoformat(data["generated_at"]),
309 scan_info=scan_info,
310 total_unsupported_architectures=data.get(
311 "total_unsupported_architectures",
312 data.get("total_unsupported", len(gaps)),
313 ),
314 total_unsupported_models=data.get(
315 "total_unsupported_models",
316 sum(g.total_models for g in gaps),
317 ),
318 gaps=gaps,
319 )
322@dataclass
323class ArchitectureStats:
324 """Statistics about an architecture including supported and gap info.
326 Attributes:
327 architecture_id: The architecture identifier
328 is_supported: Whether TransformerLens supports this architecture
329 model_count: Number of models using this architecture
330 verified_count: Number of verified models (if supported)
331 example_models: Sample model IDs for this architecture
332 """
334 architecture_id: str
335 is_supported: bool
336 model_count: int
337 verified_count: int = 0
338 example_models: list[str] = field(default_factory=list)
340 def to_dict(self) -> dict:
341 """Convert to a JSON-serializable dictionary."""
342 return {
343 "architecture_id": self.architecture_id,
344 "is_supported": self.is_supported,
345 "model_count": self.model_count,
346 "verified_count": self.verified_count,
347 "example_models": self.example_models,
348 }
351@dataclass
352class ArchitectureAnalysis:
353 """Analysis result for prioritizing architecture support.
355 Attributes:
356 architecture_id: The architecture identifier
357 total_models: Total models using this architecture
358 total_downloads: Sum of downloads across all models
359 priority_score: Computed priority score for implementation
360 top_models: Most popular models for this architecture
361 """
363 architecture_id: str
364 total_models: int
365 total_downloads: int
366 priority_score: float
367 top_models: list[str] = field(default_factory=list)
369 def to_dict(self) -> dict:
370 """Convert to a JSON-serializable dictionary."""
371 return {
372 "architecture_id": self.architecture_id,
373 "total_models": self.total_models,
374 "total_downloads": self.total_downloads,
375 "priority_score": self.priority_score,
376 "top_models": self.top_models,
377 }