This commit is contained in:
isotr0py 2024-11-04 16:19:23 +08:00
parent 205676ceb7
commit dd320df4b4

View file

@ -89,8 +89,9 @@ class GGUFReader:
def __init__(self, path: os.PathLike[str] | str, mode: Literal['r', 'r+', 'c'] = 'r'): def __init__(self, path: os.PathLike[str] | str, mode: Literal['r', 'r+', 'c'] = 'r'):
file_mode = "rb" if mode == 'r' else 'rb+' file_mode = "rb" if mode == 'r' else 'rb+'
self.mode = mode
self.data = open(path, mode=file_mode) self.data = open(path, mode=file_mode)
self.mmap = np.memmap(self.data, mode = mode) self.mmap = np.memmap(path, mode = mode)
offs = 0 offs = 0
# Check for GGUF magic # Check for GGUF magic
@ -150,10 +151,11 @@ class GGUFReader:
itemsize = np.dtype(dtype).itemsize itemsize = np.dtype(dtype).itemsize
if not lazy: if not lazy:
self.data.seek(offset) self.data.seek(offset)
return ( data = (
np.frombuffer(self.data.read(itemsize * count), dtype = dtype, count = count) np.frombuffer(self.data.read(itemsize * count), dtype = dtype, count = count)
.newbyteorder(override_order or self.byte_order) .newbyteorder(override_order or self.byte_order)
) )
return data if self.mode == 'r' else data.copy()
else: else:
return ( return (
self.mmap[offset:offset + itemsize * count] self.mmap[offset:offset + itemsize * count]