Coverage for transformer_lens/supported

1# HookedTransformers model registry. If you are looking for TransformerBridge's supported models

2# see transformer_lens/tools/model_registry/data/supported_models.json or the docs page at

3# https://transformerlensorg.github.io/TransformerLens/generated/transformer_bridge_models.html

4from transformer_lens.tools.model_registry.registry_io import load_model_aliases

6OFFICIAL_MODEL_NAMES: list[str] = [

7 "01-ai/Yi-34B",

8 "01-ai/Yi-34B-Chat",

9 "01-ai/Yi-6B",

10 "01-ai/Yi-6B-Chat",

11 "ai-forever/mGPT",

12 "allenai/OLMo-1B-hf",

13 "allenai/OLMo-2-0425-1B",

14 "allenai/OLMo-2-1124-7B",

15 "allenai/Olmo-3-1025-7B",

16 "allenai/Olmo-3-1125-32B",

17 "allenai/Olmo-3-32B-Think",

18 "allenai/Olmo-3-7B-Instruct",

19 "allenai/Olmo-3-7B-Think",

20 "allenai/Olmo-3.1-32B-Instruct",

21 "allenai/Olmo-3.1-32B-Think",

22 "allenai/OLMo-7B-hf",

23 "allenai/OLMoE-1B-7B-0924",

24 "ArthurConmy/redwood_attn_2l",

25 "Baidicoot/Othello-GPT-Transformer-Lens",

26 "bigcode/santacoder",

27 "bigscience/bloom-1b1",

28 "bigscience/bloom-1b7",

29 "bigscience/bloom-3b",

30 "bigscience/bloom-560m",

31 "bigscience/bloom-7b1",

32 "codellama/CodeLlama-7b-hf",

33 "codellama/CodeLlama-7b-Instruct-hf",

34 "codellama/CodeLlama-7b-Python-hf",

35 "distilgpt2",

36 "EleutherAI/gpt-j-6B",

37 "EleutherAI/gpt-neo-1.3B",

38 "EleutherAI/gpt-neo-125M",

39 "EleutherAI/gpt-neo-2.7B",

40 "EleutherAI/gpt-neox-20b",

41 "EleutherAI/pythia-1.4b",

42 "EleutherAI/pythia-1.4b-deduped",

43 "EleutherAI/pythia-1.4b-deduped-v0",

44 "EleutherAI/pythia-1.4b-v0",

45 "EleutherAI/pythia-12b",

46 "EleutherAI/pythia-12b-deduped",

47 "EleutherAI/pythia-12b-deduped-v0",

48 "EleutherAI/pythia-12b-v0",

49 "EleutherAI/pythia-14m",

50 "EleutherAI/pythia-160m",

51 "EleutherAI/pythia-160m-deduped",

52 "EleutherAI/pythia-160m-deduped-v0",

53 "EleutherAI/pythia-160m-seed1",

54 "EleutherAI/pythia-160m-seed2",

55 "EleutherAI/pythia-160m-seed3",

56 "EleutherAI/pythia-160m-v0",

57 "EleutherAI/pythia-1b",

58 "EleutherAI/pythia-1b-deduped",

59 "EleutherAI/pythia-1b-deduped-v0",

60 "EleutherAI/pythia-1b-v0",

61 "EleutherAI/pythia-2.8b",

62 "EleutherAI/pythia-2.8b-deduped",

63 "EleutherAI/pythia-2.8b-deduped-v0",

64 "EleutherAI/pythia-2.8b-v0",

65 "EleutherAI/pythia-31m",

66 "EleutherAI/pythia-410m",

67 "EleutherAI/pythia-410m-deduped",

68 "EleutherAI/pythia-410m-deduped-v0",

69 "EleutherAI/pythia-410m-v0",

70 "EleutherAI/pythia-6.9b",

71 "EleutherAI/pythia-6.9b-deduped",

72 "EleutherAI/pythia-6.9b-deduped-v0",

73 "EleutherAI/pythia-6.9b-v0",

74 "EleutherAI/pythia-70m",

75 "EleutherAI/pythia-70m-deduped",

76 "EleutherAI/pythia-70m-deduped-v0",

77 "EleutherAI/pythia-70m-v0",

78 "facebook/hubert-base-ls960",

79 "facebook/opt-1.3b",

80 "facebook/opt-125m",

81 "facebook/opt-13b",

82 "facebook/opt-2.7b",

83 "facebook/opt-30b",

84 "facebook/opt-6.7b",

85 "facebook/opt-66b",

86 "facebook/wav2vec2-base",

87 "facebook/wav2vec2-large",

88 "google-bert/bert-base-cased",

89 "google-bert/bert-base-uncased",

90 "google-bert/bert-large-cased",

91 "google-bert/bert-large-uncased",

92 "google-t5/t5-base",

93 "google-t5/t5-large",

94 "google-t5/t5-small",

95 "google/gemma-2-27b",

96 "google/gemma-2-27b-it",

97 "google/gemma-2-2b",

98 "google/gemma-2-2b-it",

99 "google/gemma-2-9b",

100 "google/gemma-2-9b-it",

101 "google/gemma-2b",

102 "google/gemma-2b-it",

103 "google/gemma-3-12b-it",

104 "google/gemma-3-12b-pt",

105 "google/gemma-3-1b-it",

106 "google/gemma-3-1b-pt",

107 "google/gemma-3-270m",

108 "google/gemma-3-270m-it",

109 "google/gemma-3-27b-it",

110 "google/gemma-3-27b-pt",

111 "google/gemma-3-4b-it",

112 "google/gemma-3-4b-pt",

113 "google/gemma-7b",

114 "google/gemma-7b-it",

115 "google/medgemma-27b-it",

116 "google/medgemma-27b-text-it",

117 "google/medgemma-4b-it",

118 "google/medgemma-4b-pt",

119 "gpt2",

120 "gpt2-large",

121 "gpt2-medium",

122 "gpt2-xl",

123 "llama-13b-hf",

124 "llama-30b-hf",

125 "llama-65b-hf",

126 "llama-7b-hf",

127 "meta-llama/Llama-2-13b-chat-hf",

128 "meta-llama/Llama-2-13b-hf",

129 "meta-llama/Llama-2-70b-chat-hf",

130 "meta-llama/Llama-2-7b-chat-hf",

131 "meta-llama/Llama-2-7b-hf",

132 "meta-llama/Llama-3.1-70B",

133 "meta-llama/Llama-3.1-70B-Instruct",

134 "meta-llama/Llama-3.1-8B",

135 "meta-llama/Llama-3.1-8B-Instruct",

136 "meta-llama/Llama-3.2-1B",

137 "meta-llama/Llama-3.2-1B-Instruct",

138 "meta-llama/Llama-3.2-3B",

139 "meta-llama/Llama-3.2-3B-Instruct",

140 "meta-llama/Llama-3.3-70B-Instruct",

141 "meta-llama/Meta-Llama-3-70B",

142 "meta-llama/Meta-Llama-3-70B-Instruct",

143 "meta-llama/Meta-Llama-3-8B",

144 "meta-llama/Meta-Llama-3-8B-Instruct",

145 "microsoft/phi-1",

146 "microsoft/phi-1_5",

147 "microsoft/phi-2",

148 "microsoft/Phi-3-mini-4k-instruct",

149 "microsoft/phi-4",

150 "mistralai/Mistral-7B-Instruct-v0.1",

151 "mistralai/Mistral-7B-v0.1",

152 "mistralai/Mistral-Nemo-Base-2407",

153 "mistralai/Mistral-Small-24B-Base-2501",

154 "mistralai/Mixtral-8x7B-Instruct-v0.1",

155 "mistralai/Mixtral-8x7B-v0.1",

156 "NeelNanda/Attn-Only-2L512W-Shortformer-6B-big-lr",

157 "NeelNanda/Attn_Only_1L512W_C4_Code",

158 "NeelNanda/Attn_Only_2L512W_C4_Code",

159 "NeelNanda/Attn_Only_3L512W_C4_Code",

160 "NeelNanda/Attn_Only_4L512W_C4_Code",

161 "NeelNanda/GELU_1L512W_C4_Code",

162 "NeelNanda/GELU_2L512W_C4_Code",

163 "NeelNanda/GELU_3L512W_C4_Code",

164 "NeelNanda/GELU_4L512W_C4_Code",

165 "NeelNanda/SoLU_10L1280W_C4_Code",

166 "NeelNanda/SoLU_10L_v22_old",

167 "NeelNanda/SoLU_12L1536W_C4_Code",

168 "NeelNanda/SoLU_12L_v23_old",

169 "NeelNanda/SoLU_1L512W_C4_Code",

170 "NeelNanda/SoLU_1L512W_Wiki_Finetune",

171 "NeelNanda/SoLU_1L_v9_old",

172 "NeelNanda/SoLU_2L512W_C4_Code",

173 "NeelNanda/SoLU_2L_v10_old",

174 "NeelNanda/SoLU_3L512W_C4_Code",

175 "NeelNanda/SoLU_4L512W_C4_Code",

176 "NeelNanda/SoLU_4L512W_Wiki_Finetune",

177 "NeelNanda/SoLU_4L_v11_old",

178 "NeelNanda/SoLU_6L768W_C4_Code",

179 "NeelNanda/SoLU_6L_v13_old",

180 "NeelNanda/SoLU_8L1024W_C4_Code",

181 "NeelNanda/SoLU_8L_v21_old",

182 "openai/gpt-oss-20b",

183 "Qwen/Qwen-14B",

184 "Qwen/Qwen-14B-Chat",

185 "Qwen/Qwen-1_8B",

186 "Qwen/Qwen-1_8B-Chat",

187 "Qwen/Qwen-7B",

188 "Qwen/Qwen-7B-Chat",

189 "Qwen/Qwen1.5-0.5B",

190 "Qwen/Qwen1.5-0.5B-Chat",

191 "Qwen/Qwen1.5-1.8B",

192 "Qwen/Qwen1.5-1.8B-Chat",

193 "Qwen/Qwen1.5-14B",

194 "Qwen/Qwen1.5-14B-Chat",

195 "Qwen/Qwen1.5-4B",

196 "Qwen/Qwen1.5-4B-Chat",

197 "Qwen/Qwen1.5-7B",

198 "Qwen/Qwen1.5-7B-Chat",

199 "Qwen/Qwen2-0.5B",

200 "Qwen/Qwen2-0.5B-Instruct",

201 "Qwen/Qwen2-1.5B",

202 "Qwen/Qwen2-1.5B-Instruct",

203 "Qwen/Qwen2-7B",

204 "Qwen/Qwen2-7B-Instruct",

205 "Qwen/Qwen2.5-0.5B",

206 "Qwen/Qwen2.5-0.5B-Instruct",

207 "Qwen/Qwen2.5-1.5B",

208 "Qwen/Qwen2.5-1.5B-Instruct",

209 "Qwen/Qwen2.5-14B",

210 "Qwen/Qwen2.5-14B-Instruct",

211 "Qwen/Qwen2.5-32B",

212 "Qwen/Qwen2.5-32B-Instruct",

213 "Qwen/Qwen2.5-3B",

214 "Qwen/Qwen2.5-3B-Instruct",

215 "Qwen/Qwen2.5-72B",

216 "Qwen/Qwen2.5-72B-Instruct",

217 "Qwen/Qwen2.5-7B",

218 "Qwen/Qwen2.5-7B-Instruct",

219 "Qwen/Qwen3-0.6B",

220 "Qwen/Qwen3-0.6B-Base",

221 "Qwen/Qwen3-1.7B",

222 "Qwen/Qwen3-14B",

223 "Qwen/Qwen3-4B",

224 "Qwen/Qwen3-8B",

225 "Qwen/QwQ-32B-Preview",

226 "roneneldan/TinyStories-1Layer-21M",

227 "roneneldan/TinyStories-1M",

228 "roneneldan/TinyStories-28M",

229 "roneneldan/TinyStories-2Layers-33M",

230 "roneneldan/TinyStories-33M",

231 "roneneldan/TinyStories-3M",

232 "roneneldan/TinyStories-8M",

233 "roneneldan/TinyStories-Instruct-1M",

234 "roneneldan/TinyStories-Instruct-28M",

235 "roneneldan/TinyStories-Instruct-2Layers-33M",

236 "roneneldan/TinyStories-Instruct-33M",

237 "roneneldan/TinyStories-Instruct-3M",

238 "roneneldan/TinyStories-Instruct-8M",

239 "roneneldan/TinyStories-Instuct-1Layer-21M",

240 "stabilityai/stablelm-base-alpha-3b",

241 "stabilityai/stablelm-base-alpha-7b",

242 "stabilityai/stablelm-tuned-alpha-3b",

243 "stabilityai/stablelm-tuned-alpha-7b",

244 "stanford-crfm/alias-gpt2-small-x21",

245 "stanford-crfm/arwen-gpt2-medium-x21",

246 "stanford-crfm/battlestar-gpt2-small-x49",

247 "stanford-crfm/beren-gpt2-medium-x49",

248 "stanford-crfm/caprica-gpt2-small-x81",

249 "stanford-crfm/celebrimbor-gpt2-medium-x81",

250 "stanford-crfm/darkmatter-gpt2-small-x343",

251 "stanford-crfm/durin-gpt2-medium-x343",

252 "stanford-crfm/eowyn-gpt2-medium-x777",

253 "stanford-crfm/expanse-gpt2-small-x777",

254 "swiss-ai/Apertus-8B-2509",

255 "swiss-ai/Apertus-8B-Instruct-2509",

256]

