Coverage for transformer_lens/tools/model_registry/hf_scraper.py: 0%

1#!/usr/bin/env python3 

2"""HuggingFace model scraper for discovering compatible models. 

3 

4This module queries the HuggingFace Hub API to find ALL models and categorize 

5them by architecture - those supported by TransformerLens and those not yet supported. 

6 

7The scraper works by: 

81. Scanning ALL models for one HF task tag per run (paginated; default: text-generation) 

92. Extracting the architecture class from each model's config 

103. Categorizing models into supported vs unsupported based on TransformerLens adapters 

114. Building comprehensive lists for both categories 

12 

13Sequential runs MERGE (existing rows and gap entries are preserved), so a full 

14registry refresh is layered task passes: text-generation, then text2text-generation 

15(T5/mT5), then the vision passes image-classification and image-feature-extraction 

16(ViT/DeiT). 

17 

18Output format matches the schemas defined in schemas.py exactly, so the data 

19files can be loaded by api.py without any transformation. 

20 

21Usage: 

22 # Full scan of all HuggingFace models (recommended) 

23 python -m transformer_lens.tools.model_registry.hf_scraper --full-scan 

24 

25 # Targeted scrape: only models of a specific architecture 

26 python -m transformer_lens.tools.model_registry.hf_scraper \\ 

27 --architecture LlamaForCausalLM --full-scan 

28 

29 # Vision pass: layer image-task models onto a prior text-generation scan 

30 python -m transformer_lens.tools.model_registry.hf_scraper \\ 

31 --task image-classification --full-scan 

32 

33 # Quick scan (top N models by downloads) 

34 python -m transformer_lens.tools.model_registry.hf_scraper --limit 10000 

35 

36 # Output to custom directory 

37 python -m transformer_lens.tools.model_registry.hf_scraper --full-scan --output data/ 

38""" 

39 

40import argparse 

41import json 

42import logging 

43import time 

44from datetime import date, datetime 

45from pathlib import Path 

46from typing import Optional 

47 

48from transformer_lens.benchmarks.text_quality_profiles import ( 

49 ARCHITECTURE_PROFILE_KINDS, 

50 MODEL_PROFILE_OVERRIDES, 

51 HFSignals, 

52 extract_languages, 

53 is_default_profile, 

54 profile_from_hf_signals, 

55 resolve_profile, 

56) 

57 

58from . import HF_SUPPORTED_ARCHITECTURES 

59from .registry_io import ( 

60 STATUS_UNVERIFIED, 

61 is_quantized_model, 

62 recompute_registry_totals, 

63) 

64 

65logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s") 

66logger = logging.getLogger(__name__) 

67 

68 

69def _extract_architecture(model_info) -> Optional[str]: # type: ignore[no-untyped-def] 

70 """Extract the primary architecture class from a model's inline config. 

71 

72 Args: 

73 model_info: ModelInfo object from list_models(expand=['config']) 

74 

75 Returns: 

76 Architecture class name or None if not found 

77 """ 

78 config = model_info.config 

79 if config and isinstance(config, dict): 

80 archs = config.get("architectures", []) 

81 if archs: 

82 return archs[0] 

83 return None 

84 

85 

86def _extract_param_count(model_info) -> Optional[int]: # type: ignore[no-untyped-def] 

87 """Extract parameter count from a model's safetensors metadata or config. 

88 

89 Tries safetensors metadata first (most reliable), then falls back to 

90 config fields like num_parameters or n_params. 

91 

92 Args: 

93 model_info: ModelInfo object from list_models(expand=['config', 'safetensors']) 

94 

95 Returns: 

96 Total parameter count or None if not available 

97 """ 

98 # Try safetensors metadata (most reliable source) 

99 safetensors = getattr(model_info, "safetensors", None) 

100 if safetensors and isinstance(safetensors, dict): 

101 # safetensors metadata has a 'total' field with total parameter count 

102 total = safetensors.get("total") 

103 if total is not None: 

104 try: 

105 return int(total) 

106 except (ValueError, TypeError): 

107 pass 

108 # Some models store it under 'parameters' -> 'total' 

109 params = safetensors.get("parameters") 

110 if params and isinstance(params, dict): 

111 total = params.get("total") 

112 if total is not None: 

113 try: 

114 return int(total) 

115 except (ValueError, TypeError): 

116 pass 

117 

118 # Fall back to config fields 

119 config = getattr(model_info, "config", None) 

120 if config and isinstance(config, dict): 

