Coverage for transformer_lens/lit/model.py: 19%

1"""LIT Model wrapper for TransformerLens models.

3This module provides a LIT-compatible wrapper around TransformerLens models,

4enabling the use of Google's Learning Interpretability Tool (LIT) for model visualization

5and analysis.

7The wrapper exposes:

8- Token predictions (logits, top-k tokens)

9- Per-layer embeddings (residual stream)

10- Attention patterns (all layers/heads)

11- Token gradients for salience maps

12- Loss computation

14Example usage:

15 >>> from transformer_lens import TransformerBridge # doctest: +SKIP

16 >>> from transformer_lens.lit import TransformerLensLIT # doctest: +SKIP

17 >>>

18 >>> # Load model

19 >>> model = TransformerBridge.boot_transformers("gpt2") # doctest: +SKIP

20 >>>

21 >>> # Create LIT wrapper

22 >>> lit_model = TransformerLensLIT(model) # doctest: +SKIP

23 >>>

24 >>> # Run prediction

25 >>> inputs = [{"text": "Hello, world!"}] # doctest: +SKIP

26 >>> outputs = list(lit_model.predict(inputs)) # doctest: +SKIP

28References:

29 - LIT Model API: https://pair-code.github.io/lit/documentation/api#models

30 - TransformerLens: https://github.com/TransformerLensOrg/TransformerLens

31"""

33from __future__ import annotations

35import logging

36from dataclasses import dataclass

37from typing import TYPE_CHECKING, Any, Dict, Iterable, Iterator, List, Optional

39import torch

41from .constants import DEFAULTS, ERRORS, INPUT_FIELDS, OUTPUT_FIELDS

42from .utils import (

43 check_lit_installed,

44 clean_token_strings,

45 extract_attention_from_cache,

46 get_model_info,

47 get_tokens_from_model,

48 tensor_to_numpy,

49)

51if TYPE_CHECKING:

52 from lit_nlp.api import model as lit_model_types # noqa: F401

53 from lit_nlp.api import types as lit_types_module # noqa: F401

55# Check for LIT installation and import conditionally

56if check_lit_installed(): 56 ↛ 57line 56 didn't jump to line 57 because the condition on line 56 was never true

57 from lit_nlp.api import ( # type: ignore[import-not-found] # noqa: F401

58 model as lit_model,

59 )

60 from lit_nlp.api import ( # type: ignore[import-not-found] # noqa: F401

61 types as lit_types,

62 )

63 from lit_nlp.lib import utils as lit_utils # type: ignore[import-not-found]

65 _LIT_AVAILABLE = True

66else:

67 _LIT_AVAILABLE = False

68 # Create placeholder when LIT not installed

69 lit_model = None # type: ignore[assignment]

70 lit_types = None # type: ignore[assignment]

71 lit_utils = None # type: ignore[assignment]

73logger = logging.getLogger(__name__)

76@dataclass

77class TransformerLensLITConfig:

78 """Configuration for the TransformerLensLIT wrapper."""

80 max_seq_length: int = DEFAULTS.MAX_SEQ_LENGTH

81 batch_size: int = DEFAULTS.BATCH_SIZE

82 top_k: int = DEFAULTS.TOP_K

83 compute_gradients: bool = DEFAULTS.COMPUTE_GRADIENTS

84 output_attention: bool = DEFAULTS.OUTPUT_ATTENTION

85 output_embeddings: bool = DEFAULTS.OUTPUT_EMBEDDINGS

86 output_all_layers: bool = DEFAULTS.OUTPUT_ALL_LAYERS

87 embedding_layers: Optional[List[int]] = None

88 prepend_bos: bool = DEFAULTS.PREPEND_BOS

89 device: Optional[str] = None

92def _ensure_lit_available():

93 """Raise ImportError if LIT is not available."""

94 if not _LIT_AVAILABLE:

95 raise ImportError(ERRORS.LIT_NOT_INSTALLED)

98# Create base class dynamically based on LIT availability

