From fb89971e73a50508afcb19fb0aa4afa3a509e74c Mon Sep 17 00:00:00 2001
From: leopf <43857362+leopf@users.noreply.github.com>
Date: Wed, 4 Dec 2024 12:08:54 +0100
Subject: [PATCH] use BufferedReader (#8032)

---
 tinygrad/nn/state.py | 17 ++++++-----------
 1 file changed, 6 insertions(+), 11 deletions(-)

diff --git a/tinygrad/nn/state.py b/tinygrad/nn/state.py
index e29df2dfd5..dfd7d9d4be 100644
--- a/tinygrad/nn/state.py
+++ b/tinygrad/nn/state.py
@@ -297,14 +297,9 @@ def gguf_load(tensor: Tensor) -> Tuple[Dict, Dict[str, Tensor]]:
   kv_data, state_dict = gguf_load(gguf_tensor)
   ```
   """
-  if tensor.dtype != dtypes.uint8 or len(tensor.shape) != 1: raise ValueError("GGUF tensor must be 1d and of dtype uint8!")
-  pos, read_buffer, rb_start, kv_data, state_dict = 0, memoryview(bytes()), 0, {}, {}
-  def read_bytes(n: int):
-    nonlocal pos, read_buffer, rb_start
-    if rb_start + len(read_buffer) < pos + n: rb_start, read_buffer = pos, tensor[pos:(pos+max(n, 1000_000))].data()
-    return read_buffer[pos-rb_start:(pos:=pos+n)-rb_start]
-  def read_unpack(fmt: str, n: int): return struct.unpack(fmt, read_bytes(n))[0]
-  def read_str(): return str(read_bytes(read_uint64()), "utf-8")
+  reader, kv_data, state_dict = io.BufferedReader(TensorIO(tensor), 1000_000), {}, {}
+  def read_unpack(fmt: str, n: int): return struct.unpack(fmt, reader.read(n))[0]
+  def read_str(): return str(reader.read(read_uint64()), "utf-8")
   def read_arr():
     reader, n = readers[read_int32()], read_uint64()
     return [ reader() for _ in range(n) ]
@@ -313,15 +308,15 @@ def gguf_load(tensor: Tensor) -> Tuple[Dict, Dict[str, Tensor]]:
     (1,"b",1), (2,"H",2), (3,"h",2), (4,"I",4), (5,"i",4), (6,"f",4), (7,"?",1), (10,"Q",8), (11,"q",8), (12,"d",8) ] } }
   read_uint32, read_int32, read_uint64, read_int64 = readers[4], readers[5], readers[10], readers[11]
 
-  magic, version, n_tensors, n_kv = read_bytes(4), read_int32(), read_int64(), read_int64()
+  magic, version, n_tensors, n_kv = reader.read(4), read_int32(), read_int64(), read_int64()
   if magic != b"GGUF" or version not in [2, 3]: raise ValueError("Invalid GGUF format!")
   for _ in range(n_kv):
     k, typ = read_str(), read_int32()
     kv_data[k] = readers[typ]()
 
   t_infos = [ (read_str(), tuple(read_uint64() for _ in range(read_uint32())), read_int32(), read_uint64()) for _ in range(n_tensors) ]
-  alignment = kv_data.get("general.alignment", 32)
-  data_start = pos = pos + (alignment - pos % alignment if pos % alignment != 0 else 0)
+  alignment, pos = kv_data.get("general.alignment", 32), reader.tell()
+  data_start = pos + (alignment - pos % alignment if pos % alignment != 0 else 0)
 
   for name, dims, typ, off in t_infos: state_dict[name] = ggml_data_to_tensor(tensor[data_start + off:], prod(dims), typ).reshape(*reversed(dims))