Coverage for transformer_lens/lit/__init__.py: 25%

1"""LIT (Learning Interpretability Tool) integration for TransformerLens. 

2 

3This module provides integration between TransformerLens and Google's Learning 

4Interpretability Tool (LIT), enabling interactive visualization and analysis 

5of transformer models. 

6 

7Quick Start: 

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

9 >>> from transformer_lens.lit import TransformerLensLIT, SimpleTextDataset, serve # doctest: +SKIP 

10 >>> 

11 >>> # Load model and create LIT wrapper 

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

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

14 >>> 

15 >>> # Create a dataset 

16 >>> dataset = SimpleTextDataset.from_strings([ # doctest: +SKIP 

17 ... "The capital of France is Paris.", 

18 ... "Machine learning is a field of AI.", 

19 ... ]) 

20 >>> 

21 >>> # Start LIT server 

22 >>> serve({"gpt2": lit_model}, {"examples": dataset}) # doctest: +SKIP 

23 

24For Colab/Jupyter notebooks: 

25 >>> from transformer_lens.lit import LITWidget # doctest: +SKIP 

26 >>> 

27 >>> widget = LITWidget({"gpt2": lit_model}, {"examples": dataset}) # doctest: +SKIP 

28 >>> widget.render() # doctest: +SKIP 

29 

30Features: 

31 - Interactive token predictions and top-k analysis 

32 - Attention pattern visualization across all layers and heads 

33 - Embedding projector for layer-wise representations 

34 - Token salience/gradient visualization 

35 - Support for IOI and Induction datasets 

36 

37Requirements: 

38 - lit-nlp >= 1.0 (install with: pip install lit-nlp) 

39 

40References: 

41 - LIT: https://pair-code.github.io/lit/ 

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

43 

44Note: 

45 This module requires the optional `lit-nlp` dependency. Install it with: 

46 ``` 

47 pip install lit-nlp 

48 ``` 

49 or 

50 ``` 

51 pip install transformer-lens[lit] 

52 ``` 

53""" 

54 

55from __future__ import annotations 

56 

57import logging 

58from typing import Any, Dict, Union 

59 

60# Check if LIT is available 

61from .utils import check_lit_installed 

62 

63__all__ = [ 

64 # Model wrappers 

65 "HookedTransformerLIT", 

66 "HookedTransformerLITConfig", 

67 "TransformerLensLIT", 

68 "TransformerLensLITBatched", 

69 "TransformerLensLITConfig", 

70 # Datasets 

71 "SimpleTextDataset", 

72 "PromptCompletionDataset", 

73 "IOIDataset", 

74 "InductionDataset", 

75 "wrap_for_lit", 

76 # Server utilities 

77 "serve", 

78 "LITWidget", 

79 # Constants 

80 "INPUT_FIELDS", 

81 "OUTPUT_FIELDS", 

82 # Utilities 

83 "check_lit_installed", 

84] 

85 

86logger = logging.getLogger(__name__) 

87 

88# Import constants (always available) 

89from .constants import ERRORS, INPUT_FIELDS, OUTPUT_FIELDS, SERVER_CONFIG # noqa: E402 

90 

91# Import datasets (handles LIT availability internally) 

92from .dataset import ( # noqa: E402 

93 InductionDataset, 

94 IOIDataset, 

95 PromptCompletionDataset, 

96 SimpleTextDataset, 

97 wrap_for_lit, 

98) 

99 

100# Import model wrapper (handles LIT availability internally) 

101from .model import ( # noqa: E402 

102 HookedTransformerLIT, 

103 HookedTransformerLITConfig, 

104 TransformerLensLIT, 

105 TransformerLensLITConfig, 

106) 

107 

108# Conditional imports that require LIT 

109_LIT_AVAILABLE = check_lit_installed() 

110 

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

112 from .model import TransformerLensLITBatched # noqa: E402 

113else: 

114 TransformerLensLITBatched = None # type: ignore[misc, assignment] 

115 

116 

117def serve( 

118 models: Union[Dict[str, Any], Any], 

119 datasets: Union[Dict[str, Any], Any], 

120 port: int = SERVER_CONFIG.DEFAULT_PORT, 

121 host: str = SERVER_CONFIG.DEFAULT_HOST, 

122 page_title: str = SERVER_CONFIG.DEFAULT_TITLE, 

123 **kwargs, 

124) -> None: 

125 """Start a LIT server with the given models and datasets. 

126 

127 This is a convenience function to quickly start a LIT server 

128 for interactive model exploration. 

129 

130 Args: 

131 models: Either a single TransformerLens model/TransformerLensLIT, or 

132 a dictionary mapping model names to model wrappers. 

133 datasets: Either a single dataset, or a dictionary mapping 

134 dataset names to datasets. 

135 port: Port number for the server. 

136 host: Host address for the server. 

137 page_title: Title shown in the browser tab. 

138 **kwargs: Additional arguments passed to LIT server. 

139 

140 Example: 

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

142 >>> from transformer_lens.lit import SimpleTextDataset, serve # doctest: +SKIP 

143 >>> 

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

145 >>> dataset = SimpleTextDataset.from_strings(["Hello world!"]) # doctest: +SKIP 

146 >>> 

147 >>> # Simple usage with single model and dataset 

148 >>> serve(model, dataset) # doctest: +SKIP 

149 >>> 

150 >>> # Or with explicit names 

151 >>> serve({"gpt2": model}, {"examples": dataset}) # doctest: +SKIP 

152 

153 Note: 

154 This function will block and run the server. Press Ctrl+C to stop. 

155 """ 