99if _LIT_AVAILABLE: 99 ↛ 100line 99 didn't jump to line 100 because the condition on line 99 was never true

100 _LITModelBase = lit_model.Model

101else:

102 _LITModelBase = object # type: ignore[misc,assignment]

103

104

105class TransformerLensLIT(_LITModelBase): # type: ignore[valid-type,misc]

106 """LIT Model wrapper for TransformerLens models.

107

108 This wrapper implements the LIT Model API, enabling the use of LIT's

109 visualization and analysis tools with TransformerLens models.

110

111 The wrapper provides:

112 - Token predictions with top-k probabilities

113 - Per-layer embeddings for embedding projector

114 - Attention patterns for attention visualization

115 - Token gradients for salience maps

116

117 Example:

118 >>> model = TransformerBridge.boot_transformers("gpt2") # doctest: +SKIP

119 >>> lit_model = TransformerLensLIT(model) # doctest: +SKIP

120 >>> lit_model.input_spec() # doctest: +SKIP

121 {'text': TextSegment(), ...}

122 """

123

124 def __init__(

125 self,

126 model: Any,

127 config: Optional[TransformerLensLITConfig] = None,

128 ):

129 """Initialize the LIT wrapper.

130

131 Args:

132 model: A TransformerLens model (e.g. TransformerBridge).

133 config: Optional configuration. Uses defaults if not provided.

134

135 Raises:

136 ImportError: If lit-nlp is not installed.

137 TypeError: If model does not implement the TransformerLens model interface.

138 """

139 _ensure_lit_available()

140

141 # Validate model type (structural: accepts TransformerBridge and any

142 # other conforming TransformerLens model)

143 from transformer_lens.model_protocol import TransformerLensModel

144

145 if not isinstance(model, TransformerLensModel):

146 raise TypeError(ERRORS.INVALID_MODEL.format(model_type=type(model)))

147

148 # Deliberately Any: the wrapper reaches beyond the minimal protocol

149 # surface (tokenizer, embed, callables) on whichever model conforms.

150 self.model: Any = model

151 self.config = config or TransformerLensLITConfig()

152

153 # Gradients require embeddings to be output (for alignment)

154 if self.config.compute_gradients and not self.config.output_embeddings:

155 logger.info("Enabling output_embeddings (required for compute_gradients)")

156 self.config.output_embeddings = True

157

158 # Set device

159 if self.config.device is None:

160 self.config.device = str(model.cfg.device)

161

162 # Cache model info

163 self._model_info = get_model_info(model)

164

165 logger.info(f"Created TransformerLensLIT wrapper for {self._model_info['model_name']}")

166

167 @property

168 def supports_concurrent_predictions(self) -> bool:

169 """Whether this model supports concurrent predictions.

170

171 Returns False as PyTorch models typically aren't thread-safe.

172 """

173 return False

174

175 def description(self) -> str:

176 """Return a human-readable description of the model.

177

178 Returns:

179 Model description string.

180 """

181 info = self._model_info

182 return (

183 f"TransformerLens: {info['model_name']} "

184 f"({info['n_layers']}L, {info['n_heads']}H, d={info['d_model']})"

185 )

186

187 @classmethod

188 def init_spec(cls) -> Dict[str, Any]:

189 """Return spec for model initialization in LIT UI.

190

191 This allows loading new models through the LIT interface.

192

193 Returns:

194 Specification for initialization parameters.

195 """

196 _ensure_lit_available()

197 return {

198 "model_name": lit_types.String( # type: ignore[union-attr]

199 default="gpt2-small",

200 required=True,

201 ),

202 "max_seq_length": lit_types.Integer( # type: ignore[union-attr]

203 default=DEFAULTS.MAX_SEQ_LENGTH,

204 min_val=1,

205 max_val=2048,

206 required=False,

207 ),

208 "compute_gradients": lit_types.Boolean( # type: ignore[union-attr]

209 default=DEFAULTS.COMPUTE_GRADIENTS,

210 required=False,

211 ),

212 "output_attention": lit_types.Boolean( # type: ignore[union-attr]

213 default=DEFAULTS.OUTPUT_ATTENTION,

214 required=False,

215 ),

216 "output_embeddings": lit_types.Boolean( # type: ignore[union-attr]

217 default=DEFAULTS.OUTPUT_EMBEDDINGS,

218 required=False,

219 ),

220 }

221

222 def input_spec(self) -> Dict[str, Any]:

223 """Return spec describing the model inputs.

224

225 Defines the expected input format for the model. LIT uses this

226 to validate inputs and generate appropriate UI controls.

227

228 Returns:

229 Dictionary mapping field names to LIT type specs.

230 """

231 _ensure_lit_available()

232

233 spec = {

234 # Primary text input

235 INPUT_FIELDS.TEXT: lit_types.TextSegment(), # type: ignore[union-attr]

236 # Optional pre-tokenized input (for Integrated Gradients)

237 INPUT_FIELDS.TOKENS: lit_types.Tokens( # type: ignore[union-attr]

238 parent=INPUT_FIELDS.TEXT,

239 required=False,

240 ),

241 }

242

243 # Add optional embeddings input for Integrated Gradients

244 if self.config.output_embeddings:

245 spec[INPUT_FIELDS.TOKEN_EMBEDDINGS] = lit_types.TokenEmbeddings( # type: ignore[union-attr]

246 align=INPUT_FIELDS.TOKENS,

247 required=False,

248 )

249

250 # Add target mask for sequence salience

251 if self.config.compute_gradients:

252 spec[INPUT_FIELDS.TARGET_MASK] = lit_types.Tokens( # type: ignore[union-attr]

253 parent=INPUT_FIELDS.TEXT,

254 required=False,

255 )

256

257 return spec

258

259 def output_spec(self) -> Dict[str, Any]:

260 """Return spec describing the model outputs.

261

262 Defines all the outputs that the model produces. LIT uses this

263 to determine which visualizations to show.

264

265 Returns:

266 Dictionary mapping field names to LIT type specs.

267 """

268 _ensure_lit_available()

269

270 spec = {}

271

272 # Tokens (always output)

273 spec[OUTPUT_FIELDS.TOKENS] = lit_types.Tokens( # type: ignore[union-attr]

274 parent=INPUT_FIELDS.TEXT,

275 )

276

277 # Top-K predictions for next token

278 spec[OUTPUT_FIELDS.TOP_K_TOKENS] = lit_types.TokenTopKPreds( # type: ignore[union-attr]

279 align=OUTPUT_FIELDS.TOKENS,

280 )

281

282 # Embeddings

283 if self.config.output_embeddings:

284 # Input embeddings (for Integrated Gradients)

285 spec[OUTPUT_FIELDS.INPUT_EMBEDDINGS] = lit_types.TokenEmbeddings( # type: ignore[union-attr]

286 align=OUTPUT_FIELDS.TOKENS,

287 )

288

289 # Final layer embedding (CLS-style)

290 spec[OUTPUT_FIELDS.CLS_EMBEDDING] = lit_types.Embeddings() # type: ignore[union-attr]

291

292 # Mean pooled embedding

293 spec[OUTPUT_FIELDS.MEAN_EMBEDDING] = lit_types.Embeddings() # type: ignore[union-attr]

294

295 # Per-layer embeddings

296 layers_to_output = self._get_embedding_layers()

297 for layer in layers_to_output:

298 field_name = OUTPUT_FIELDS.LAYER_EMB_TEMPLATE.format(layer=layer)

299 spec[field_name] = lit_types.Embeddings() # type: ignore[union-attr]

300

301 # Attention patterns

302 if self.config.output_attention:

303 for layer in range(self._model_info["n_layers"]):

304 field_name = OUTPUT_FIELDS.LAYER_ATTENTION_TEMPLATE.format(layer=layer)

305 spec[field_name] = lit_types.AttentionHeads( # type: ignore[union-attr]

306 align_in=OUTPUT_FIELDS.TOKENS,

307 align_out=OUTPUT_FIELDS.TOKENS,

308 )

