(model, pretrained_model, logger)
| 78 | |
| 79 | |
| 80 | def load_pretrained_params(model, pretrained_model, logger): |
| 81 | if pretrained_model.endswith(".safetensors"): |
| 82 | from safetensors.torch import load_file |
| 83 | logger.info(f"Loading weights from safetensors: {pretrained_model}") |
| 84 | checkpoint = load_file(pretrained_model) |
| 85 | else: |
| 86 | logger.info(f"Loading weights using torch.load: {pretrained_model}") |
| 87 | checkpoint = torch.load(pretrained_model, map_location=torch.device("cpu")) |
| 88 | |
| 89 | if "state_dict" in checkpoint: |
| 90 | state_dict = checkpoint["state_dict"] |
| 91 | else: |
| 92 | state_dict = checkpoint |
| 93 | |
| 94 | model.load_state_dict(state_dict, strict=False) |
| 95 | model_keys = model.state_dict().keys() |
| 96 | for name in model_keys: |
| 97 | if name not in state_dict: |
| 98 | logger.info(f"{name} is not in pretrained model") |
| 99 |
no test coverage detected