156 if not _LIT_AVAILABLE: 

157 raise ImportError(ERRORS.LIT_NOT_INSTALLED) 

158 

159 from lit_nlp import dev_server 

160 

161 # Handle single model vs dictionary of models 

162 if not isinstance(models, dict): 

163 # Single model passed - check if it's a TransformerLens model that needs wrapping 

164 model = models 

165 if hasattr(model, "cfg") and hasattr(model, "run_with_cache"): 

166 # It's a TransformerLens model, wrap it 

167 model = TransformerLensLIT(model) 

168 models = {"model": model} 

169 

170 # Handle single dataset vs dictionary of datasets 

171 if not isinstance(datasets, dict): 

172 datasets = {"dataset": datasets} 

173 

174 # Wrap datasets if needed 

175 wrapped_datasets = {} 

176 for name, dataset in datasets.items(): 

177 if hasattr(dataset, "_examples"): 

178 # Our custom dataset, wrap it 

179 wrapped_datasets[name] = wrap_for_lit(dataset) 

180 else: 

181 # Already a LIT dataset 

182 wrapped_datasets[name] = dataset 

183 

184 # Get the LIT client root path and layout 

185 import os 

186 

187 import lit_nlp 

188 from lit_nlp.api import layout as lit_layout 

189 

190 client_root = os.path.join(os.path.dirname(lit_nlp.__file__), "client", "build", "default") 

191 

192 # Use default layouts if not provided 

193 if "layouts" not in kwargs: 

194 kwargs["layouts"] = lit_layout.DEFAULT_LAYOUTS 

195 if "default_layout" not in kwargs: 

196 kwargs["default_layout"] = "default" 

197 

198 # Create and start server 

199 server = dev_server.Server( 

200 models, 

201 wrapped_datasets, 

202 port=port, 

203 host=host, 

204 page_title=page_title, 

205 client_root=client_root, 

206 **kwargs, 

207 ) 

208 

209 logger.info(f"Starting LIT server at http://{host}:{port}") 

210 server.serve() 

211 

212 

213class LITWidget: 

214 """LIT Widget for Jupyter/Colab notebooks. 

215 

216 This class provides an easy way to use LIT within notebook environments 

217 without needing to run a separate server. 

218 

219 Example: 

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

221 >>> from transformer_lens.lit import TransformerLensLIT, SimpleTextDataset, LITWidget # doctest: +SKIP 

222 >>> 

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

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

225 >>> dataset = SimpleTextDataset.from_strings(["Hello world!"]) # doctest: +SKIP 

226 >>> 

227 >>> widget = LITWidget({"gpt2": lit_model}, {"examples": dataset}) # doctest: +SKIP 

228 >>> widget.render() # Displays in the notebook # doctest: +SKIP 

229 

230 Note: 

231 VSCode notebooks don't support iframe rendering. Use `widget.url` to 

232 get the URL and open it manually in your browser. 

233 """ 

234 

235 def __init__( 

236 self, 

237 models: Dict[str, Any], 

238 datasets: Dict[str, Any], 

239 height: int = 800, 

240 **kwargs, 

241 ): 

242 """Initialize the LIT widget. 

243 

244 Args: 

245 models: Dictionary mapping model names to model wrappers. 

246 datasets: Dictionary mapping dataset names to datasets. 

247 height: Height of the widget in pixels. 

248 **kwargs: Additional arguments for the LIT widget. 

249 """ 

250 if not _LIT_AVAILABLE: 

251 raise ImportError(ERRORS.LIT_NOT_INSTALLED) 

252 

253 from lit_nlp import notebook 

254 

255 # Wrap datasets if needed 

256 wrapped_datasets = {} 

257 for name, dataset in datasets.items(): 

258 if hasattr(dataset, "_examples"): 

259 wrapped_datasets[name] = wrap_for_lit(dataset) 

260 else: 

261 wrapped_datasets[name] = dataset 

262 

263 # LitWidget expects models and datasets as positional args 

264 # Remove default_layout from kwargs as it's handled internally by LitWidget 

265 kwargs.pop("default_layout", None) 

266 

267 self._widget = notebook.LitWidget( 

268 models, 

269 wrapped_datasets, 

270 height=height, 

271 render=False, # Don't auto-render 

272 **kwargs, 

273 ) 

274 

275 @property 

276 def url(self) -> str: 

277 """Get the URL of the LIT server. 

278 

279 Use this to manually open LIT in a browser when notebook 

280 rendering doesn't work (e.g., in VSCode). 

281 

282 Returns: 

283 The URL to access the LIT UI. 

284 """ 

285 port = self._widget._server.port 

286 return f"http://localhost:{port}" 

287 

288 def render(self, open_in_new_tab: bool = False, **kwargs): 

289 """Render the LIT widget. 

290 

291 Args: 

292 open_in_new_tab: If True, opens in a new browser tab. 

293 **kwargs: Additional render arguments. 

294 

295 Note: 

296 If rendering doesn't work in your environment (e.g., VSCode), 

297 use `print(widget.url)` and open that URL in your browser. 

298 """ 

299 self._widget.render(open_in_new_tab=open_in_new_tab, **kwargs) 

300 

301 def stop(self): 

302 """Stop the widget's server and free resources.""" 

303 self._widget.stop() 

304 

305 

306# Version info 

307__version__ = "1.0.0"