309

310 # Gradients for salience

311 if self.config.compute_gradients:

312 # TokenGradients spec requirements (per LIT API):

313 # - align: must point to a Tokens field (for token alignment)

314 # - grad_for: must point to a TokenEmbeddings field (for grad-dot-input)

315 # LIT's GradientNorm component computes L2 norm internally

316 # LIT's GradientDotInput component computes dot product with embeddings

317 spec[OUTPUT_FIELDS.GRAD_L2] = lit_types.TokenGradients( # type: ignore[union-attr]

318 align=OUTPUT_FIELDS.TOKENS,

319 grad_for=OUTPUT_FIELDS.INPUT_EMBEDDINGS,

320 )

321 # Gradient dot input uses same format

322 spec[OUTPUT_FIELDS.GRAD_DOT_INPUT] = lit_types.TokenGradients( # type: ignore[union-attr]

323 align=OUTPUT_FIELDS.TOKENS,

324 grad_for=OUTPUT_FIELDS.INPUT_EMBEDDINGS,

325 )

326

327 return spec

328

329 def _get_embedding_layers(self) -> List[int]:

330 """Get the layers to output embeddings for.

331

332 Returns:

333 List of layer indices.

334 """

335 if self.config.embedding_layers is not None:

336 return self.config.embedding_layers

337

338 n_layers = self._model_info["n_layers"]

339

340 if self.config.output_all_layers:

341 return list(range(n_layers))

342 else:

343 # Output first, middle, and last layers by default

344 if n_layers <= 3:

345 return list(range(n_layers))

