Merge branch 'gguf-no-tempfile'

Author: cebtenzzre

convert-llama-ggml-to-gguf.py

@@ -235,10 +235,7 @@ def __init__(self, ggml_model, data, cfg, params_override = None, vocab_override
 
     def save(self):
         print('* Preparing to save GGUF file')
-        gguf_writer = gguf.GGUFWriter(
-            self.cfg.output,
-            gguf.MODEL_ARCH_NAMES[gguf.MODEL_ARCH.LLAMA],
-            use_temp_file = False )
+        gguf_writer = gguf.GGUFWriter(self.cfg.output, gguf.MODEL_ARCH_NAMES[gguf.MODEL_ARCH.LLAMA])
         self.add_params(gguf_writer)
         self.add_vocab(gguf_writer)
         if self.special_vocab is not None:

examples/finetune/convert-finetune-checkpoint-to-gguf.py

@@ -475,7 +475,7 @@ def main():
     # we should have read all available data
     assert(offset == len(data))
 
-    gguf_writer = gguf.GGUFWriter(cfg.output, gguf.MODEL_ARCH_NAMES[gguf.MODEL_ARCH.LLAMA], use_temp_file = False)
+    gguf_writer = gguf.GGUFWriter(cfg.output, gguf.MODEL_ARCH_NAMES[gguf.MODEL_ARCH.LLAMA])
     chk.save_gguf(gguf_writer)
     print("    gguf: write header")
     gguf_writer.write_header_to_file()

examples/train-text-from-scratch/convert-train-checkpoint-to-gguf.py

@@ -485,7 +485,7 @@ def main():
     # we should have read all available data
     assert(offset == len(data))
 
-    gguf_writer = gguf.GGUFWriter(cfg.output, gguf.MODEL_ARCH_NAMES[gguf.MODEL_ARCH.LLAMA], use_temp_file = False)
+    gguf_writer = gguf.GGUFWriter(cfg.output, gguf.MODEL_ARCH_NAMES[gguf.MODEL_ARCH.LLAMA])
     chk.save_gguf(gguf_writer)
     print("    gguf: write header")
     gguf_writer.write_header_to_file()

gguf-py/gguf/gguf.py

@@ -636,18 +636,16 @@ def get_type(val):
             sys.exit()
 
 
+class WriterState:
+    EMPTY   = auto()
+    HEADER  = auto()
+    KV_DATA = auto()
+    TI_DATA = auto()
+
+
 class GGUFWriter:
     fout: BufferedWriter
-    arch: str
-    offset_tensor = 0
-    data_alignment = GGUF_DEFAULT_ALIGNMENT
-    kv_data = b""
-    kv_data_count = 0
-    ti_data = b""
-    ti_data_count = 0
-    use_temp_file: bool
-    temp_file: tempfile.SpooledTemporaryFile[bytes] | None = None
-    tensors: list[tuple[np.ndarray[Any, Any], int]]
+    tensors: list[np.ndarray[Any, Any]]
 
     @property
     def pack_prefix(self):
@@ -656,9 +654,15 @@ def pack_prefix(self):
         else:
             return ">"
 
-    def __init__(self, path: os.PathLike[str] | str, arch: str, use_temp_file = True, endianess=GGUFEndian.LITTLE):
+    def __init__(self, path: os.PathLike[str] | str, arch: str, endianess=GGUFEndian.LITTLE):
         self.fout = open(path, "wb")
         self.arch = arch
+        self.offset_tensor = 0
+        self.data_alignment = GGUF_DEFAULT_ALIGNMENT
+        self.kv_data = b""
+        self.kv_data_count = 0
+        self.ti_data = b""
+        self.ti_data_count = 0
         self.endianess = endianess
         self._simple_value_packing = {
             GGUFValueType.UINT8:   f"{self.pack_prefix}B",
@@ -673,27 +677,41 @@ def __init__(self, path: os.PathLike[str] | str, arch: str, use_temp_file = True
             GGUFValueType.FLOAT64: f"{self.pack_prefix}d",
             GGUFValueType.BOOL:    "?" ,
         }
-        self.add_architecture()
-        self.use_temp_file = use_temp_file
         self.tensors = []
+        self.state = WriterState.EMPTY
+
+        self.add_architecture()
+
         endianess_str = "Big Endian" if self.endianess == GGUFEndian.BIG else "Little Endian"
         print(f"This gguf file is for {endianess_str} only")
 
     def write_header_to_file(self):
+        if self.state is not WriterState.EMPTY:
+            raise ValueError(f'Expected output file to be empty, got {self.state}')
+
         self.fout.write(struct.pack("<I", GGUF_MAGIC))
         self.fout.write(struct.pack(f"{self.pack_prefix}I", GGUF_VERSION))
         self.fout.write(struct.pack(f"{self.pack_prefix}Q", self.ti_data_count))
         self.fout.write(struct.pack(f"{self.pack_prefix}Q", self.kv_data_count))
         self.flush()
-#        print("tensors " + str(self.ti_data_count) + " kv " + str(self.kv_data_count))
+        #print("tensors " + str(self.ti_data_count) + " kv " + str(self.kv_data_count))
+        self.state = WriterState.HEADER
 
     def write_kv_data_to_file(self):
+        if self.state is not WriterState.HEADER:
+            raise ValueError(f'Expected output file to contain the header, got {self.state}')
+
         self.fout.write(self.kv_data)
         self.flush()
+        self.state = WriterState.KV_DATA
 
     def write_ti_data_to_file(self):
+        if self.state is not WriterState.KV_DATA:
+            raise ValueError(f'Expected output file to contain KV data, got {self.state}')
+
         self.fout.write(self.ti_data)
         self.flush()
+        self.state = WriterState.TI_DATA
 
     def add_key(self, key: str):
         self.add_val(key, GGUFValueType.STRING, add_vtype=False)
@@ -807,33 +825,24 @@ def add_tensor_info(self, name: str, tensor_shape: Sequence[int], tensor_dtype:
     def add_tensor(self, name: str, tensor: np.ndarray[Any, Any], raw_shape: Sequence[int] | None = None, raw_dtype: GGMLQuantizationType | None = None):
         if self.endianess == GGUFEndian.BIG:
             tensor.byteswap(inplace=True)
-        if self.use_temp_file and self.temp_file is None:
-            fp = tempfile.SpooledTemporaryFile(mode="w+b", max_size=256*1024*1024)
-            fp.seek(0)
-            self.temp_file = fp
 
         shape: Sequence[int] = raw_shape if raw_shape is not None else tensor.shape
         self.add_tensor_info(name, shape, tensor.dtype, tensor.nbytes, raw_dtype = raw_dtype)
 
-        pad = GGUFWriter.ggml_pad(tensor.nbytes, self.data_alignment) - tensor.nbytes
-
-        if  self.temp_file is None:
-            self.tensors.append((tensor, pad))
-            return
-
-        tensor.tofile(self.temp_file)
-
-        if pad != 0:
-            self.temp_file.write(bytes([0] * pad))
+        self.tensors.append(tensor)
 
     def write_padding(self, fp: BinaryIO, n: int, align: int | None = None):
         pad = GGUFWriter.ggml_pad(n, align if align is not None else self.data_alignment) - n
         if pad != 0:
             fp.write(bytes([0] * pad))
 
     def write_tensor_data(self, tensor: np.ndarray[Any, Any]):
+        if self.state is not WriterState.TI_DATA:
+            raise ValueError(f'Expected output file to contain tensor info, got {self.state}')
+
         if self.endianess==GGUFEndian.BIG:
             tensor.byteswap(inplace=True)
+
         self.write_padding(self.fout, self.fout.tell())
         tensor.tofile(self.fout)
         self.write_padding(self.fout, tensor.nbytes)
@@ -843,18 +852,13 @@ def write_tensors_to_file(self):
 
         self.write_padding(self.fout, self.fout.tell())
 
-        if self.temp_file is None:
-            for (currtensor, currpad) in self.tensors:
-                currtensor.tofile(self.fout)
-                if currpad != 0:
-                    self.fout.write(bytes([0] * currpad))
-            return
-
-        self.temp_file.seek(0)
-
-        shutil.copyfileobj(self.temp_file, self.fout)
-        self.flush()
-        self.temp_file.close()
+        while True:
+            try:
+                tensor = self.tensors.pop(0)
+            except IndexError:
+                break
+            tensor.tofile(self.fout)
+            self.write_padding(self.fout, tensor.nbytes)
 
     def flush(self):
         self.fout.flush()
@@ -983,11 +987,11 @@ def add_pad_token_id(self, id: int):
 
 
 class SpecialVocab:
-    load_merges: bool = False
-    merges: list[str] = []
-    special_token_types: tuple[str, ...] = ('bos', 'eos', 'unk', 'sep', 'pad')
-    special_token_ids: dict[str, int] = {}
-    n_vocab: int | None = None
+    load_merges: bool
+    merges: list[str]
+    special_token_types: tuple[str, ...]
+    special_token_ids: dict[str, int]
+    n_vocab: int | None
 
     def __init__(
         self, path: str | os.PathLike[str], load_merges: bool = False,
@@ -997,8 +1001,11 @@ def __init__(
         self.special_token_ids = {}
         self.n_vocab = n_vocab
         self.load_merges = load_merges
+        self.merges = []
         if special_token_types is not None:
             self.special_token_types = special_token_types
+        else:
+            self.special_token_types = ('bos', 'eos', 'unk', 'sep', 'pad')
         self._load(Path(path))
 
     def _load(self, path: Path) -> None: