Coverage for transformer_lens/model_bridge/sources/transformers/source.py: 90%

1"""``boot`` — load a model via HuggingFace ``transformers`` and wrap it in a TransformerBridge."""

2from __future__ import annotations

4import contextlib

5import copy

6import logging

7import warnings

8from typing import Any

10import torch

11from transformers import (

12 AutoConfig,

13 AutoTokenizer,

14 PreTrainedModel,

15 PreTrainedTokenizerBase,

16)

18# Module import (not a name import): architecture_adapter_factory imports model_bridge,

19# which imports this module, so its names may not be defined yet at this point.

20from transformer_lens.factories import architecture_adapter_factory

21from transformer_lens.model_bridge.bridge import TransformerBridge

22from transformer_lens.model_bridge.sources._bridge_builder import (

23 build_bridge_config_from_hf,

24 configure_tokenizer,

25 skip_tokenizer_for_modality,

26)

27from transformer_lens.model_bridge.sources._hf_format import (

28 determine_architecture_from_hf_config,

29)

30from transformer_lens.tools.model_registry.registry_io import resolve_model_alias

31from transformer_lens.utilities import get_device

32from transformer_lens.utilities.attn_implementation import force_eager_attention

34from .helpers import (

35 _resolve_checkpoint_to_revision,

36 get_hf_model_class_for_architecture,

37 load_modality_processor,

38)

40# Suppress transformers warnings that go to stderr; otherwise notebook tests fail

41# on unexpected stderr output.

42warnings.filterwarnings("ignore", message=".*generation flags.*not valid.*")

43logging.getLogger("transformers").setLevel(logging.ERROR)

46def boot(

47 model_name: str,

48 hf_config_overrides: dict | None = None,

49 device: str | torch.device | None = None,

50 dtype: torch.dtype = torch.float32,

51 tokenizer: PreTrainedTokenizerBase | None = None,

52 load_weights: bool = True,

53 trust_remote_code: bool = False,

54 model_class: Any | None = None,

55 hf_model: Any | None = None,

56 n_ctx: int | None = None,

57 revision: str | None = None,

58 code_revision: str | None = None,

59 checkpoint_index: int | None = None,

60 checkpoint_value: int | None = None,

61 # Multi-device placement (accelerate-dispatched). Mutually exclusive with device.

62 device_map: str | dict[str, str | int] | None = None,

63 n_devices: int | None = None,

64 max_memory: dict[str | int, str | int] | None = None,

65 offload_folder: str | None = None,

66) -> TransformerBridge:

67 """Boot a model from HuggingFace (exposed as ``TransformerBridge.boot_transformers``).

69 Returns raw HF weights by default — logits/activations match HF, *not*

70 legacy ``HookedTransformer`` (which folds LayerNorm + centers weights).

71 Call ``enable_compatibility_mode()`` on the result for HookedTransformer-

72 equivalent numerics. Generation, argmax, and CE loss are unaffected.

74 Attention implementation is forced to ``"eager"`` so hooks can capture scores

75 and patterns. For an apples-to-apples HF comparison, load the HF model with

76 ``attn_implementation="eager"`` too; comparing against the default ``"sdpa"``

77 shows ~1e-3 fp32 drift from kernel-level op reordering, not a bridge bug.

79 Args:

80 model_name: The name of the model to load.

81 hf_config_overrides: Optional overrides applied to the HuggingFace config before model load.

82 device: The device to use. If None, will be determined automatically. Mutually exclusive

83 with ``device_map``.

84 dtype: The dtype to use for the model.

85 tokenizer: Optional pre-initialized tokenizer to use; if not provided one will be created.

86 load_weights: If False, build the model with ``from_config`` instead of loading the

87 checkpoint: real parameters with the architecture's random initialisation, on the CPU,

88 in the config's dtype and in training mode. ``device`` is recorded in ``cfg.device``

89 but not applied, and neither are ``dtype`` nor ``device_map``; move, cast and

90 ``eval()`` the model before running it on another device or in another dtype.

91 model_class: Optional HuggingFace model class to use instead of the default auto-detected

92 class. When the class name matches a key in SUPPORTED_ARCHITECTURES, the corresponding

93 adapter is selected automatically (e.g., BertForNextSentencePrediction).

94 hf_model: Optional pre-loaded HuggingFace model to use instead of loading one. Useful for

95 models loaded with custom configurations (e.g., quantization via BitsAndBytesConfig).

96 When provided, load_weights is ignored.

97 device_map: HuggingFace-style device map (``"auto"``, ``"balanced"``, dict, etc.) for

98 dispatched inference. Explicit maps may include CPU and disk targets; meta targets

99 are still rejected when ``load_weights=True`` (meta has no real data to offload from,

100 unlike disk/cpu). Mixed CPU/disk + GPU maps are rejected too, not because they're

101 known to be broken but because CPU/disk offload has only been verified on CPU-only

102 hardware — no GPU to mix in. ``bridge.enable_compatibility_mode()`` (with weight

103 processing, i.e. not ``no_processing=True``) is unsupported on a CPU/disk-offloaded

104 bridge and raises immediately rather than mid-fold; the default (non-compat-mode)

105 forward pass, ``run_with_cache``, and hooks all work normally under offload. Mutually

106 exclusive with ``device``.

107 n_devices: Convenience: split the model across this many CUDA devices (translated to a

108 ``max_memory`` dict internally). Requires CUDA with at least this many visible devices.

109 max_memory: Optional per-device memory budget for HF's dispatcher.

110 offload_folder: Directory for disk-offloaded weight shards when ``device_map`` includes

111 a ``"disk"`` target. Defaults to a temporary directory (HF's own default) if omitted.

112 n_ctx: Optional context length override. The bridge normally uses the model's documented

113 max context from the HF config. Setting this writes to whichever HF field the model

114 uses (n_positions / max_position_embeddings / etc.), so callers don't need to know

115 the field name. If larger than the model's default, a warning is emitted — quality

116 may degrade past the trained length for rotary models.

117 revision: Optional HF revision string (branch, tag, or commit). Forwarded to

118 config, model, and tokenizer loading.

119 Mutually exclusive with ``checkpoint_index`` and ``checkpoint_value``.

120 code_revision: Optional HF revision string pinning the ``trust_remote_code``

121 modeling file separately from the weights (HF's own ``code_revision``).

122 Forwarded alongside ``revision`` to config and model loading and to the

123 adapters' remote-code patching; weights and tokenizer keep ``revision``.

124 checkpoint_index: Index into the available training checkpoints for the model family.

125 Convenience over ``revision`` for checkpointed models like EleutherAI/pythia* and

126 stanford-crfm/*. Resolved to a revision string via the known per-family naming

127 conventions (``step{value}`` for Pythia, ``checkpoint-{value}`` for stanford-crfm).

128 checkpoint_value: Training step or token count of the desired checkpoint. Alternative to

129 ``checkpoint_index``; must be one of the labels returned by ``get_checkpoint_labels``.

130

131 Returns:

132 The bridge to the loaded model.

133 """

134 official_name = resolve_model_alias(model_name)

135 if official_name is not None:

136 logging.warning(

137 f"DEPRECATED: You are using a deprecated, model_name alias '{model_name}'. TransformerLens will now load the official transformers model name, '{official_name}' instead.\n Please update your code to use the official name by changing model_name from '{model_name}' to '{official_name}'.\nSince TransformerLens v3, all model names should be the official transformers model names.\nThe aliases may be removed in a future version of TransformerLens, so please do the update now."

138 )

139 model_name = official_name

140 if checkpoint_index is not None or checkpoint_value is not None:

141 if revision is not None:

142 raise ValueError(

143 "Specify either revision= or checkpoint_index/checkpoint_value, not both."

144 )

145 revision = _resolve_checkpoint_to_revision(model_name, checkpoint_index, checkpoint_value)

146 # Pass HF token for gated model access (e.g. meta-llama/*)

147 from transformer_lens.utilities.hf_utils import (

148 autoconfig_with_remote_post_init_compat,

149 autotokenizer_with_special_token_compat,

150 get_hf_token,

151 )

152

153 _hf_token = get_hf_token()

154 if hf_model is not None:

155 # Reuse the pre-loaded model's config to avoid a Hub call when model_name

156 # is a Hub repo ID but the model is already loaded locally.

157 hf_config = copy.deepcopy(hf_model.config)

158 else:

159 # Compat wrapper: 4.x-era remote-code configs (OpenELM) define an

160 # argless __post_init__ that 5.x's dataclass machinery calls with the

161 # class's own fields as kwargs — unloadable without the shim.