346 return [0, n_layers // 2, n_layers - 1]

347

348 def predict(

349 self,

350 inputs: Iterable[Dict[str, Any]],

351 ) -> Iterator[Dict[str, Any]]:

352 """Run prediction on a sequence of inputs.

353

354 This is the main entry point for LIT to get model outputs.

355

356 Args:

357 inputs: Iterable of input dictionaries, each with fields

358 matching input_spec().

359

360 Yields:

361 Output dictionaries for each input, with fields matching

362 output_spec().

363 """

364 for example in inputs:

365 yield self._predict_single(example)

366

367 def _predict_single(

368 self,

369 example: Dict[str, Any],

370 ) -> Dict[str, Any]:

371 """Run prediction on a single example.

372

373 Args:

374 example: Input dictionary with text field.

375

376 Returns:

377 Output dictionary with predictions.

378 """

379 text = example[INPUT_FIELDS.TEXT]

380

381 # Check for pre-tokenized input (reserved for future use)

382 _ = example.get(INPUT_FIELDS.TOKENS)

383 _ = example.get(INPUT_FIELDS.TOKEN_EMBEDDINGS)

384

385 # Initialize output

386 output: Dict[str, Any] = {}

387

388 # Tokenize

389 if self.model.tokenizer is None:

390 raise ValueError(ERRORS.NO_TOKENIZER)

391

392 tokens, token_ids = get_tokens_from_model(

393 self.model,

394 text,

395 prepend_bos=self.config.prepend_bos,

396 max_length=self.config.max_seq_length,

397 )

398 output[OUTPUT_FIELDS.TOKENS] = clean_token_strings(tokens)

399

400 # Prepare input

401 input_tokens = token_ids.unsqueeze(0).to(self.config.device)

402

403 # Run with cache to get all activations

404 with torch.no_grad():

405 result, cache = self.model.run_with_cache(

406 input_tokens,

407 return_type="logits",

408 )

409 # Ensure logits is a tensor (run_with_cache returns Output type)

410 logits: torch.Tensor = (

411 result if isinstance(result, torch.Tensor) else torch.tensor(result)

412 )

413

414 # Top-K predictions

415 output[OUTPUT_FIELDS.TOP_K_TOKENS] = self._get_top_k_per_position(logits, len(tokens))

416

417 # Embeddings

418 if self.config.output_embeddings:

419 output.update(self._extract_embeddings(cache, len(tokens)))

420

421 # Attention

422 if self.config.output_attention:

423 output.update(self._extract_attention(cache))

424

425 # Gradients (requires separate forward pass with gradients enabled)

426 if self.config.compute_gradients:

427 output.update(self._compute_gradients(text, example))

428

429 return output

430

431 def _get_top_k_per_position(

432 self,

433 logits: torch.Tensor,

434 seq_len: int,

435 ) -> List[List[tuple]]:

436 """Get top-k predictions for each position.

437

438 Args:

439 logits: Model logits [batch, pos, vocab].

440 seq_len: Sequence length.

441

442 Returns:

443 List of lists of (token, probability) tuples.

444 """

445 results = []

446 # Ensure logits is a tensor (handle Output type from run_with_cache)

447 if not isinstance(logits, torch.Tensor):

448 logits = torch.tensor(logits)

449 probs = torch.softmax(logits[0], dim=-1)

450

451 for pos in range(seq_len):

452 top_probs, top_indices = torch.topk(probs[pos], self.config.top_k)

453 pos_results = []

454 for prob, idx in zip(top_probs.tolist(), top_indices.tolist()):

455 if self.model.tokenizer is not None:

456 token_str = self.model.tokenizer.decode([idx])

457 else:

458 token_str = f"<{idx}>"

459 pos_results.append((token_str, prob))

460 results.append(pos_results)

461

462 return results

463

464 def _extract_embeddings(

465 self,

466 cache: Any,

467 seq_len: int,

468 ) -> Dict[str, Any]:

469 """Extract embeddings from the activation cache.

470

471 Args:

472 cache: Activation cache from forward pass.

473 seq_len: Sequence length.

474

475 Returns:

476 Dictionary of embedding arrays.

477 """

478 output = {}

479

480 # Input embeddings (from hook_embed)

481 input_emb = cache["hook_embed"][0] # [seq_len, d_model]

482 output[OUTPUT_FIELDS.INPUT_EMBEDDINGS] = tensor_to_numpy(input_emb)

483

484 # Final layer embeddings

485 final_layer = self._model_info["n_layers"] - 1

486 final_resid = cache[f"blocks.{final_layer}.hook_resid_post"][0]

487

488 # CLS-style (first token)

489 output[OUTPUT_FIELDS.CLS_EMBEDDING] = tensor_to_numpy(final_resid[0])

490

491 # Mean pooled

492 output[OUTPUT_FIELDS.MEAN_EMBEDDING] = tensor_to_numpy(final_resid.mean(dim=0))

493

494 # Per-layer embeddings

495 for layer in self._get_embedding_layers():

496 resid = cache[f"blocks.{layer}.hook_resid_post"][0]

497 # Use mean pooled embedding for the layer

498 field_name = OUTPUT_FIELDS.LAYER_EMB_TEMPLATE.format(layer=layer)

499 output[field_name] = tensor_to_numpy(resid.mean(dim=0))

500

501 return output

502

503 def _extract_attention(

504 self,

505 cache: Any,

506 ) -> Dict[str, Any]:

507 """Extract attention patterns from the activation cache.

508

509 Args:

510 cache: Activation cache from forward pass.

511

512 Returns:

513 Dictionary of attention pattern arrays.

514 """

515 output = {}

516

517 for layer in range(self._model_info["n_layers"]):

518 # Get attention pattern for this layer

519 attn = extract_attention_from_cache(cache, layer, head=None, batch_idx=0)

520 # attn shape: [num_heads, query_pos, key_pos]

521 field_name = OUTPUT_FIELDS.LAYER_ATTENTION_TEMPLATE.format(layer=layer)

522 output[field_name] = attn

523

524 return output

525

526 def _compute_gradients(

527 self,

528 text: str,

529 example: Dict[str, Any],

530 ) -> Dict[str, Any]:

531 """Compute token gradients for salience.

532

533 Args:

534 text: Input text.

535 example: Full input example (may contain target_mask).

536

537 Returns:

538 Dictionary with gradient arrays.

539 """

540 output = {}

541

542 # Tokenize

543 tokens, token_ids = get_tokens_from_model(

544 self.model,

545 text,

546 prepend_bos=self.config.prepend_bos,

547 max_length=self.config.max_seq_length,

548 )

549 input_tokens = token_ids.unsqueeze(0).to(self.config.device)

550

551 # Get target mask if provided

552 target_mask = example.get(INPUT_FIELDS.TARGET_MASK)

553

554 # Get embeddings with gradient tracking

555 with torch.enable_grad():

556 # Get input embeddings and make them a leaf tensor for gradients

557 embed = self.model.embed(input_tokens).detach().clone()

558 embed.requires_grad_(True)

559

560 # Add positional embeddings if applicable

561 if self.model.cfg.positional_embedding_type == "standard":

562 # W_pos indexed by POSITION. Calling pos_embed(input_tokens) on

563 # a bridge routes to HF's wpe, which embeds whatever ids it is

564 # given — token ids, here — silently producing wrong positions.

565 pos_embed = self.model.W_pos[: input_tokens.shape[1]].unsqueeze(0)

566 residual = embed + pos_embed

567 else:

568 residual = embed

569

570 # Forward through the rest of the model

571 logits = self.model(residual, start_at_layer=0)

572

573 # Compute loss or target logit

574 if target_mask is not None:

575 # Use masked tokens as targets

576 # For now, use simple next-token prediction loss

577 pass

578

579 # Use last token prediction as target

580 target_idx = token_ids[-1].item() # Predict last token

581 target_logit = logits[0, -2, target_idx] # Logit at second-to-last position

582

583 # Backward pass

584 target_logit.backward()

585

586 # Get gradients - now embed is a leaf tensor so grad should be populated

587 if embed.grad is None:

588 # Fallback: return zeros if gradients couldn't be computed

589 gradients = torch.zeros_like(embed[0])

590 else:

591 gradients = embed.grad[0] # [seq_len, d_model]

592

593 # Return the full gradient tensor - LIT computes norms internally

594 # TokenGradients expects shape [num_tokens, emb_dim]

595 output[OUTPUT_FIELDS.GRAD_L2] = tensor_to_numpy(gradients)

596 output[OUTPUT_FIELDS.GRAD_DOT_INPUT] = tensor_to_numpy(gradients)

597

598 return output

599

600 def max_minibatch_size(self) -> int:

601 """Return the maximum batch size for prediction.

602

603 Returns:

604 Maximum batch size.

605 """

606 return self.config.batch_size

607

608 def get_embedding_table(self) -> tuple:

609 """Return the token embedding table.

610

611 Required by LIT for certain generators like HotFlip.

612

613 Returns:

614 Tuple of (vocab_list, embedding_matrix) where vocab_list is

615 a list of token strings and embedding_matrix is [vocab, d_model].

616 """

617 # Get the embedding matrix from the model

618 embed_weight = self.model.embed.W_E.detach().cpu().numpy()

619

620 # Get vocabulary list - use tokenizer's vocab size to avoid index errors

621 if self.model.tokenizer is not None:

622 # Use the tokenizer's actual vocabulary size

623 tokenizer_vocab_size = len(self.model.tokenizer)

624 # Use the smaller of embedding size and tokenizer vocab size

625 vocab_size = min(embed_weight.shape[0], tokenizer_vocab_size)

626 vocab_list = []

627 for i in range(vocab_size):

628 try:

629 token = self.model.tokenizer.decode([i])

630 vocab_list.append(token)

631 except Exception:

632 vocab_list.append(f"<{i}>")

633 # Truncate embedding matrix to match vocab_list

634 embed_weight = embed_weight[:vocab_size]

635 else:

636 vocab_list = [f"<{i}>" for i in range(embed_weight.shape[0])]

637

638 return vocab_list, embed_weight

639

640 @classmethod

641 def from_pretrained(

642 cls,

643 model_name: str,

644 config: Optional[TransformerLensLITConfig] = None,

645 **model_kwargs,

646 ) -> "TransformerLensLIT":

647 """Create a LIT wrapper from a pretrained model name.

648

649 Convenience method that boots a TransformerBridge (with compatibility

650 mode enabled for the legacy hook/weight surface) and wraps it for LIT.

651

652 Args:

653 model_name: Name of the pretrained model (e.g., "gpt2").

654 config: Optional wrapper configuration.

655 **model_kwargs: Additional arguments for TransformerBridge.boot_transformers.

656

657 Returns:

658 TransformerLensLIT wrapper instance.

659

660 Example:

661 >>> lit_model = TransformerLensLIT.from_pretrained("gpt2") # doctest: +SKIP

662 """

663 from transformer_lens.model_bridge import TransformerBridge

664

665 model = TransformerBridge.boot_transformers(model_name, **model_kwargs)

666 model.enable_compatibility_mode()

667 return cls(model, config=config)

668

669

670# If LIT is available, register as a proper LIT BatchedModel subclass

671if _LIT_AVAILABLE: 671 ↛ 673line 671 didn't jump to line 673 because the condition on line 671 was never true

672

673 class TransformerLensLITBatched(lit_model.BatchedModel): # type: ignore[union-attr]

674 """Batched version of TransformerLensLIT for better performance.

675

676 This class implements the BatchedModel interface for efficient

677 batch processing. Use this for production deployments.

678 """

679

680 def __init__(

681 self,

682 model: Any,

683 config: Optional[TransformerLensLITConfig] = None,

684 ):

685 """Initialize the batched LIT wrapper.

686

687 Args:

688 model: A TransformerLens model (e.g. TransformerBridge).

689 config: Optional configuration.

690 """

691 # Use the non-batched wrapper internally

692 self._wrapper = TransformerLensLIT(model, config)

693 self.model = model

694 self.config = self._wrapper.config

695

696 def description(self) -> str:

697 return self._wrapper.description()

698

699 @classmethod

700 def init_spec(cls) -> Dict[str, Any]:

701 return TransformerLensLIT.init_spec()

702

703 def input_spec(self) -> Dict[str, Any]:

704 return self._wrapper.input_spec()

705

706 def output_spec(self) -> Dict[str, Any]:

707 return self._wrapper.output_spec()

708

709 def max_minibatch_size(self) -> int:

710 return self._wrapper.max_minibatch_size()

711

712 def predict_minibatch( # type: ignore[union-attr]

713 self,

714 inputs, # type: ignore[override]

715 ):

716 """Run prediction on a minibatch of inputs.

717

718 Args:

719 inputs: List of input dictionaries.

720

721 Returns:

722 List of output dictionaries.

723 """

724 # For now, just iterate (can be optimized for true batching)

725 return [self._wrapper._predict_single(ex) for ex in inputs] # type: ignore[union-attr]

726

727 @classmethod

728 def from_pretrained(

729 cls,

730 model_name: str,

731 config: Optional[TransformerLensLITConfig] = None,

732 **model_kwargs,

733 ) -> "TransformerLensLITBatched":

734 """Create a batched LIT wrapper from a pretrained model.

735

736 Args:

737 model_name: Name of the pretrained model.

738 config: Optional wrapper configuration.

739 **model_kwargs: Additional arguments for model loading.

740

741 Returns:

742 TransformerLensLITBatched instance.

743 """

744 from transformer_lens.model_bridge import TransformerBridge

745

746 model = TransformerBridge.boot_transformers(model_name, **model_kwargs)

747 model.enable_compatibility_mode()

748 return cls(model, config=config)

749

750

751# Legacy aliases (deprecated names kept for one transition release).

752HookedTransformerLITConfig = TransformerLensLITConfig

753HookedTransformerLIT = TransformerLensLIT

754if _LIT_AVAILABLE: 754 ↛ 755line 754 didn't jump to line 755 because the condition on line 754 was never true

755 HookedTransformerLITBatched = TransformerLensLITBatched