257"""Official model names for models on HuggingFace."""

258

259# Model Aliases:

260# Canonical data lives in the bridge model registry

261# (tools/model_registry/data/model_aliases.json); re-exported here for the

262# frozen legacy ledger and get_official_model_name.

263MODEL_ALIASES: dict[str, list[str]] = load_model_aliases()

264"""Model aliases for models on HuggingFace."""

265

266

267# Sets a default model alias, by convention the first one in the model alias table, else the official name if it has no aliases

268DEFAULT_MODEL_ALIASES: list[str] = [

269 MODEL_ALIASES[name][0] if name in MODEL_ALIASES else name for name in OFFICIAL_MODEL_NAMES

270]

271

272

273def _model_alias_map() -> dict[str, str]:

274 """Lowercased-key map from every alias (and official name) to its official

275 name — mirrors the deleted loading_from_pretrained.make_model_alias_map so

276 resolution stays case-insensitive."""

277 alias_map: dict[str, str] = {}

278 for official_name in OFFICIAL_MODEL_NAMES:

279 for alias in MODEL_ALIASES.get(official_name, []):

280 alias_map[alias.lower()] = official_name

281 alias_map[official_name.lower()] = official_name

282 return alias_map

283

284

285def get_official_model_name(model_name: str) -> str:

286 """Resolve a HookedTransformer-era alias to its official HF name.

287

288 Case-insensitive (as the deleted ``loading_from_pretrained`` resolver was);

289 identity for an already-official name; raises for an unknown one. Rehomed so

290 the legacy-compatibility ledger and alias-drift tooling keep a canonical

291 resolver.

292 """

293 official_name = _model_alias_map().get(model_name.lower())

294 if official_name is None:

295 raise ValueError(f"{model_name!r} is not an official model name or a known alias.")

296 return official_name

Coverage for transformer_lens/supported_models.py: 100%

18 statements