| """ | |
| Pure python version of Safetensors safe_open | |
| From https://gist.github.com/Narsil/3edeec2669a5e94e4707aa0f901d2282 | |
| """ | |
| import json | |
| import mmap | |
| import os | |
| import torch | |
| class SafetensorsWrapper: | |
| def __init__(self, metadata, tensors): | |
| self._metadata = metadata | |
| self._tensors = tensors | |
| def metadata(self): | |
| return self._metadata | |
| def keys(self): | |
| return self._tensors.keys() | |
| def get_tensor(self, k): | |
| return self._tensors[k] | |
| DTYPES = { | |
| "F32": torch.float32, | |
| "F16": torch.float16, | |
| "BF16": torch.bfloat16, | |
| } | |
| def create_tensor(storage, info, offset): | |
| dtype = DTYPES[info["dtype"]] | |
| shape = info["shape"] | |
| start, stop = info["data_offsets"] | |
| return ( | |
| torch.asarray(storage[start + offset : stop + offset], dtype=torch.uint8) | |
| .view(dtype=dtype) | |
| .reshape(shape) | |
| ) | |
| def safe_open(filename, framework="pt", device="cpu"): | |
| if framework != "pt": | |
| raise ValueError("`framework` must be 'pt'") | |
| with open(filename, mode="r", encoding="utf8") as file_obj: | |
| with mmap.mmap(file_obj.fileno(), length=0, access=mmap.ACCESS_READ) as m: | |
| header = m.read(8) | |
| n = int.from_bytes(header, "little") | |
| metadata_bytes = m.read(n) | |
| metadata = json.loads(metadata_bytes) | |
| size = os.stat(filename).st_size | |
| storage = torch.ByteStorage.from_file(filename, shared=False, size=size).untyped() | |
| offset = n + 8 | |
| return SafetensorsWrapper( | |
| metadata=metadata.get("__metadata__", {}), | |
| tensors={ | |
| name: create_tensor(storage, info, offset).to(device) | |
| for name, info in metadata.items() | |
| if name != "__metadata__" | |
| }, | |
| ) | |