121 for key in ("num_parameters", "n_params", "num_params"): 

122 val = config.get(key) 

123 if val is not None: 

124 try: 

125 return int(val) 

126 except (ValueError, TypeError): 

127 pass 

128 

129 return None 

130 

131 

132def _load_existing_models(output_dir: Path) -> tuple[set[str], list[dict]]: 

133 """Load model IDs and data already in supported_models.json. 

134 

135 Args: 

136 output_dir: Directory containing the data files 

137 

138 Returns: 

139 Tuple of (set of existing model IDs, list of existing model dicts) 

140 """ 

141 existing_ids: set[str] = set() 

142 existing_models: list[dict] = [] 

143 supported_path = output_dir / "supported_models.json" 

144 

145 if supported_path.exists(): 

146 try: 

147 with open(supported_path) as f: 

148 data = json.load(f) 

149 for model in data.get("models", []): 

150 if "model_id" in model: 

151 existing_ids.add(model["model_id"]) 

152 existing_models.append(model) 

153 logger.info(f"Loaded {len(existing_ids)} existing models from {supported_path}") 

154 except (json.JSONDecodeError, KeyError) as e: 

155 logger.warning(f"Could not load existing models: {e}") 

156 

157 return existing_ids, existing_models 

158 

159 

160def _load_existing_gaps(output_dir: Path) -> dict[str, dict]: 

161 """Load existing per-architecture gap entries keyed by architecture_id. 

162 

163 Lets a new scrape merge instead of overwrite — without this, each later pass of a 

164 sequential multi-task run (text-generation, then text2text-generation, then the 

165 image tasks) wipes the earlier passes' gap data. 

166 """ 

167 gaps_path = output_dir / "architecture_gaps.json" 

168 by_arch: dict[str, dict] = {} 

169 if not gaps_path.exists(): 

170 return by_arch 

171 try: 

172 data = json.loads(gaps_path.read_text()) 

173 except (json.JSONDecodeError, OSError) as e: 

174 logger.warning(f"Could not load existing gaps: {e}") 

175 return by_arch 

176 for entry in data.get("gaps", []): 

177 if isinstance(entry, dict) and "architecture_id" in entry: 

178 by_arch[entry["architecture_id"]] = entry 

179 if by_arch: 

180 logger.info(f"Loaded {len(by_arch)} existing architecture gaps from {gaps_path}") 

181 return by_arch 

182 

183 

184def _extract_profile_signals(model_info) -> HFSignals: # type: ignore[no-untyped-def] 

185 """Distill pipeline_tag/tags/cardData off a listing payload (no extra request). 

186 

187 Args: 

188 model_info: ModelInfo object from list_models(expand=[..., 'pipeline_tag', 

189 'tags', 'cardData']) 

190 """ 

191 pipeline_tag = getattr(model_info, "pipeline_tag", None) 

192 tags = tuple(getattr(model_info, "tags", None) or []) 

193 card_data = getattr(model_info, "card_data", None) 

194 card_language = getattr(card_data, "language", None) if card_data is not None else None 

195 if card_language is None and card_data is not None and hasattr(card_data, "get"): 

196 card_language = card_data.get("language") 

197 languages = extract_languages(card_language, tags) 

198 return HFSignals(pipeline_tag=pipeline_tag, languages=languages, tags=tags) 

199 

200 

201def _build_model_entry( 

202 model_id: str, architecture_id: str, signals: Optional[HFSignals] = None 

203) -> dict: 

204 """Build a model entry dict matching the ModelEntry schema. 

205 

206 ``signals``, when given, resolves and stores a sparse ``prompt_profile`` key 

207 (omitted when it's just the default) and warns on tag/curation disagreement 

208 — the warning is how curation gaps (missing override/architecture rule) surface. 

209 """ 

210 entry = { 

211 "architecture_id": architecture_id, 

212 "model_id": model_id, 

213 "status": STATUS_UNVERIFIED, 

214 "verified_date": None, 

215 "metadata": None, 

216 "note": None, 

217 "phase1_score": None, 

218 "phase2_score": None, 

219 "phase3_score": None, 

220 "phase4_score": None, 

221 "phase7_score": None, 

222 "phase8_score": None, 

223 "phase9_score": None, 

224 } 

225 if signals is not None: 

226 hinted = profile_from_hf_signals(model_id, architecture_id, signals) 

227 resolved = resolve_profile(model_id, architecture_id, signals=signals) 

