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