GGUF : Support writing tensors in Python

This commit is contained in:
M. Yusuf Sarıgöz 2023-07-27 23:09:53 +03:00
parent 464192b9be
commit bb54d1700e

View File

@ -62,7 +62,7 @@ class GGUFWriter:
def __init__(self, fout: IO): def __init__(self, fout: IO):
self.fout = fout self.fout = fout
self.offset_tensor = 0 self.offset_tensor = 0
self.tensors = [] self.tensors: List[np.ndarray] = []
def write_header(self, tensor_count: int, metadata_kv_count: int): def write_header(self, tensor_count: int, metadata_kv_count: int):
self.fout.write(struct.pack("<I", constants.GGUF_MAGIC)) self.fout.write(struct.pack("<I", constants.GGUF_MAGIC))
@ -176,6 +176,12 @@ class GGUFWriter:
self.tensors.append(tensor) self.tensors.append(tensor)
def write_tensors(self):
for tensor in self.tensors:
tensor.tofile(self.fout)
pad = GGUFWriter.ggml_pad(tensor.nbytes, constants.GGUF_DEFAULT_ALIGNMENT) - tensor.nbytes
self.fout.write(bytes([0] * pad))
def flush(self): def flush(self):
self.fout.flush() self.fout.flush()