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 text-generation models on HuggingFace (paginated) 

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 

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

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

15 

16Usage: 

17 # Full scan of all HuggingFace models (recommended) 

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

19 

20 # Targeted scrape: only models of a specific architecture 

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

22 --architecture LlamaForCausalLM --full-scan 

23 

24 # Quick scan (top N models by downloads) 

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

26 

27 # Output to custom directory 

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

29""" 

30 

31import argparse 

32import json 

33import logging 

34import time 

35from datetime import date, datetime 

36from pathlib import Path 

37from typing import Optional 

38 

39from transformer_lens.benchmarks.text_quality_profiles import ( 

40 ARCHITECTURE_PROFILE_KINDS, 

41 MODEL_PROFILE_OVERRIDES, 

42 HFSignals, 

43 extract_languages, 

44 is_default_profile, 

45 profile_from_hf_signals, 

46 resolve_profile, 

47) 

48 

49from . import HF_SUPPORTED_ARCHITECTURES 

50from .registry_io import is_quantized_model 

51 

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

53logger = logging.getLogger(__name__) 

54 

55 

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

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

58 

59 Args: 

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

61 

62 Returns: 

63 Architecture class name or None if not found 

64 """ 

65 config = model_info.config 

66 if config and isinstance(config, dict): 

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

68 if archs: 

69 return archs[0] 

70 return None 

71 

72 

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

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

75 

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

77 config fields like num_parameters or n_params. 

78 

79 Args: 

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

81 

82 Returns: 

83 Total parameter count or None if not available 

84 """ 

85 # Try safetensors metadata (most reliable source) 

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

87 if safetensors and isinstance(safetensors, dict): 

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

89 total = safetensors.get("total") 

90 if total is not None: 

91 try: 

92 return int(total) 

93 except (ValueError, TypeError): 

94 pass 

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

96 params = safetensors.get("parameters") 

97 if params and isinstance(params, dict): 

98 total = params.get("total") 

99 if total is not None: 

100 try: 

101 return int(total) 

102 except (ValueError, TypeError): 

103 pass 

104 

105 # Fall back to config fields 

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

107 if config and isinstance(config, dict): 

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

109 val = config.get(key) 

110 if val is not None: 

111 try: 

112 return int(val) 

113 except (ValueError, TypeError): 

114 pass 

115 

116 return None 

117 

118 

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

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

121 

122 Args: 

123 output_dir: Directory containing the data files 

124 

125 Returns: 

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

127 """ 

128 existing_ids: set[str] = set() 

129 existing_models: list[dict] = [] 

130 supported_path = output_dir / "supported_models.json" 

131 

132 if supported_path.exists(): 

133 try: 

134 with open(supported_path) as f: 

135 data = json.load(f) 

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

137 if "model_id" in model: 

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

139 existing_models.append(model) 

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

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

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

143 

144 return existing_ids, existing_models 

145 

146 

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

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

149 

150 Lets a new scrape merge instead of overwrite — without this, the second of two 

151 sequential scrapes (e.g. text-generation then text2text-generation) wipes the 

152 first run's gap data. 

153 """ 

154 gaps_path = output_dir / "architecture_gaps.json" 

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

156 if not gaps_path.exists(): 

157 return by_arch 

158 try: 

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

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

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

162 return by_arch 

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

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

165 by_arch[entry["architecture_id"]] = entry 

166 if by_arch: 

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

168 return by_arch 

169 

170 

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

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

173 

174 Args: 

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

176 'tags', 'cardData']) 

177 """ 

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

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

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

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

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

183 card_language = card_data.get("language") 

184 languages = extract_languages(card_language, tags) 

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

186 

187 

188def _build_model_entry( 

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

190) -> dict: 

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

192 

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

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

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

196 """ 

197 entry = { 

198 "architecture_id": architecture_id, 

199 "model_id": model_id, 

200 "status": 0, 

201 "verified_date": None, 

202 "metadata": None, 

203 "note": None, 

204 "phase1_score": None, 

205 "phase2_score": None, 

206 "phase3_score": None, 

207 "phase4_score": None, 

208 "phase7_score": None, 

209 "phase8_score": None, 

210 "phase9_score": None, 

211 } 

212 if signals is not None: 

213 hinted = profile_from_hf_signals(model_id, architecture_id, signals) 

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

215 deliberately_curated = ( 

216 model_id in MODEL_PROFILE_OVERRIDES or architecture_id in ARCHITECTURE_PROFILE_KINDS 

217 ) 

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

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

220 logger.warning( 

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

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

223 ) 

224 if not is_default_profile(resolved): 

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

226 items = list(entry.items()) 

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

228 entry.clear() 

229 entry.update(items) 

230 return entry 

231 

232 

233def _canonical_author_sweep( 

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

235 supported_models: list[dict], 

236 seen_models: set[str], 

237 architecture: Optional[str] = None, 

238 refresh_profiles: bool = False, 

239) -> int: 

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

241 

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

243 only admit models whose extracted arch matches it. 

244 """ 