162 hf_config = autoconfig_with_remote_post_init_compat(

163 model_name,

164 auto_config=AutoConfig,

165 output_attentions=True,

166 trust_remote_code=trust_remote_code,

167 token=_hf_token,

168 revision=revision,

169 code_revision=code_revision,

170 )

171 _n_ctx_field: str | None = None

172 if n_ctx is not None:

173 if n_ctx <= 0:

174 raise ValueError(f"n_ctx must be a positive integer, got n_ctx={n_ctx}.")

175 # Resolve n_ctx to whichever HF config field this model uses. Mirrors the order in

176 # map_default_transformer_lens_config so the TL config derivation picks up the override.

177 for _field in ( 177 ↛ 188line 177 didn't jump to line 188 because the loop on line 177 didn't complete

178 "n_positions",

179 "max_position_embeddings",

180 "max_context_length",

181 "max_length",

182 "seq_length",

183 "max_sequence_length",

184 ):

185 if hasattr(hf_config, _field):

186 _n_ctx_field = _field

187 break

188 if _n_ctx_field is None: 188 ↛ 189line 188 didn't jump to line 189 because the condition on line 188 was never true

189 raise ValueError(

190 f"Cannot apply n_ctx={n_ctx}: no recognized context-length field on "

191 f"HF config for {model_name}. Use hf_config_overrides instead."

192 )

193 _default_n_ctx = getattr(hf_config, _n_ctx_field)

194 if _default_n_ctx is not None and n_ctx > _default_n_ctx:

195 logging.warning(

196 "Setting n_ctx=%d which is larger than the model's default "

197 "context length of %d. The model was not trained on sequences "

198 "this long and may produce unreliable results (especially for "

199 "rotary models without RoPE scaling).",

200 n_ctx,

201 _default_n_ctx,

202 )

203 # Warn if the caller also set the same field via hf_config_overrides — explicit n_ctx wins.

204 if hf_config_overrides and _n_ctx_field in hf_config_overrides:

205 _conflicting_value = hf_config_overrides[_n_ctx_field]

206 if _conflicting_value != n_ctx:

207 logging.warning(

208 "Both n_ctx=%d and hf_config_overrides['%s']=%s were provided. "

209 "The explicit n_ctx takes precedence.",

210 n_ctx,

211 _n_ctx_field,

212 _conflicting_value,

213 )

214 hf_config_overrides = dict(hf_config_overrides or {})

215 hf_config_overrides[_n_ctx_field] = n_ctx

216 if hf_config_overrides:

217 hf_config.__dict__.update(hf_config_overrides)

218 architecture = determine_architecture_from_hf_config(hf_config)

219 bridge_config = build_bridge_config_from_hf(hf_config, architecture, model_name, dtype)

220 bridge_config.trust_remote_code = trust_remote_code

221 adapter = architecture_adapter_factory.ArchitectureAdapterFactory.select_architecture_adapter(

222 bridge_config

223 )

224 # Pre-loaded models carry their own weight placement (possibly set by the caller via

225 # device_map). Passing device_map / n_devices / max_memory alongside hf_model= is ambiguous

226 # and would silently be ignored, so fail loudly.

227 if hf_model is not None and (

228 device_map is not None or n_devices is not None or max_memory is not None

229 ):

230 raise ValueError(

231 "device_map / n_devices / max_memory are only supported when the bridge loads "

232 "the HF model itself. When passing hf_model=..., apply device_map via "

233 "AutoModel.from_pretrained before handing the model to the bridge."

234 )

235 # Stateful/SSM (Mamba) models keep a per-layer recurrent cache that must live on that

236 # layer's device. The bridge allocates the stateful cache on a single cfg.device, so

237 # cross-device splits would silently misplace the cache. Blocked until v2.

238 if (n_devices is not None and n_devices > 1) or device_map is not None:

239 if getattr(bridge_config, "is_stateful", False): 239 ↛ 240line 239 didn't jump to line 240 because the condition on line 239 was never true

240 raise ValueError(

241 "Multi-device splits are not yet supported for stateful (SSM / Mamba) "

242 "architectures: the stateful cache allocation is single-device. "

243 "Load on one device, or wait for v2 support."

244 )

245 # Resolve device_map before defaulting `device` — the two are mutually exclusive and the

246 # resolver raises on conflict. If n_devices>1 is passed it's translated into a device_map +

247 # max_memory pair here so downstream code only needs to check the resolved values.