228 deliberately_curated = ( 

229 model_id in MODEL_PROFILE_OVERRIDES or architecture_id in ARCHITECTURE_PROFILE_KINDS 

230 ) 

231 if hinted is not None and hinted.kind != resolved.kind and not deliberately_curated: 

232 # A disagreement nothing deliberate explains is a curation gap. 

233 logger.warning( 

234 f"Profile mismatch for {model_id} ({architecture_id}): Hub tags say " 

235 f"{hinted.kind!r}, curation resolves {resolved.kind!r}" 

236 ) 

237 if not is_default_profile(resolved): 

238 # Keep key position consistent with ModelEntry.to_dict (after note). 

239 items = list(entry.items()) 

240 items.insert([k for k, _ in items].index("note") + 1, ("prompt_profile", str(resolved))) 

241 entry.clear() 

242 entry.update(items) 

243 return entry 

244 

245 

246def _canonical_author_sweep( 

247 api, # type: ignore[no-untyped-def] 

248 supported_models: list[dict], 

249 seen_models: set[str], 

250 architecture: Optional[str] = None, 

251 refresh_profiles: bool = False, 

252) -> int: 

253 """Admit canonical-org supported-arch models regardless of downloads. Returns count added. 

254 

255 When ``architecture`` is set, only sweep authors canonical for that architecture and 

256 only admit models whose extracted arch matches it. 

257 """ 

258 from . import CANONICAL_AUTHORS_BY_ARCH, HF_SUPPORTED_ARCHITECTURES 

259 

260 # Same author can be canonical for multiple archs (e.g. google: T5 + MT5 + Gemma). 

261 authors_to_archs: dict[str, set[str]] = {} 

262 for arch, authors in CANONICAL_AUTHORS_BY_ARCH.items(): 

263 for author in authors: 

264 authors_to_archs.setdefault(author, set()).add(arch) 

265 

266 added = 0 

267 for author, expected_archs in sorted(authors_to_archs.items()): 

268 if architecture is not None and architecture not in expected_archs: 

269 continue 

270 try: 

271 models_iter = api.list_models( 

272 author=author, 

273 expand=["config", "safetensors", "pipeline_tag", "tags", "cardData"], 

274 ) 

275 except Exception as exc: # pragma: no cover — network/transient 

276 logger.warning(f"Canonical sweep: list_models(author={author!r}) failed: {exc}") 

277 continue 

278 

279 # Iterate paginated results; a single timeout shouldn't lose every prior author. 

280 existing_by_id = {m["model_id"]: m for m in supported_models} if refresh_profiles else {} 

281 try: 

282 for model in models_iter: 

283 if model.id in seen_models: 

284 # Below-threshold canonical models are reachable only here; 

285 # the main scan's backfill never sees them. 

286 if refresh_profiles: 

287 existing_entry = existing_by_id.get(model.id) 

288 if existing_entry is not None and "prompt_profile" not in existing_entry: 

289 resolved = resolve_profile( 

290 model.id, 

291 existing_entry.get("architecture_id"), 

292 signals=_extract_profile_signals(model), 

293 ) 

294 if not is_default_profile(resolved): 

295 existing_entry["prompt_profile"] = str(resolved) 

296 continue 

297 if is_quantized_model(model.id): 

298 continue 

299 model_arch: Optional[str] = _extract_architecture(model) 

300 if model_arch is None or model_arch not in HF_SUPPORTED_ARCHITECTURES: 

301 continue 

302 if architecture is not None and model_arch != architecture: 

303 continue 

304 # Reject e.g. mistralai's non-Mistral checkpoints. 

305 if model_arch not in expected_archs: 

306 continue 

307 signals = _extract_profile_signals(model) 

308 supported_models.append(_build_model_entry(model.id, model_arch, signals)) 

309 seen_models.add(model.id) 

310 added += 1 

311 logger.info(f"Canonical sweep added: {model.id} ({model_arch})") 

312 except Exception as exc: # pragma: no cover — network/transient 

313 logger.warning( 

314 f"Canonical sweep: pagination for {author!r} failed mid-iteration: {exc}" 

315 ) 

316 continue 

317 return added 

318 

319 

320def scrape_all_models( 

321 output_dir: Path, 

322 max_models: Optional[int] = None, 

323 task: str = "text-generation", 

324 batch_size: int = 1000, 

325 checkpoint_interval: int = 5000, 

326 min_downloads: int = 500, 

327 canonical_sweep: bool = True, 

328 architecture: Optional[str] = None, 

329 refresh_profiles: bool = False, 

330) -> tuple[dict, dict]: 