245 from . import CANONICAL_AUTHORS_BY_ARCH, HF_SUPPORTED_ARCHITECTURES 

246 

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

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

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

250 for author in authors: 

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

252 

253 added = 0 

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

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

256 continue 

257 try: 

258 models_iter = api.list_models( 

259 author=author, 

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

261 ) 

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

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

264 continue 

265 

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

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

268 try: 

269 for model in models_iter: 

270 if model.id in seen_models: 

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

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

273 if refresh_profiles: 

274 existing_entry = existing_by_id.get(model.id) 

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

276 resolved = resolve_profile( 

277 model.id, 

278 existing_entry.get("architecture_id"), 

279 signals=_extract_profile_signals(model), 

280 ) 

281 if not is_default_profile(resolved): 

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

283 continue 

284 if is_quantized_model(model.id): 

285 continue 

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

287 if model_arch is None or model_arch not in HF_SUPPORTED_ARCHITECTURES: 

288 continue 

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

290 continue 

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

292 if model_arch not in expected_archs: 

293 continue 

294 signals = _extract_profile_signals(model) 

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

296 seen_models.add(model.id) 

297 added += 1 

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

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

300 logger.warning( 

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

302 ) 

303 continue 

304 return added 

305 

306 

307def scrape_all_models( 

308 output_dir: Path, 

309 max_models: Optional[int] = None, 

310 task: str = "text-generation", 

311 batch_size: int = 1000, 

312 checkpoint_interval: int = 5000, 

313 min_downloads: int = 500, 

314 canonical_sweep: bool = True, 

315 architecture: Optional[str] = None, 

316 refresh_profiles: bool = False, 

317) -> tuple[dict, dict]: 

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

319 

320 This is the comprehensive scraper that: 

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

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

323 3. Iterates through ALL models for a given task 

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

325 5. Categorizes into supported vs unsupported 

326 6. Saves checkpoints periodically for long runs 

327 

328 Output format matches schemas.py exactly (SupportedModelsReport and 

329 ArchitectureGapsReport). 

330 

331 Args: 

332 output_dir: Directory to write JSON data files 

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

334 task: HuggingFace task filter (default: text-generation) 

335 batch_size: Log progress every N models 

336 checkpoint_interval: Save checkpoint every N models 

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

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

339 below the download threshold (default: True). 

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

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

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

343 a single new adapter without rescanning every architecture. 

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

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

346 extra requests (default: False). 

347 

348 Returns: 

349 Tuple of (supported_models_dict, architecture_gaps_dict) 

