File size: 5,074 Bytes
a449848 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 | """GGUF v3 writer and q4 tensor decoder reused from the v2 release."""
import argparse
import hashlib
import json
import os
import struct
import sys
import numpy as np
MAGIC = b"GGUF"
VERSION = 3
ALIGNMENT = 32
# GGUF metadata value types
T_UINT8, T_INT8, T_UINT16, T_INT16 = 0, 1, 2, 3
T_UINT32, T_INT32, T_FLOAT32, T_BOOL = 4, 5, 6, 7
T_STRING, T_ARRAY, T_UINT64 = 8, 9, 10
# GGML tensor types
GGML_F16, GGML_I8 = 1, 24
# ---------------------------------------------------------------------------- q4 decoding (mirrors bhc/xlit_q.py)
def read_q4(path):
blob = open(path, "rb").read()
n = int.from_bytes(blob[:4], "little")
header = blob[4:4 + n]
meta = json.loads(header)
return blob, header, meta, blob[4 + n:]
def dequantise(meta, data):
"""Dequantise every tensor of an xlit.q4 payload to float32 (same arithmetic as bhc/xlit_q.py)."""
bits_default, G = int(meta["bits"]), int(meta["G"])
out = []
for k, shape, kind, off, ln in meta["tensors"]:
cnt = int(np.prod(shape))
if kind.startswith("q"):
bits = int(kind[1:]) if len(kind) > 1 else bits_default
if bits == 4:
raw = np.frombuffer(data, np.uint8, ln, off)
q = np.stack([raw >> 4, raw & 15], 1).reshape(-1)[:cnt].astype(np.int16) - 8
else:
q = np.frombuffer(data, np.int8, ln, off).astype(np.int16)
r, c = shape
ng = (c + G - 1) // G
sc = np.frombuffer(data, np.float16, r * ng, off + ln).astype(np.float32).reshape(r, ng)
w = q.reshape(r, c).astype(np.float32) * np.repeat(sc, G, 1)[:, :c]
else:
w = np.frombuffer(data, np.float16, cnt, off).astype(np.float32).reshape(shape).copy()
out.append((k, list(shape), w))
return out
# ---------------------------------------------------------------------------- GGUF writer (spec-conformant, v3)
class Writer:
def __init__(self):
self.kv = []
# -- typed metadata helpers
def add(self, key, vtype, value):
self.kv.append((key, vtype, value))
return self
def string(self, k, v): return self.add(k, T_STRING, v)
def uint32(self, k, v): return self.add(k, T_UINT32, int(v))
def uint64(self, k, v): return self.add(k, T_UINT64, int(v))
def float32(self, k, v): return self.add(k, T_FLOAT32, float(v))
def boolean(self, k, v): return self.add(k, T_BOOL, bool(v))
def strings(self, k, v): return self.add(k, T_ARRAY, (T_STRING, list(v)))
@staticmethod
def _kv_bytes(key, vtype, value):
out = bytearray()
def s(x):
b = x.encode("utf-8")
return struct.pack("<Q", len(b)) + b
out += s(key)
out += struct.pack("<I", vtype)
if vtype == T_STRING:
out += s(value)
elif vtype == T_ARRAY:
sub, vals = value
out += struct.pack("<I", sub) + struct.pack("<Q", len(vals))
for v in vals:
if sub == T_STRING:
out += s(v)
else:
raise ValueError("unsupported array element type %d" % sub)
elif vtype == T_BOOL:
out += struct.pack("<?", value)
elif vtype == T_UINT32:
out += struct.pack("<I", value)
elif vtype == T_UINT64:
out += struct.pack("<Q", value)
elif vtype == T_FLOAT32:
out += struct.pack("<f", value)
else:
raise ValueError("unsupported metadata type %d" % vtype)
return bytes(out)
def write(self, path, tensors):
"""tensors: list of (name, np.ndarray, ggml_type_int)."""
# tensor table: dims are stored fastest-varying first
info = bytearray()
payload = bytearray()
for name, arr, gtype in tensors:
raw = arr.tobytes(order="C")
payload += raw
payload += b"\0" * ((-len(payload)) % ALIGNMENT)
offsets = []
off = 0
for name, arr, gtype in tensors:
offsets.append(off)
nbytes = arr.size * arr.dtype.itemsize
off += nbytes + ((-nbytes) % ALIGNMENT)
for (name, arr, gtype), o in zip(tensors, offsets):
b = name.encode("utf-8")
info += struct.pack("<Q", len(b)) + b
dims = list(arr.shape)
info += struct.pack("<I", len(dims))
for d in reversed(dims):
info += struct.pack("<Q", int(d))
info += struct.pack("<I", gtype)
info += struct.pack("<Q", o)
kv = bytearray()
for key, vtype, value in self.kv:
kv += self._kv_bytes(key, vtype, value)
head = MAGIC + struct.pack("<I", VERSION) + struct.pack("<Q", len(tensors)) + struct.pack("<Q", len(self.kv))
body = bytes(head) + bytes(kv) + bytes(info)
body += b"\0" * ((-len(body)) % ALIGNMENT)
with open(path, "wb") as f:
f.write(body)
f.write(bytes(payload))
return len(body) + len(payload), len(body)
|