GGUF : Support writing tensors in Python
This commit is contained in:
parent
464192b9be
commit
bb54d1700e
1 changed files with 7 additions and 1 deletions
8
gguf.py
8
gguf.py
|
@ -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()
|
||||||
|
|
||||||
|
|
Loading…
Add table
Add a link
Reference in a new issue