Load model, add forward method, and move to device. Args: model_name_or_checkpoint_path: Name or path of pretrained model or checkpoint. device: Device to move the model to. decoder_loss_weight: Decoder loss weight.
(
configs,
checkpoint_path: Union[str, Path] = None,
device: Union[str, torch.device] = "cuda",
)
| 223 | |
| 224 | |
| 225 | def load_model( |
| 226 | configs, |
| 227 | checkpoint_path: Union[str, Path] = None, |
| 228 | device: Union[str, torch.device] = "cuda", |
| 229 | ) -> Model: |
| 230 | """Load model, add forward method, and move to device. |
| 231 | |
| 232 | Args: |
| 233 | model_name_or_checkpoint_path: Name or path of pretrained model or checkpoint. |
| 234 | device: Device to move the model to. |
| 235 | decoder_loss_weight: Decoder loss weight. |
| 236 | """ |
| 237 | |
| 238 | model_arg = ModelArgs( |
| 239 | backbone_flavor=configs["models"]["backbone_flavor"], |
| 240 | decoder_flavor=configs["models"]["decoder_flavor"], |
| 241 | text_vocab_size=configs["models"]["text_vocab_size"], |
| 242 | audio_vocab_size=configs["models"]["audio_vocab_size"], |
| 243 | audio_num_codebooks=configs["models"]["audio_num_codebooks"], |
| 244 | decoder_loss_weight=configs["models"]["decoder_loss_weight"], |
| 245 | use_text_loss=True, |
| 246 | ) |
| 247 | model = Model(model_arg) |
| 248 | |
| 249 | if checkpoint_path and os.path.exists(checkpoint_path): |
| 250 | state_dict = torch.load( |
| 251 | checkpoint_path, map_location="cpu", weights_only=False |
| 252 | )["model"] |
| 253 | model.load_state_dict(state_dict) |
| 254 | else: |
| 255 | model = init_weights(model) |
| 256 | |
| 257 | model = model.to(device=device) |
| 258 | return model |
| 259 | |
| 260 | |
| 261 | def load_llm_model( |
no test coverage detected