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, load model without weights (on meta device) for config inspection only.

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

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

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

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

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

92 When provided, load_weights is ignored.

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

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

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

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

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

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

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

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

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

102 exclusive with ``device``.

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

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

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

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

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

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

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

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

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

112 may degrade past the trained length for rotary models.

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

114 config, model, and tokenizer loading.

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

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

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

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

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

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

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

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

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

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

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

126

127 Returns:

128 The bridge to the loaded model.

129 """

130 official_name = resolve_model_alias(model_name)

131 if official_name is not None:

132 logging.warning(

133 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."

134 )

135 model_name = official_name

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

137 if revision is not None:

138 raise ValueError(

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

140 )

141 revision = _resolve_checkpoint_to_revision(model_name, checkpoint_index, checkpoint_value)

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

143 from transformer_lens.utilities.hf_utils import (

144 autoconfig_with_remote_post_init_compat,

145 autotokenizer_with_special_token_compat,

146 get_hf_token,

147 )

148

149 _hf_token = get_hf_token()

150 if hf_model is not None:

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

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

153 hf_config = copy.deepcopy(hf_model.config)

154 else:

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

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

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

158 hf_config = autoconfig_with_remote_post_init_compat(

159 model_name,

160 auto_config=AutoConfig,

161 output_attentions=True,

162 trust_remote_code=trust_remote_code,

163 token=_hf_token,

164 revision=revision,

165 code_revision=code_revision,

166 )

167 _n_ctx_field: str | None = None

168 if n_ctx is not None:

169 if n_ctx <= 0:

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

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

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

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

174 "n_positions",

175 "max_position_embeddings",

176 "max_context_length",

177 "max_length",

178 "seq_length",

179 "max_sequence_length",

180 ):

181 if hasattr(hf_config, _field):

182 _n_ctx_field = _field

183 break

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

185 raise ValueError(

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

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

188 )

189 _default_n_ctx = getattr(hf_config, _n_ctx_field)

190 if _default_n_ctx is not None and n_ctx > _default_n_ctx:

191 logging.warning(

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

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

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

195 "rotary models without RoPE scaling).",

196 n_ctx,

197 _default_n_ctx,

198 )

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

200 if hf_config_overrides and _n_ctx_field in hf_config_overrides:

201 _conflicting_value = hf_config_overrides[_n_ctx_field]

202 if _conflicting_value != n_ctx:

203 logging.warning(

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

205 "The explicit n_ctx takes precedence.",

206 n_ctx,

207 _n_ctx_field,

208 _conflicting_value,

209 )

210 hf_config_overrides = dict(hf_config_overrides or {})

211 hf_config_overrides[_n_ctx_field] = n_ctx

212 if hf_config_overrides:

213 hf_config.__dict__.update(hf_config_overrides)

214 architecture = determine_architecture_from_hf_config(hf_config)

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

216 bridge_config.trust_remote_code = trust_remote_code

217 adapter = architecture_adapter_factory.ArchitectureAdapterFactory.select_architecture_adapter(

218 bridge_config

219 )

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

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

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

223 if hf_model is not None and (

224 device_map is not None or n_devices is not None or max_memory is not None

225 ):

226 raise ValueError(

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

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

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

230 )

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

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

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

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

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

236 raise ValueError(

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

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

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

240 )

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

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

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

244 from transformer_lens.utilities.multi_gpu import (

245 MIXED_OFFLOAD_GPU_ERROR,

246 count_unique_devices,

247 find_embedding_device,

248 find_misplaced_modules,

249 is_mixed_offload_gpu,

250 maybe_cast_floating_params,

251 resolve_device_map,

252 )

253

254 resolved_device_map, resolved_max_memory = resolve_device_map(

255 n_devices, device_map, device, max_memory

256 )

257 if resolved_device_map is None:

258 if device is None:

259 device = get_device()

260 adapter.cfg.device = str(device)

261 else:

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

263 adapter.cfg.device = None

264 if model_class is None:

265 model_class = get_hf_model_class_for_architecture(architecture)

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

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

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

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

270 if isinstance(fallback_pad, list):

271 fallback_pad = fallback_pad[0] if fallback_pad else None

272 hf_config.pad_token_id = fallback_pad

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

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

275 model_kwargs["token"] = _hf_token

276 if trust_remote_code:

277 model_kwargs["trust_remote_code"] = True

278 if revision is not None:

279 model_kwargs["revision"] = revision

280 if code_revision is not None:

281 model_kwargs["code_revision"] = code_revision

282 if resolved_device_map is not None:

283 model_kwargs["device_map"] = resolved_device_map

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

285 model_kwargs["max_memory"] = resolved_max_memory

286 if offload_folder is not None:

287 model_kwargs["offload_folder"] = offload_folder

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

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

290 else:

291 # Eager is required for output_attentions hooks.

292 model_kwargs["attn_implementation"] = "eager"

293 adapter.prepare_loading(model_name, model_kwargs)

294 # Meta device_map targets crash at boot when loading weights

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

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

297 if load_weights and isinstance(resolved_device_map, dict):

298 _meta_targets = [

299 k

300 for k, v in resolved_device_map.items()

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

302 ]

303 if _meta_targets:

304 raise ValueError(

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

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

307 "Set load_weights=False for config inspection only "

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

309 )

310 if hf_model is not None:

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

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

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

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

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

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

317 force_eager_attention(hf_model)

318 elif not load_weights:

319 from_config_kwargs = {}

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

321 from_config_kwargs["trust_remote_code"] = True

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

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

324 # same config the load-weights path would.

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

326 with contextlib.redirect_stdout(None):

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

328 else:

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

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

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

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

333 model_kwargs.pop("code_revision", None)

334 try:

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

336 except RuntimeError as e:

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

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

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

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

341 raise RuntimeError(

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

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

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

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

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

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

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

349 f"weight mismatch."

350 ) from e

351 raise

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

353 if resolved_device_map is None and device is not None:

354 hf_model = hf_model.to(device)

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

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

357 # parameters before we touch them.

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

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

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

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

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

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

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

365 maybe_cast_floating_params(hf_model, dtype)

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

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

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

369 if hf_device_map_post:

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

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

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

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

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

375 unsupported = offload_values & {"meta"}

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

377 raise ValueError(

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

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

380 )

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

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

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

384 "cpu" in offload_values

385 and device_map is None

386 and n_devices is not None

387 and n_devices > 1

388 ):

389 raise ValueError(

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

391 "explicitly for CPU placement."

392 )

393 misplaced = find_misplaced_modules(hf_model)

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

395 details = "; ".join(

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

397 for name, mapped, actual in misplaced

398 )

399 raise ValueError(

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

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

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

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

404 "Map tied modules to the same device."

405 )

406 embedding_device = find_embedding_device(hf_model)

407 if embedding_device is not None:

408 adapter.cfg.device = str(embedding_device)

409 adapter.cfg.n_devices = count_unique_devices(hf_model)

410 elif adapter.cfg.device is None:

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

412 try:

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

414 except StopIteration:

415 adapter.cfg.device = "cpu"

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

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

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

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

420 if _actual != n_ctx:

421 logging.warning(

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

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

424 "accept sequences longer than %s.",

425 n_ctx,

426 _n_ctx_field,

427 _actual,

428 _actual,

429 )

430 adapter.prepare_model(hf_model)

431 if tokenizer is not None:

432 tokenizer = configure_tokenizer(tokenizer, adapter.cfg)

433 elif not skip_tokenizer_for_modality(adapter.cfg):

434 token_arg = get_hf_token()

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

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

437 tokenizer_source = model_name

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

439 tokenizer_source = adapter.cfg.tokenizer_name

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

441 try:

442 base_tokenizer = autotokenizer_with_special_token_compat(

443 tokenizer_source,

444 auto_tokenizer=AutoTokenizer,

445 add_bos_token=True,

446 use_fast=use_fast,

447 token=token_arg,

448 trust_remote_code=trust_remote_code,

449 revision=revision,

450 )

451 except ValueError:

452 base_tokenizer = AutoTokenizer.from_pretrained(

453 tokenizer_source,

454 use_fast=use_fast,

455 token=token_arg,

456 trust_remote_code=trust_remote_code,

457 revision=revision,

458 )

459 tokenizer = configure_tokenizer(base_tokenizer, adapter.cfg)

460 from transformer_lens.model_bridge.sources.transformers_driver import (

461 TransformersDriver,

462 )

463

464 driver = TransformersDriver(hf_model, adapter, tokenizer)

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

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

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

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

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

470 from transformer_lens.tools.model_registry.checkpoints import (

471 get_checkpoint_labels,

472 )

473

474 labels, _label_kind = get_checkpoint_labels(model_name)

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

476 resolved_value = checkpoint_value

477 else:

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

479 resolved_value = labels[checkpoint_index]

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

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

482 return bridge