| 237 | self.initialize_others() |
| 238 | |
| 239 | def initialize_esm_model(self): |
| 240 | # load ESM2 model |
| 241 | esm_model, esm_dict = esm_registry["esm2_3B"]() |
| 242 | af2_to_esm = _af2_to_esm(esm_dict) |
| 243 | |
| 244 | if self.backend == "torch": |
| 245 | esm_model = esm_model.to(self.device) |
| 246 | af2_to_esm = af2_to_esm.to(self.device) |
| 247 | elif self.backend == "mlx": |
| 248 | esm_model_mlx = ESM2MLX(num_layers=36, embed_dim=2560, attention_heads=40) |
| 249 | esm_state_dict_torch = esm_model.cpu().state_dict() |
| 250 | |
| 251 | esm_state_dict_torch = { |
| 252 | k: mx.array(v) |
| 253 | for k, v in starmap(map_torch_to_mlx, esm_state_dict_torch.items()) |
| 254 | if k is not None |
| 255 | } |
| 256 | esm_model_mlx.update(tree_unflatten(list(esm_state_dict_torch.items()))) |
| 257 | esm_model = esm_model_mlx |
| 258 | print(f"pLM ESM-3B loaded with {self.backend} backend.") |
| 259 | |
| 260 | self.esm_model = esm_model.eval() |
| 261 | self.esm_dict = esm_dict |
| 262 | self.af2_to_esm = af2_to_esm |
| 263 | |
| 264 | def initialize_others(self): |
| 265 | # prepare data tokenizer, featurizer, and processor |