350 """ 

351 try: 

352 from huggingface_hub import HfApi 

353 except ImportError: 

354 raise ImportError( 

355 "huggingface_hub is required for scraping. " 

356 "Install it with: pip install huggingface_hub" 

357 ) 

358 

359 from transformer_lens.utilities.hf_utils import get_hf_token 

360 

361 api = HfApi(token=get_hf_token()) 

362 output_dir = Path(output_dir) 

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

364 

365 # Load existing models from supported_models.json 

366 existing_model_ids, existing_models = _load_existing_models(output_dir) 

367 

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

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

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

371 # is reflected in the final write. 

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

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

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

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

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

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

378 

379 scanned = 0 

380 skipped = 0 

381 new_supported = 0 

382 errors = 0 

383 start_time = time.time() 

384 

385 # Check for existing checkpoint to resume from 

386 checkpoint_path = output_dir / "scrape_checkpoint.json" 

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

388 

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

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

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

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

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

394 # client-side filter. 

395 from . import CANONICAL_AUTHORS_BY_ARCH 

396 

397 skip_main_scan = architecture is not None and architecture in CANONICAL_AUTHORS_BY_ARCH 

398 if skip_main_scan: 

399 assert architecture is not None # narrowed by skip_main_scan 

400 logger.info( 

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

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

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

404 ) 

405 if not canonical_sweep: 

406 logger.warning( 

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

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

409 "discover models." 

410 ) 

411 

412 if not skip_main_scan and checkpoint_path.exists(): 

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

414 with open(checkpoint_path) as f: 

415 checkpoint = json.load(f) 

416 # Merge checkpoint data with existing 

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

418 for model in checkpoint_supported: 

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

420 supported_models.append(model) 

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

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

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

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

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

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

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

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

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

430 

431 if not skip_main_scan: 

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

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

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

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

436 if max_models: 

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

438 else: 

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

440 

441 try: 

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

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

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

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

446 # 1000 req / 5 min limit). 

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

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

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

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

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

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

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

454 list_kwargs: dict = { 

455 "filter": task, 

456 "sort": "downloads", 

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

458 } 

459 if max_models is not None: 

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

461 

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

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

464 max_retries = 10 

465 for attempt in range(max_retries + 1): 

466 if skip_main_scan: 

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

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

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

470 # fine-tunes for the same arch. 

471 break 

472 try: 

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

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

475 if model.id in seen_models: 

476 skipped += 1 

477 if refresh_profiles: 

478 existing_entry = existing_by_id.get(model.id) 

479 if ( 

480 existing_entry is not None 

481 and "prompt_profile" not in existing_entry 

482 ): 

483 resolved = resolve_profile( 

484 model.id, 

485 existing_entry.get("architecture_id"), 

486 signals=_extract_profile_signals(model), 

487 ) 

488 if not is_default_profile(resolved): 

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

490 continue 

491 

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

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

494 # all remaining models will also be below it. 

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

496 if downloads < min_downloads: 

497 logger.info( 

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

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

500 f"Stopping scan." 

501 ) 

502 break 

503 

504 scanned += 1 

505 seen_models.add(model.id) 

506 

507 if max_models and scanned > max_models: 

508 break 

509 

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

511 # TransformerLens requires full-precision weights. 

512 if is_quantized_model(model.id): 

513 continue 

514 

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

516 arch = _extract_architecture(model) 

517 

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

519 # Applied before classification so the unsupported counters reflect 

520 # only the architecture under inspection. 

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

522 continue 

523 

524 if arch is None: 

525 errors += 1 

526 elif arch in HF_SUPPORTED_ARCHITECTURES: 

527 signals = _extract_profile_signals(model) 

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

529 new_supported += 1 

530 else: 

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

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

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

534 if len(samples) < max_samples: 

535 samples.append(model.id) 

536 # Accumulate downloads for relevancy scoring 

537 unsupported_arch_downloads[arch] = ( 

538 unsupported_arch_downloads.get(arch, 0) + downloads 

539 ) 

540 # Track smallest model per arch for benchmarkability 

541 param_count = _extract_param_count(model) 

542 if param_count is not None: 

543 current_min = unsupported_arch_min_params.get(arch) 

544 if current_min is None or param_count < current_min: 

545 unsupported_arch_min_params[arch] = param_count 

546 

547 # Progress logging 

548 if scanned % batch_size == 0: 

549 elapsed = time.time() - start_time 

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

551 logger.info( 

552 f"Scanned {scanned} new | " 

553 f"Skipped {skipped} existing | " 

554 f"New supported: {new_supported} | " 

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

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

557 f"Errors: {errors} | " 

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

559 ) 

560 

561 # Save checkpoint periodically 

562 if scanned % checkpoint_interval == 0: 

563 _save_checkpoint( 

564 checkpoint_path, 

565 supported_models, 

566 unsupported_arch_counts, 

567 unsupported_arch_samples, 

568 list(seen_models), 

569 scanned, 

570 skipped, 

571 unsupported_arch_downloads, 

572 unsupported_arch_min_params, 

573 ) 

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

575 

576 break # Iteration completed successfully, exit retry loop 

577 

578 except Exception as exc: 

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

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

581 logger.warning( 

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

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

584 ) 

585 _save_checkpoint( 

586 checkpoint_path, 

587 supported_models, 

588 unsupported_arch_counts, 

589 unsupported_arch_samples, 

590 list(seen_models), 

591 scanned, 

592 skipped, 

593 unsupported_arch_downloads, 

594 unsupported_arch_min_params, 

595 ) 

596 time.sleep(wait) 

597 skipped = 0 # Reset skip counter for restart 

598 else: 

599 raise 

600 

601 except KeyboardInterrupt: 

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

603 _save_checkpoint( 

604 checkpoint_path, 

605 supported_models, 

606 unsupported_arch_counts, 

607 unsupported_arch_samples, 

608 list(seen_models), 

609 scanned, 

610 skipped, 

611 unsupported_arch_downloads, 

612 unsupported_arch_min_params, 

613 ) 

614 raise 

615 except Exception as e: 

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

617 _save_checkpoint( 

618 checkpoint_path, 

619 supported_models, 

620 unsupported_arch_counts, 

621 unsupported_arch_samples, 

622 list(seen_models), 

623 scanned, 

624 skipped, 

625 unsupported_arch_downloads, 

626 unsupported_arch_min_params, 

627 ) 

628 raise 

629 

630 if canonical_sweep: 

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

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

633 try: 

634 canonical_added = _canonical_author_sweep( 

635 api, 

636 supported_models, 

637 seen_models, 

638 architecture=architecture, 

639 refresh_profiles=refresh_profiles, 

640 ) 

641 new_supported += canonical_added 

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

643 except Exception as exc: 

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

645 

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

647 elapsed = time.time() - start_time 

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

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

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

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

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

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

654 

655 # Count unique supported architectures and verified/provisional models 

656 supported_arch_ids: set[str] = set() 

657 total_verified = 0 

658 total_provisional = 0 

659 for model in supported_models: 

660 supported_arch_ids.add(model["architecture_id"]) 

661 if model.get("status", 0) == 1: 

662 total_verified += 1 

663 elif model.get("status", 0) == 4: 

664 total_provisional += 1 

665 

666 # Build scan info (shared by both reports) 

667 scan_info = { 

668 "total_scanned": scanned, 

669 "task_filter": task, 

670 "min_downloads": min_downloads, 

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

672 } 

673 

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

675 supported_report = { 

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

677 "scan_info": scan_info, 

678 "total_architectures": len(supported_arch_ids), 

679 "total_models": len(supported_models), 

680 "total_verified": total_verified, 

681 "total_provisional": total_provisional, 

682 "models": supported_models, 

683 } 

684 

685 # Write supported models (single file) 

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

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

688 f.write("\n") 

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

690 

691 # Build architecture gaps report (matches ArchitectureGapsReport schema) 

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

693 from transformer_lens.tools.model_registry.relevancy import compute_scores_for_gaps 

694 

695 gaps: list[dict] = [ 

696 { 

697 "architecture_id": arch, 

698 "total_models": count, 

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

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

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

702 } 

703 for arch, count in unsupported_arch_counts.items() 

704 ] 

705 

706 # Merge with gaps from prior scrapes so a sequential text-generation + 

707 # text2text-generation run doesn't lose the first pass's data. For overlapping 

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

709 # union sample_models (capped at 10). 

710 existing_gaps = _load_existing_gaps(output_dir) 

711 if existing_gaps: 

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

713 merged: list[dict] = [] 

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

715 o = existing_gaps.get(arch) 

716 n = new_by_arch.get(arch) 

717 if o is None and n is not None: 

718 merged.append(n) 

719 continue 

720 if n is None and o is not None: 

721 merged.append(o) 

722 continue 

723 assert o is not None and n is not None 

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

725 merged_samples: list[str] = [] 

726 seen_samples: set[str] = set() 

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

728 if s not in seen_samples: 

729 merged_samples.append(s) 

730 seen_samples.add(s) 

731 if len(merged_samples) >= 10: 

732 break 

733 min_p = [ 

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

735 ] 

736 merged.append( 

737 { 

738 "architecture_id": arch, 

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

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

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

742 "sample_models": merged_samples, 

743 } 

744 ) 

745 gaps = merged 

746 

747 # Compute relevancy scores and sort by score descending 

748 compute_scores_for_gaps(gaps) 

749 

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

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

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

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

754 seen_archs: set[str] = set() 

755 deduped: list[dict] = [] 

756 for g in gaps: 

757 arch_id = g["architecture_id"] 

758 if arch_id in HF_SUPPORTED_ARCHITECTURES: 

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

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

761 continue 

762 if arch_id in seen_archs: 

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

764 continue 

765 seen_archs.add(arch_id) 

766 deduped.append(g) 

767 gaps = deduped 

768 

769 gaps_report = { 

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

771 "scan_info": scan_info, 

772 "total_unsupported_architectures": len(gaps), 

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

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

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

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

777 "gaps": gaps, 

778 } 

779 

780 gaps_path = output_dir / "architecture_gaps.json" 

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

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

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

784 

785 # Write verification history placeholder (single file) 

786 verification_path = output_dir / "verification_history.json" 

787 if not verification_path.exists(): 

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

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

790 f.write("\n") 

791 

792 # Clean up checkpoint on successful completion 

793 if checkpoint_path.exists(): 

794 checkpoint_path.unlink() 

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

796 

797 # Print summary 

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

799 logger.info("SCAN SUMMARY") 

800 logger.info("=" * 70) 

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

802 logger.info(f"\nSUPPORTED ARCHITECTURES ({len(supported_arch_ids)}):") 

803 

804 # Count models per supported architecture 

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

806 for model in supported_models: 

807 arch = model["architecture_id"] 

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

809 

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

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

812 

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

814 for gap in gaps[:20]: 

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

816 logger.info( 

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

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

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

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

821 ) 

822 

823 if len(gaps) > 20: 

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

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

826 

827 logger.info("=" * 70)