331 """Scrape ALL models from HuggingFace and categorize by architecture. 

332 

333 This is the comprehensive scraper that: 

334 1. Loads existing models from supported_models.json to preserve them 

335 2. Skips models already in the JSON (only scans new models) 

336 3. Iterates through ALL models for a given task 

337 4. Fetches the architecture from each model's config 

338 5. Categorizes into supported vs unsupported 

339 6. Saves checkpoints periodically for long runs 

340 

341 Output format matches schemas.py exactly (SupportedModelsReport and 

342 ArchitectureGapsReport). 

343 

344 Args: 

345 output_dir: Directory to write JSON data files 

346 max_models: Maximum NEW models to scan (None = unlimited/all) 

347 task: HuggingFace task tag to filter by (default: text-generation). Sequential 

348 runs merge, so layer extra passes for non-text architectures: 

349 ``text2text-generation`` (T5/mT5), ``image-classification`` and 

350 ``image-feature-extraction`` (ViT/DeiT). 

351 batch_size: Log progress every N models 

352 checkpoint_interval: Save checkpoint every N models 

353 min_downloads: Minimum download count to include a model (default: 500) 

354 canonical_sweep: If True, run the post-scrape pass that admits canonical-org models 

355 below the download threshold (default: True). 

356 architecture: If set, only include models whose ``config.architectures[0]`` matches 

357 this class (e.g. ``"LlamaForCausalLM"``). Applies to both the main scan and 

358 the canonical-author sweep. Useful for populating the registry after adding 

359 a single new adapter without rescanning every architecture. 

360 refresh_profiles: If True, backfill a missing ``prompt_profile`` key onto 

361 already-seen registry entries using the listing payload already in hand — no 

362 extra requests (default: False). 

363 

364 Returns: 

365 Tuple of (supported_models_dict, architecture_gaps_dict) 

366 """ 

367 try: 

368 from huggingface_hub import HfApi 

369 except ImportError: 

370 raise ImportError( 

371 "huggingface_hub is required for scraping. " 

372 "Install it with: pip install huggingface_hub" 

373 ) 

374 

375 from transformer_lens.utilities.hf_utils import get_hf_token 

376 

377 api = HfApi(token=get_hf_token()) 

378 output_dir = Path(output_dir) 

379 output_dir.mkdir(parents=True, exist_ok=True) 

380 

381 # Load existing models from supported_models.json 

382 existing_model_ids, existing_models = _load_existing_models(output_dir) 

383 

384 # Track all models by architecture (start with existing models) 

385 supported_models: list[dict] = list(existing_models) # Preserve existing 

386 # Same dict objects as supported_models — mutating via this index (--refresh-profiles) 

387 # is reflected in the final write. 

388 existing_by_id: dict[str, dict] = {m["model_id"]: m for m in supported_models} 

389 unsupported_arch_counts: dict[str, int] = {} # arch -> count 

390 unsupported_arch_samples: dict[str, list[str]] = {} # arch -> top model IDs 

391 unsupported_arch_downloads: dict[str, int] = {} # arch -> total downloads 

392 unsupported_arch_min_params: dict[str, int] = {} # arch -> smallest param count 

393 max_samples = 10 # Keep top N sample models per unsupported architecture 

394 

395 scanned = 0 

396 skipped = 0 

397 new_supported = 0 

398 errors = 0 

399 start_time = time.time() 

400 

401 # Check for existing checkpoint to resume from 

402 checkpoint_path = output_dir / "scrape_checkpoint.json" 

403 seen_models: set[str] = set(existing_model_ids) # Include existing as "seen" 

404 

405 # When `architecture` is set AND we have canonical orgs for it, skip the global 

406 # text-generation scan: the canonical sweep already exhausts those orgs and is 

407 # exact (`author=` is a server-side filter). The main scan would only add 

408 # community fine-tunes of that arch, which are rarely worth verifying. For 

409 # archs with no canonical orgs registered, fall back to the main scan + 

410 # client-side filter. 

411 from . import CANONICAL_AUTHORS_BY_ARCH 

412 

413 skip_main_scan = architecture is not None and architecture in CANONICAL_AUTHORS_BY_ARCH 

414 if skip_main_scan: 

415 assert architecture is not None # narrowed by skip_main_scan 

416 logger.info( 

417 f"Targeted scrape for architecture={architecture!r}: skipping the global " 

418 f"'{task}' scan; relying on canonical-author sweep over " 

419 f"{sorted(CANONICAL_AUTHORS_BY_ARCH[architecture])}." 

420 ) 