248 from transformer_lens.utilities.multi_gpu import (

249 MIXED_OFFLOAD_GPU_ERROR,

250 count_unique_devices,

251 find_embedding_device,

252 find_misplaced_modules,

253 is_mixed_offload_gpu,

254 maybe_cast_floating_params,

255 resolve_device_map,

256 )

257

258 resolved_device_map, resolved_max_memory = resolve_device_map(

259 n_devices, device_map, device, max_memory

260 )

261 if resolved_device_map is None:

262 if device is None:

263 device = get_device()

264 adapter.cfg.device = str(device)

265 else:

266 # cfg.device set from hf_device_map after the model is loaded; provisionally None.

267 adapter.cfg.device = None

268 if model_class is None:

269 model_class = get_hf_model_class_for_architecture(architecture)

270 # Ensure pad_token_id exists (v5 raises AttributeError if missing).

271 if not hasattr(hf_config, "pad_token_id") or "pad_token_id" not in hf_config.__dict__:

272 fallback_pad = getattr(hf_config, "eos_token_id", None)

273 # eos_token_id can be a list (Gemma3 uses [1, 106]); take the first.

274 if isinstance(fallback_pad, list):

275 fallback_pad = fallback_pad[0] if fallback_pad else None

276 hf_config.pad_token_id = fallback_pad

277 model_kwargs = {"config": hf_config, "torch_dtype": dtype}

278 if _hf_token: 278 ↛ 280line 278 didn't jump to line 280 because the condition on line 278 was always true

279 model_kwargs["token"] = _hf_token

280 if trust_remote_code:

281 model_kwargs["trust_remote_code"] = True

282 if revision is not None:

283 model_kwargs["revision"] = revision

284 if code_revision is not None:

285 model_kwargs["code_revision"] = code_revision

286 if resolved_device_map is not None:

287 model_kwargs["device_map"] = resolved_device_map

288 if resolved_max_memory is not None: 288 ↛ 289line 288 didn't jump to line 289 because the condition on line 288 was never true

289 model_kwargs["max_memory"] = resolved_max_memory

290 if offload_folder is not None:

291 model_kwargs["offload_folder"] = offload_folder

292 if hasattr(adapter.cfg, "attn_implementation") and adapter.cfg.attn_implementation is not None:

293 model_kwargs["attn_implementation"] = adapter.cfg.attn_implementation

294 else:

295 # Eager is required for output_attentions hooks.

296 model_kwargs["attn_implementation"] = "eager"

297 adapter.prepare_loading(model_name, model_kwargs)

298 # Meta device_map targets crash at boot when loading weights

299 # (NotImplementedError in HF tie_weights, KeyError in Accelerate offload hooks).

300 # Only accepted with load_weights=False (config inspection; map not applied).

301 if load_weights and isinstance(resolved_device_map, dict):

302 _meta_targets = [

303 k

304 for k, v in resolved_device_map.items()

305 if isinstance(v, str) and v.strip().lower() == "meta"

306 ]

307 if _meta_targets:

308 raise ValueError(

309 f"device_map contains meta target(s): {_meta_targets}. "

310 "Meta device_map values crash at boot when loading weights. "

311 "Set load_weights=False for config inspection only "

312 "(the map is not applied; parameters load on CPU via from_config)."

313 )

314 if hf_model is not None:

315 # Pre-loaded weights stay untouched (quantized models with custom device_map are

316 # why this branch takes the model as-is), but attn-impl selection is runtime

317 # dispatch — safe to switch in place. Without it, transformers' default sdpa

318 # returns attn_weights=None and hook_pattern / hook_attn_scores silently never

319 # fire. An explicit adapter/user choice wins, matching the load path above.

320 if getattr(adapter.cfg, "attn_implementation", None) is None: 320 ↛ 372line 320 didn't jump to line 372 because the condition on line 320 was always true

321 force_eager_attention(hf_model)

322 elif not load_weights:

323 from_config_kwargs = {}

324 if trust_remote_code: 324 ↛ 325line 324 didn't jump to line 325 because the condition on line 324 was never true

325 from_config_kwargs["trust_remote_code"] = True

326 # adapter.prepare_loading may have replaced model_kwargs["config"] (e.g. Qwen3.5

327 # text-only extraction); honor that here so the no-weights path uses the

328 # same config the load-weights path would.

329 prepared_config = model_kwargs.get("config", hf_config)

330 with contextlib.redirect_stdout(None):

331 hf_model = model_class.from_config(prepared_config, **from_config_kwargs)

