File size: 5,582 Bytes
ddf8c5b | 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 143 144 | #!/usr/bin/env python3
"""Add tokenizer.ggml.mask_token_id to the DSpark draft GGUF.
llama.cpp's DSpark spec reads the draft's mask token via
llama_vocab_mask(vocab) -> tokenizer.ggml.mask_token_id. Our rewritten draft
copied K3's tokenizer, which has no mask token, so it returns -1 and the draft
batch gets fed token -1 -> 'invalid token[1] = -1' -> draft decode fails every step.
The draft's *trained* mask id lives in dflash.mask_token_id (163824). We copy it
into tokenizer.ggml.mask_token_id so llama_vocab_mask() returns the right id.
"""
import struct, sys
SRC = "/root/models/k3-draft/draft_final.gguf"
DST = "/root/models/k3-draft/draft_masked.gguf"
MASK_KEY = b"tokenizer.ggml.mask_token_id"
DFLASH_MASK_KEY = b"dflash.mask_token_id"
def read_kv(data, pos):
klen = struct.unpack_from("<Q", data, pos)[0]; pos += 8
key = data[pos:pos+klen]; pos += klen
vtype = struct.unpack_from("<I", data, pos)[0]; pos += 4
if vtype == 8:
vlen = struct.unpack_from("<Q", data, pos)[0]; pos += 8
val = data[pos:pos+vlen]; pos += vlen
elif vtype in (0, 1, 7):
val = data[pos:pos+1]; pos += 1
elif vtype in (2, 3):
val = data[pos:pos+2]; pos += 2
elif vtype in (4, 5, 6):
val = data[pos:pos+4]; pos += 4
elif vtype in (10, 11, 12):
val = data[pos:pos+8]; pos += 8
elif vtype == 9:
atype = struct.unpack_from("<I", data, pos)[0]; pos += 4
alen = struct.unpack_from("<Q", data, pos)[0]; pos += 8
if atype == 8:
vals = []
for _ in range(alen):
slen = struct.unpack_from("<Q", data, pos)[0]; pos += 8
vals.append(data[pos:pos+slen]); pos += slen
val = (atype, vals)
elif atype in (0, 1, 7):
val = (atype, data[pos:pos+alen]); pos += alen
elif atype in (2, 3):
val = (atype, data[pos:pos+alen*2]); pos += alen*2
elif atype in (4, 5, 6):
val = (atype, data[pos:pos+alen*4]); pos += alen*4
elif atype in (10, 11, 12):
val = (atype, data[pos:pos+alen*8]); pos += alen*8
else:
raise ValueError(f"array atype {atype}")
else:
raise ValueError(f"vtype {vtype} key {key}")
return key, vtype, val, pos
def write_kv(out, key, vtype, val):
out += struct.pack("<Q", len(key)) + key + struct.pack("<I", vtype)
if vtype == 8:
out += struct.pack("<Q", len(val)) + val
elif vtype in (0, 1, 7, 2, 3, 4, 5, 6, 10, 11, 12):
out += val
elif vtype == 9:
atype, d = val
out += struct.pack("<I", atype)
if atype == 8:
out += struct.pack("<Q", len(d))
for s in d:
out += struct.pack("<Q", len(s)) + s
else:
sz = {0:1,1:1,7:1,2:2,3:2,4:4,5:4,6:4,10:8,11:8,12:8}[atype]
out += struct.pack("<Q", len(d)//sz) + d
def main():
data = open(SRC, "rb").read()
pos = 0
magic = data[0:4]; pos = 4
version = struct.unpack_from("<I", data, pos)[0]; pos += 4
n_tensors = struct.unpack_from("<Q", data, pos)[0]; pos += 8
n_kv = struct.unpack_from("<Q", data, pos)[0]; pos += 8
kvs = []
mask_val = None
have_tok_mask = False
for _ in range(n_kv):
key, vtype, val, pos = read_kv(data, pos)
if key == DFLASH_MASK_KEY:
mask_val = struct.unpack("<I", val)[0] if vtype in (4,5) else None
print(f"found {DFLASH_MASK_KEY.decode()} vtype={vtype} val={mask_val}")
if key == MASK_KEY:
have_tok_mask = True
print(f"WARNING: {MASK_KEY.decode()} already present")
kvs.append((key, vtype, val))
if have_tok_mask:
print("mask token already present, nothing to do"); return
if mask_val is None:
print("ERROR: dflash.mask_token_id not found"); sys.exit(1)
# Append the new key
kvs.append((MASK_KEY, 4, struct.pack("<I", mask_val))) # uint32
print(f"adding {MASK_KEY.decode()} = {mask_val} (uint32)")
# Rebuild header + KV
out = bytearray()
out += magic + struct.pack("<I", version) + struct.pack("<Q", n_tensors) + struct.pack("<Q", len(kvs))
for key, vtype, val in kvs:
write_kv(out, key, vtype, val)
kv_end_new = len(out)
# tensor infos + data follow unchanged from original pos
rest = data[pos:]
# The tensor data section is aligned; tensor info offsets are relative to data start.
# We must recompute alignment: original data start vs new data start.
alignment = 32
# original tensor-data start
orig_data_start = (pos + 0) # pos == end of KV in original
# tensor infos occupy from pos; find their end by walking n_tensors
ti_pos = pos
for _ in range(n_tensors):
nl = struct.unpack_from("<Q", data, ti_pos)[0]; ti_pos += 8 + nl
nd = struct.unpack_from("<I", data, ti_pos)[0]; ti_pos += 4 + nd*8
ti_pos += 4 # dtype
ti_pos += 8 # offset
orig_tensor_data_start = (ti_pos + alignment - 1)//alignment*alignment
# new tensor-data start
new_kv_and_ti_len = len(out) + (ti_pos - pos)
new_tensor_data_start = (new_kv_and_ti_len + alignment - 1)//alignment*alignment
# append tensor infos verbatim
out += data[pos:ti_pos]
while len(out) < new_tensor_data_start:
out += b"\x00"
out += data[orig_tensor_data_start:]
open(DST, "wb").write(out)
print(f"wrote {DST}: {len(out)} bytes (src {len(data)})")
print("verify: tokenizer.ggml.mask_token_id should now read", mask_val)
if __name__ == "__main__":
main()
|