421 if not canonical_sweep: 

422 logger.warning( 

423 "skip_main_scan is set but --no-canonical-sweep was passed. No HF " 

424 "queries will run. Re-run without --no-canonical-sweep to actually " 

425 "discover models." 

426 ) 

427 

428 if not skip_main_scan and checkpoint_path.exists(): 

429 logger.info(f"Found checkpoint at {checkpoint_path}, loading...") 

430 with open(checkpoint_path) as f: 

431 checkpoint = json.load(f) 

432 # Merge checkpoint data with existing 

433 checkpoint_supported = checkpoint.get("supported_models", []) 

434 for model in checkpoint_supported: 

435 if model["model_id"] not in existing_model_ids: 

436 supported_models.append(model) 

437 existing_model_ids.add(model["model_id"]) 

438 unsupported_arch_counts = checkpoint.get("unsupported_arch_counts", {}) 

439 unsupported_arch_samples = checkpoint.get("unsupported_arch_samples", {}) 

440 unsupported_arch_downloads = checkpoint.get("unsupported_arch_downloads", {}) 

441 unsupported_arch_min_params = checkpoint.get("unsupported_arch_min_params", {}) 

442 seen_models.update(checkpoint.get("seen_models", [])) 

443 scanned = checkpoint.get("scanned", 0) 

444 skipped = checkpoint.get("skipped", 0) 

445 logger.info(f"Resumed from checkpoint: {scanned} models already scanned") 

446 

447 if not skip_main_scan: 

448 logger.info(f"Starting comprehensive HuggingFace scan for task='{task}'...") 

449 logger.info(f"Skipping {len(existing_model_ids)} models already in supported_models.json") 

450 logger.info(f"Supported architectures: {len(HF_SUPPORTED_ARCHITECTURES)}") 

451 logger.info(f"Minimum downloads threshold: {min_downloads:,}") 

452 if max_models: 

453 logger.info(f"Will scan up to {max_models} NEW models") 

454 else: 

455 logger.info("Will scan ALL new models (this may take a while)") 

456 

457 try: 

458 # Use expand=['config', 'safetensors', 'pipeline_tag', 'tags', 'cardData'] to get 

459 # architecture, parameter count, and prompt-profile signals inline with the 

460 # listing, avoiding per-model API calls. With ~1000 models per page, a full 

461 # scan of 200K+ models needs only ~200 paginated requests (well within the 

462 # 1000 req / 5 min limit). 

463 # Use ``filter`` rather than ``pipeline_tag`` (the query param) so 

464 # encoder-decoder models are discoverable: HF assigns T5/mT5 a primary 

465 # pipeline_tag of "translation" (or None for mT5) and only lists 

466 # "text2text-generation" in the broader tag list. ``filter`` matches against 

467 # tags, ``pipeline_tag`` only against the canonical primary tag. The 

468 # expanded ``pipeline_tag`` *field* below is a different thing — it's per-model 

469 # metadata fed to profile_from_hf_signals, not a query filter. 

470 list_kwargs: dict = { 

471 "filter": task, 

472 "sort": "downloads", 

473 "expand": ["config", "safetensors", "pipeline_tag", "tags", "cardData"], 

474 } 

475 if max_models is not None: 

476 list_kwargs["limit"] = max_models + len(seen_models) 

477 

478 # Retry loop: if we hit a 429 mid-pagination, save checkpoint, wait, 

479 # and restart iteration. Already-seen models are skipped automatically. 

480 max_retries = 10 

481 for attempt in range(max_retries + 1): 

482 if skip_main_scan: 

483 # Targeted scrape with canonical orgs available — the sweep below is 

484 # exhaustive within those orgs and exact (server-side `author=`), so 

485 # the global text-generation pagination would only add community 

486 # fine-tunes for the same arch. 

487 break 

488 try: 

489 for model in api.list_models(**list_kwargs): 

490 # Skip if already in our JSON or processed in this run 

491 if model.id in seen_models: 

492 skipped += 1 

493 if refresh_profiles: 

494 existing_entry = existing_by_id.get(model.id) 

495 if ( 

496 existing_entry is not None 

497 and "prompt_profile" not in existing_entry 

498 ): 

499 resolved = resolve_profile( 

500 model.id, 

501 existing_entry.get("architecture_id"), 

502 signals=_extract_profile_signals(model), 

503 ) 

504 if not is_default_profile(resolved): 