332 else:

333 # Only Auto classes resolve remote code, so only they accept code_revision;

334 # a concrete PreTrainedModel subclass rejects it with a TypeError. The pin

335 # already did its work in AutoConfig and prepare_loading above.

336 if isinstance(model_class, type) and issubclass(model_class, PreTrainedModel):

337 model_kwargs.pop("code_revision", None)

338 try:

339 hf_model = model_class.from_pretrained(model_name, **model_kwargs)

340 except RuntimeError as e:

341 # HF refuses to load when positional-weight shapes don't match. If the user

342 # requested an n_ctx that conflicts with the saved weights (common for

343 # learned-pos-embed models like GPT-2), re-raise with a clearer message.

344 if n_ctx is not None and "ignore_mismatched_sizes" in str(e): 344 ↛ 355line 344 didn't jump to line 355 because the condition on line 344 was always true

345 raise RuntimeError(

346 f"Failed to load {model_name} with n_ctx={n_ctx}: the pretrained "

347 f"weights' positional-embedding shape does not match the requested "

348 f"context length. This affects models with learned positional "

349 f"embeddings (e.g. GPT-2, OPT). Options: (1) use the model's "

350 f"default n_ctx, (2) pass load_weights=False if you only need "

351 f"config inspection, or (3) choose a rotary-embedding model "

352 f"(e.g. Llama, Mistral) which supports n_ctx changes without "

353 f"weight mismatch."

354 ) from e

355 raise

356 # Skip explicit .to(device) when accelerate has placed weights via device_map.

357 if resolved_device_map is None and device is not None:

358 hf_model = hf_model.to(device)

359 # Cast params to dtype; preserve float32 buffers (e.g. RotaryEmbedding.inv_freq).

360 # Use module-level alignment so Accelerate can temporarily materialize offloaded

361 # parameters before we touch them.

362 # Skip dtype normalization entirely when the model has an active quantizer.

363 # from_pretrained settled the load dtype above, and on a quantized checkpoint

364 # that is the quantizer's effective dtype rather than necessarily the requested

365 # one (HfQuantizer.update_dtype lets AWQ, fbgemm-FP8 and FP-Quant override it;

366 # AWQ keys off CUDA/XPU being available at all, not off placement).

367 # Re-casting would overwrite those choices along with quantizer-owned storage,

368 # whose width is the quantizer's to choose (FP8 *or* float32 scales).

369 maybe_cast_floating_params(hf_model, dtype)

370 # Derive cfg.device / cfg.n_devices from hf_device_map when present. Covers fresh loads

371 # with a resolved device_map AND pre-loaded models with caller-dispatched device_map="auto".

372 hf_device_map_post = getattr(hf_model, "hf_device_map", None)

373 if hf_device_map_post:

374 # All-CPU placement, and disk offload, are supported (GeneralizedComponent.__call__

375 # wraps forward in Accelerate's align_module_device, so wrapped components reading

376 # raw params directly still see materialized data). meta remains rejected here: it

377 # has no real data to materialize even under offload/dispatch, unlike disk/cpu.

378 offload_values = {str(v).lower() for v in hf_device_map_post.values() if isinstance(v, str)}

379 unsupported = offload_values & {"meta"}

380 if unsupported: 380 ↛ 381line 380 didn't jump to line 381 because the condition on line 380 was never true

381 raise ValueError(

382 f"hf_device_map contains unsupported offload targets: {sorted(unsupported)}. "

383 "TransformerBridge currently supports CPU and disk device_map targets."

384 )

385 if is_mixed_offload_gpu(hf_device_map_post.values()): 385 ↛ 386line 385 didn't jump to line 386 because the condition on line 385 was never true

386 raise ValueError(f"Realized hf_device_map is unsupported: {MIXED_OFFLOAD_GPU_ERROR}")

387 if ( 387 ↛ 393line 387 didn't jump to line 393 because the condition on line 387 was never true

388 "cpu" in offload_values

389 and device_map is None

390 and n_devices is not None

391 and n_devices > 1

392 ):

393 raise ValueError(

394 "hf_device_map contains CPU targets. n_devices is GPU-only; pass device_map "

395 "explicitly for CPU placement."

396 )

397 misplaced = find_misplaced_modules(hf_model)

398 if misplaced: 398 ↛ 399line 398 didn't jump to line 399 because the condition on line 398 was never true

