From 5ce2dbcf386e5a2667b0d36a0d6d29f59ab6c423 Mon Sep 17 00:00:00 2001
From: isotr0py <2037008807@qq.com>
Date: Thu, 31 Oct 2024 20:19:24 +0800
Subject: [PATCH] refactor gguf reader

---
 gguf-py/gguf/gguf_reader.py | 30 +++++++++++++++++++-----------
 1 file changed, 19 insertions(+), 11 deletions(-)

diff --git a/gguf-py/gguf/gguf_reader.py b/gguf-py/gguf/gguf_reader.py
index e8e61abf8..92119640b 100644
--- a/gguf-py/gguf/gguf_reader.py
+++ b/gguf-py/gguf/gguf_reader.py
@@ -87,7 +87,8 @@ class GGUFReader:
     }
 
     def __init__(self, path: os.PathLike[str] | str, mode: Literal['r', 'r+', 'c'] = 'r'):
-        self.data = np.memmap(path, mode = mode)
+        self.data = open(path, mode="rb")
+        self.mmap = np.memmap(path, mode = mode)
         offs = 0
 
         # Check for GGUF magic
@@ -127,7 +128,8 @@ class GGUFReader:
         if padding != 0:
             offs += self.alignment - padding
         self.data_offset = offs
-        self._build_tensors(offs, tensors_fields)
+        # self._build_tensors(offs, tensors_fields)
+        self.data.close()
 
     _DT = TypeVar('_DT', bound = npt.DTypeLike)
 
@@ -140,16 +142,22 @@ class GGUFReader:
         return self.tensors[idx]
 
     def _get(
-        self, offset: int, dtype: npt.DTypeLike, count: int = 1, override_order: None | Literal['I', 'S', '<'] = None,
+        self, offset: int, dtype: npt.DTypeLike, count: int = 1, override_order: None | Literal['I', 'S', '<'] = None, lazy: bool = False,
     ) -> npt.NDArray[Any]:
         count = int(count)
-        itemsize = int(np.empty([], dtype = dtype).itemsize)
-        end_offs = offset + itemsize * count
-        return (
-            self.data[offset:end_offs]
-            .view(dtype = dtype)[:count]
-            .newbyteorder(override_order or self.byte_order)
-        )
+        itemsize = np.dtype(dtype).itemsize
+        if not lazy:
+            self.data.seek(offset)
+            return (
+                np.frombuffer(self.data.read(itemsize * count), dtype = dtype, count = count)
+                .newbyteorder(override_order or self.byte_order)
+            )
+        else:
+            return (
+                self.mmap[offset:offset + itemsize * count]
+                .view(dtype = dtype)[:count]
+                .newbyteorder(override_order or self.byte_order)
+            )
 
     def _push_field(self, field: ReaderField, skip_sum: bool = False) -> int:
         if field.name in self.fields:
@@ -311,7 +319,7 @@ class GGUFReader:
                 n_elements = n_elems,
                 n_bytes = n_bytes,
                 data_offset = data_offs,
-                data = self._get(data_offs, item_type, item_count).reshape(np_dims),
+                data = self._get(data_offs, item_type, item_count, lazy=True).reshape(np_dims),
                 field = field,
             ))
         self.tensors = tensors