505 existing_entry["prompt_profile"] = str(resolved) 

506 continue 

507 

508 # Filter by minimum download count. Since results are sorted 

509 # by downloads descending, once we drop below the threshold 

510 # all remaining models will also be below it. 

511 downloads = getattr(model, "downloads", None) or 0 

512 if downloads < min_downloads: 

513 logger.info( 

514 f"Reached download threshold ({downloads:,} < " 

515 f"{min_downloads:,}) after {scanned} models. " 

516 f"Stopping scan." 

517 ) 

518 break 

519 

520 scanned += 1 

521 seen_models.add(model.id) 

522 

523 if max_models and scanned > max_models: 

524 break 

525 

526 # Skip quantized models (AWQ, GPTQ, GGUF, bnb, FP8, etc.) 

527 # TransformerLens requires full-precision weights. 

528 if is_quantized_model(model.id): 

529 continue 

530 

531 # Extract architecture from inline config (no extra API call) 

532 arch = _extract_architecture(model) 

533 

534 # Targeted scrape: drop everything that isn't the requested arch. 

535 # Applied before classification so the unsupported counters reflect 

536 # only the architecture under inspection. 

537 if architecture is not None and arch != architecture: 

538 continue 

539 

540 if arch is None: 

541 errors += 1 

542 elif arch in HF_SUPPORTED_ARCHITECTURES: 

543 signals = _extract_profile_signals(model) 

544 supported_models.append(_build_model_entry(model.id, arch, signals)) 

545 new_supported += 1 

546 else: 

547 unsupported_arch_counts[arch] = unsupported_arch_counts.get(arch, 0) + 1 

548 # Track top models per arch (sorted by downloads since list is sorted) 

549 samples = unsupported_arch_samples.setdefault(arch, []) 

550 if len(samples) < max_samples: 

551 samples.append(model.id) 

552 # Accumulate downloads for relevancy scoring 

553 unsupported_arch_downloads[arch] = ( 

554 unsupported_arch_downloads.get(arch, 0) + downloads 

555 ) 

556 # Track smallest model per arch for benchmarkability 

557 param_count = _extract_param_count(model) 

558 if param_count is not None: 

559 current_min = unsupported_arch_min_params.get(arch) 

560 if current_min is None or param_count < current_min: 

561 unsupported_arch_min_params[arch] = param_count 

562 

563 # Progress logging 

564 if scanned % batch_size == 0: 

565 elapsed = time.time() - start_time 

566 rate = scanned / elapsed if elapsed > 0 else 0 

567 logger.info( 

568 f"Scanned {scanned} new | " 

569 f"Skipped {skipped} existing | " 

570 f"New supported: {new_supported} | " 

571 f"Total supported: {len(supported_models)} | " 

572 f"Unsupported archs: {len(unsupported_arch_counts)} | " 

573 f"Errors: {errors} | " 

574 f"Rate: {rate:.1f}/s" 

575 ) 

576 

577 # Save checkpoint periodically 

578 if scanned % checkpoint_interval == 0: 

579 _save_checkpoint( 

580 checkpoint_path, 

581 supported_models, 

582 unsupported_arch_counts, 

583 unsupported_arch_samples, 

584 list(seen_models), 

585 scanned, 

586 skipped, 

587 unsupported_arch_downloads, 

588 unsupported_arch_min_params, 

589 ) 

590 logger.info(f"Saved checkpoint at {scanned} models") 

591 

592 break # Iteration completed successfully, exit retry loop 

593 

594 except Exception as exc: 

595 if "429" in str(exc) and attempt < max_retries: 

596 wait = min(10 * (attempt + 1), 60) 

597 logger.warning( 

598 f"Rate limited (429). Saving checkpoint and waiting {wait}s " 

599 f"before retry ({attempt + 1}/{max_retries})..." 

600 ) 

601 _save_checkpoint( 

602 checkpoint_path, 

603 supported_models, 

604 unsupported_arch_counts, 

605 unsupported_arch_samples, 

606 list(seen_models), 

607 scanned, 

608 skipped, 

609 unsupported_arch_downloads, 

610 unsupported_arch_min_params, 

611 ) 

612 time.sleep(wait) 

613 skipped = 0 # Reset skip counter for restart 

614 else: 

615 raise 

616 

617 except KeyboardInterrupt: 

618 logger.warning("Interrupted! Saving checkpoint...") 