399 details = "; ".join(

400 f"{name!r} mapped to {mapped} but loaded on {actual}"

401 for name, mapped, actual in misplaced

402 )

403 raise ValueError(

404 f"device_map entries were not honored: {details}. This usually means the "

405 "map splits tied parameters (e.g. GPT-2's wte/lm_head share one tensor) "

406 "across devices — accelerate places a tied parameter once, leaving a module "

407 "executing on a device its weights aren't on, which crashes mid-forward. "

408 "Map tied modules to the same device."

409 )

410 embedding_device = find_embedding_device(hf_model)

411 if embedding_device is not None:

412 adapter.cfg.device = str(embedding_device)

413 adapter.cfg.n_devices = count_unique_devices(hf_model)

414 elif adapter.cfg.device is None:

415 # Pre-loaded single-device model with no hf_device_map — fall back to first param.

416 try:

417 adapter.cfg.device = str(next(hf_model.parameters()).device)

418 except StopIteration:

419 adapter.cfg.device = "cpu"

420 # Verify the n_ctx override actually took effect on the loaded model. If HF's config class

421 # silently dropped or normalized the value, warn so the user isn't misled.

422 if n_ctx is not None and _n_ctx_field is not None and hf_model is not None:

423 _actual = getattr(hf_model.config, _n_ctx_field, None)

424 if _actual != n_ctx:

425 logging.warning(

426 "n_ctx=%d was requested but hf_model.config.%s=%s after load. "

427 "The override may not have taken effect; the model may not "

428 "accept sequences longer than %s.",

429 n_ctx,

430 _n_ctx_field,

431 _actual,

432 _actual,

433 )

434 adapter.prepare_model(hf_model)

435 if tokenizer is not None:

436 tokenizer = configure_tokenizer(tokenizer, adapter.cfg)

437 elif not skip_tokenizer_for_modality(adapter.cfg):

438 token_arg = get_hf_token()

439 use_fast = getattr(adapter.cfg, "use_fast", True)

440 # Some adapters override tokenizer source (e.g. OpenELM has no tokenizer of its own).

441 tokenizer_source = model_name

442 if hasattr(adapter.cfg, "tokenizer_name") and adapter.cfg.tokenizer_name is not None: 442 ↛ 443line 442 didn't jump to line 443 because the condition on line 442 was never true

443 tokenizer_source = adapter.cfg.tokenizer_name

444 # Encoder-decoder models like T5 don't have a BOS token and raise on add_bos_token=True.

445 try:

446 base_tokenizer = autotokenizer_with_special_token_compat(

447 tokenizer_source,

448 auto_tokenizer=AutoTokenizer,

449 add_bos_token=True,

450 use_fast=use_fast,

451 token=token_arg,

452 trust_remote_code=trust_remote_code,

453 revision=revision,

454 )

455 except ValueError:

456 base_tokenizer = AutoTokenizer.from_pretrained(

457 tokenizer_source,

458 use_fast=use_fast,

459 token=token_arg,

460 trust_remote_code=trust_remote_code,

461 revision=revision,

462 )

463 tokenizer = configure_tokenizer(base_tokenizer, adapter.cfg)

464 from transformer_lens.model_bridge.sources.transformers_driver import (

465 TransformersDriver,

466 )

467

468 driver = TransformersDriver(hf_model, adapter, tokenizer)

469 bridge = TransformerBridge(hf_model, adapter, tokenizer, driver=driver)

470 load_modality_processor(bridge, adapter.cfg, model_name, trust_remote_code, _hf_token)

471 # Mirror the legacy loader: record which training checkpoint this is, so

472 # checkpoint sweeps (e.g. induction-head formation) can read it back.

473 if checkpoint_index is not None or checkpoint_value is not None:

474 from transformer_lens.tools.model_registry.checkpoints import (

475 get_checkpoint_labels,

476 )

477

478 labels, _label_kind = get_checkpoint_labels(model_name)

479 if checkpoint_value is not None: 479 ↛ 480line 479 didn't jump to line 480 because the condition on line 479 was never true

480 resolved_value = checkpoint_value

481 else:

482 assert checkpoint_index is not None # narrowed by the enclosing if

483 resolved_value = labels[checkpoint_index]

484 bridge.cfg.checkpoint_value = resolved_value # type: ignore[attr-defined]

485 bridge.cfg.checkpoint_index = labels.index(resolved_value) # type: ignore[attr-defined]

486 return bridge