(file_path: str)
| 28 | |
| 29 | |
| 30 | def _load_checkpoint_file(file_path: str) -> tuple[int, dict[str, tuple["FileMeta", torch.Tensor]]]: |
| 31 | def _safetensors_load(fn: str) -> dict[str, tuple["FileMeta", torch.Tensor]]: |
| 32 | ret = {} |
| 33 | with safe_open(fn, framework="pt") as f: |
| 34 | for name in f.keys(): # noqa: SIM118 |
| 35 | weight = f.get_tensor(name) |
| 36 | meta = { |
| 37 | "key": name, |
| 38 | "dtype": weight.dtype, |
| 39 | "shape": weight.shape, |
| 40 | "type": type(weight), |
| 41 | "tp_concat_dim": -1, # safetensors does not support tp_concat_dim |
| 42 | } |
| 43 | ret[name] = (meta, weight) |
| 44 | return ret |
| 45 | |
| 46 | # deprecated, will be removed in the future |
| 47 | def _fast_np_load(fn: str) -> dict[str, tuple["FileMeta", torch.Tensor]]: |
| 48 | """load *.np file and return memmap and related tensor meta""" |
| 49 | |
| 50 | def parse_npy_header(fin: BinaryIO) -> dict[str, Any]: |
| 51 | start = fin.tell() |
| 52 | major, minor = np.lib.format.read_magic(fin) |
| 53 | if major == 1 and minor == 0: |
| 54 | read_header_fn = np.lib.format.read_array_header_1_0 |
| 55 | elif major == 2 and minor == 0: |
| 56 | read_header_fn = np.lib.format.read_array_header_2_0 |
| 57 | else: |
| 58 | raise ValueError( |
| 59 | f"unknown version {major}.{minor} when parsing npy header from {fn}" |
| 60 | ) |
| 61 | shape, is_fortran, dtype = read_header_fn(fin) |
| 62 | return { |
| 63 | "shape": shape, |
| 64 | "is_fortran": is_fortran, |
| 65 | "dtype": dtype, |
| 66 | "header_length": fin.tell() - start, |
| 67 | } |
| 68 | |
| 69 | meta_fn = fn + ".meta" |
| 70 | with open(meta_fn, "rb") as fin: |
| 71 | meta_lst = pickle.load(fin) |
| 72 | |
| 73 | tensors = [] |
| 74 | offset = 0 |
| 75 | with open(fn, "rb") as fin: |
| 76 | fin.seek(0, os.SEEK_END) |
| 77 | filesize = fin.tell() |
| 78 | fin.seek(0) |
| 79 | while fin.tell() < filesize: |
| 80 | tensor_meta = parse_npy_header(fin) |
| 81 | tensor = np.memmap( |
| 82 | fn, |
| 83 | dtype=tensor_meta["dtype"], |
| 84 | mode="c", |
| 85 | offset=offset + tensor_meta["header_length"], |
| 86 | shape=tensor_meta["shape"], |
| 87 | ) |
no test coverage detected