619 _save_checkpoint( 

620 checkpoint_path, 

621 supported_models, 

622 unsupported_arch_counts, 

623 unsupported_arch_samples, 

624 list(seen_models), 

625 scanned, 

626 skipped, 

627 unsupported_arch_downloads, 

628 unsupported_arch_min_params, 

629 ) 

630 raise 

631 except Exception as e: 

632 logger.error(f"Error during scan: {e}") 

633 _save_checkpoint( 

634 checkpoint_path, 

635 supported_models, 

636 unsupported_arch_counts, 

637 unsupported_arch_samples, 

638 list(seen_models), 

639 scanned, 

640 skipped, 

641 unsupported_arch_downloads, 

642 unsupported_arch_min_params, 

643 ) 

644 raise 

645 

646 if canonical_sweep: 

647 logger.info("\nRunning canonical-author sweep (bypasses download threshold)...") 

648 # Don't lose the main-scan registry on a sweep-time failure. 

649 try: 

650 canonical_added = _canonical_author_sweep( 

651 api, 

652 supported_models, 

653 seen_models, 

654 architecture=architecture, 

655 refresh_profiles=refresh_profiles, 

656 ) 

657 new_supported += canonical_added 

658 logger.info(f"Canonical sweep added {canonical_added} models.") 

659 except Exception as exc: 

660 logger.warning(f"Canonical sweep aborted: {exc}. Main-scan results preserved.") 

661 

662 # Build final reports (matching schemas.py exactly) 

663 elapsed = time.time() - start_time 

664 logger.info(f"\nScan complete in {elapsed:.1f}s") 

665 logger.info(f"New models scanned: {scanned}") 

666 logger.info(f"Existing models skipped: {skipped}") 

667 logger.info(f"New supported models found: {new_supported}") 

668 logger.info(f"Total supported models: {len(supported_models)}") 

669 logger.info(f"Unsupported architectures found: {len(unsupported_arch_counts)}") 

670 

671 # Build scan info (shared by both reports) 

672 scan_info = { 

673 "total_scanned": scanned, 

674 "task_filter": task, 

675 "min_downloads": min_downloads, 

676 "scan_duration_seconds": round(elapsed, 1), 

677 } 

678 

679 # Build supported models report dict (for return value) 

680 registry_totals = recompute_registry_totals(supported_models) 

681 supported_report = { 

682 "generated_at": date.today().isoformat(), 

683 "scan_info": scan_info, 

684 **registry_totals, 

685 "models": supported_models, 

686 } 

687 

688 # Write supported models (single file) 

689 with open(output_dir / "supported_models.json", "w") as f: 

690 json.dump(supported_report, f, indent=2) 

691 f.write("\n") 

692 logger.info(f"Wrote {len(supported_models)} supported models to supported_models.json") 

693 

694 # Build architecture gaps report (matches ArchitectureGapsReport schema) 

695 # Include download and param count data, then compute relevancy scores 

696 from transformer_lens.tools.model_registry.relevancy import compute_scores_for_gaps 

697 

698 gaps: list[dict] = [ 

699 { 

700 "architecture_id": arch, 

701 "total_models": count, 

702 "total_downloads": unsupported_arch_downloads.get(arch, 0), 

703 "min_param_count": unsupported_arch_min_params.get(arch), 

704 "sample_models": unsupported_arch_samples.get(arch, []), 

705 } 

706 for arch, count in unsupported_arch_counts.items() 

707 ] 

708 

709 # Merge with gaps from prior scrapes so sequential task passes (text-generation 

710 # + text2text-generation + image tasks) don't lose earlier data. For overlapping 

711 # architectures, sum counts/downloads, take the smaller min_param_count, and 

712 # union sample_models (capped at 10). 

713 existing_gaps = _load_existing_gaps(output_dir) 

714 if existing_gaps: 

715 new_by_arch = {g["architecture_id"]: g for g in gaps} 

716 merged: list[dict] = [] 

717 for arch in set(existing_gaps) | set(new_by_arch): 

718 o = existing_gaps.get(arch) 

719 n = new_by_arch.get(arch) 

720 if o is None and n is not None: 

721 merged.append(n) 

722 continue 

723 if n is None and o is not None: 

724 merged.append(o) 

725 continue 

726 assert o is not None and n is not None 

727 # Both present: combine counts/downloads, dedupe samples (cap 10). 

728 merged_samples: list[str] = [] 

729 seen_samples: set[str] = set() 

730 for s in o.get("sample_models", []) + n.get("sample_models", []): 

731 if s not in seen_samples: 

732 merged_samples.append(s) 

