Reads a checkpoint file, returning properly formatted errors if they arise.
(checkpoint_file: Union[str, os.PathLike], variant: Optional[str] = None)
| 97 | |
| 98 | |
| 99 | def load_state_dict(checkpoint_file: Union[str, os.PathLike], variant: Optional[str] = None): |
| 100 | """ |
| 101 | Reads a checkpoint file, returning properly formatted errors if they arise. |
| 102 | """ |
| 103 | try: |
| 104 | file_extension = os.path.basename(checkpoint_file).split(".")[-1] |
| 105 | if file_extension == SAFETENSORS_FILE_EXTENSION: |
| 106 | return safetensors.torch.load_file(checkpoint_file, device="cpu") |
| 107 | else: |
| 108 | weights_only_kwarg = {"weights_only": True} if is_torch_version(">=", "1.13") else {} |
| 109 | return torch.load( |
| 110 | checkpoint_file, |
| 111 | map_location="cpu", |
| 112 | **weights_only_kwarg, |
| 113 | ) |
| 114 | except Exception as e: |
| 115 | try: |
| 116 | with open(checkpoint_file) as f: |
| 117 | if f.read().startswith("version"): |
| 118 | raise OSError( |
| 119 | "You seem to have cloned a repository without having git-lfs installed. Please install " |
| 120 | "git-lfs and run `git lfs install` followed by `git lfs pull` in the folder " |
| 121 | "you cloned." |
| 122 | ) |
| 123 | else: |
| 124 | raise ValueError( |
| 125 | f"Unable to locate the file {checkpoint_file} which is necessary to load this pretrained " |
| 126 | "model. Make sure you have saved the model properly." |
| 127 | ) from e |
| 128 | except (UnicodeDecodeError, ValueError): |
| 129 | raise OSError( |
| 130 | f"Unable to load weights from checkpoint file for '{checkpoint_file}' " f"at '{checkpoint_file}'. " |
| 131 | ) |
| 132 | |
| 133 | |
| 134 | def load_model_dict_into_meta( |
no test coverage detected