733 seen_samples.add(s) 

734 if len(merged_samples) >= 10: 

735 break 

736 min_p = [ 

737 p for p in (o.get("min_param_count"), n.get("min_param_count")) if p is not None 

738 ] 

739 merged.append( 

740 { 

741 "architecture_id": arch, 

742 "total_models": o["total_models"] + n["total_models"], 

743 "total_downloads": o["total_downloads"] + n["total_downloads"], 

744 "min_param_count": min(min_p) if min_p else None, 

745 "sample_models": merged_samples, 

746 } 

747 ) 

748 gaps = merged 

749 

750 # Compute relevancy scores and sort by score descending 

751 compute_scores_for_gaps(gaps) 

752 

753 # Guard the load-bearing invariant: each architecture appears at most once in 

754 # the gaps list. The merge above produces unique-by-arch entries by 

755 # construction, but the report header reads from this list — so an explicit 

756 # dedup keeps the header consistent if the merge ever drifts. 

757 seen_archs: set[str] = set() 

758 deduped: list[dict] = [] 

759 for g in gaps: 

760 arch_id = g["architecture_id"] 

761 if arch_id in HF_SUPPORTED_ARCHITECTURES: 

762 # Gained an adapter since a prior scrape; the merge above carries the 

763 # stale entry forward, so drop it here — it's no longer a gap. 

764 continue 

765 if arch_id in seen_archs: 

766 logger.warning(f"Dropping duplicate gap entry for architecture {arch_id!r}") 

767 continue 

768 seen_archs.add(arch_id) 

769 deduped.append(g) 

770 gaps = deduped 

771 

772 gaps_report = { 

773 "generated_at": date.today().isoformat(), 

774 "scan_info": scan_info, 

775 "total_unsupported_architectures": len(gaps), 

776 # Sum from the merged+deduped list so the header stays consistent with 

777 # its own gaps[*].total_models — the prior `sum(unsupported_arch_counts...)` 

778 # only reflected this run, while the list also carried prior-scrape data. 

779 "total_unsupported_models": sum(g["total_models"] for g in gaps), 

780 "gaps": gaps, 

781 } 

782 

783 gaps_path = output_dir / "architecture_gaps.json" 

784 with open(gaps_path, "w") as f: 

785 json.dump(gaps_report, f, indent=2) 

786 logger.info(f"Wrote {len(gaps)} architecture gaps to {gaps_path}") 

787 

788 # Write verification history placeholder (single file) 

789 verification_path = output_dir / "verification_history.json" 

790 if not verification_path.exists(): 

791 with open(verification_path, "w") as f: 

792 json.dump({"last_updated": None, "records": []}, f, indent=2) 

793 f.write("\n") 

794 

795 # Clean up checkpoint on successful completion 

796 if checkpoint_path.exists(): 

797 checkpoint_path.unlink() 

798 logger.info("Removed checkpoint file (scan complete)") 

799 

800 # Print summary 

801 logger.info("\n" + "=" * 70) 

802 logger.info("SCAN SUMMARY") 

803 logger.info("=" * 70) 

804 logger.info(f"Total models scanned: {scanned}") 

805 logger.info(f"\nSUPPORTED ARCHITECTURES ({registry_totals['total_architectures']}):") 

806 

807 # Count models per supported architecture 

808 supported_arch_counts: dict[str, int] = {} 

809 for model in supported_models: 

810 arch = model["architecture_id"] 

811 supported_arch_counts[arch] = supported_arch_counts.get(arch, 0) + 1 

812 

813 for arch, count in sorted(supported_arch_counts.items(), key=lambda x: -x[1]): 

814 logger.info(f" {arch}: {count} models") 

815 

816 logger.info(f"\nTOP 20 UNSUPPORTED ARCHITECTURES by relevancy (of {len(gaps)}):") 

817 for gap in gaps[:20]: 

818 score = gap.get("relevancy_score", 0) 

819 logger.info( 

820 f" {gap['architecture_id']}: " 

821 f"score={score:.1f}, " 

822 f"{gap['total_models']} models, " 

823 f"{gap.get('total_downloads', 0):,} downloads" 

824 ) 

825 

826 if len(gaps) > 20: 

827 remaining = sum(g["total_models"] for g in gaps[20:]) 

828 logger.info(f" ... and {len(gaps) - 20} more architectures ({remaining} models)") 

829 

830 logger.info("=" * 70) 

831 

832 return supported_report, gaps_report 

833 

834 

835def _save_checkpoint(