method/gguf_extract.py

#!/usr/bin/env python
"""Extract selected tensors from a (split) GGUF into a single, loadable GGUF. Pure Python, no numpy.

The output carries the full metadata of the first shard (so llama-quantize can load it), with the
split.* keys rewritten to describe a single file, followed by only the selected tensors.

Usage:
  python gguf_extract.py <first shard or single .gguf> <output.gguf> <regex on tensor name>
Example (the 2026-09-03 build, all expert down projections):
  python gguf_extract.py Qwen3.8-Flash-Next-UD-Q4_K_XL-00001-of-00004.gguf downs.gguf "ffn_down_exps"

Written 2026-09-03 for the "Q4_K_XL with IQ4_NL downs" build (see studies/). See gguf_splice.py for the way back.
"""
import glob, os, re, struct, sys

ALIGN = 32
T = {0: 'B', 1: 'b', 2: 'H', 3: 'h', 4: 'I', 5: 'i', 6: 'f', 7: '?', 10: 'Q', 11: 'q', 12: 'd'}
# bytes per block, elements per block (for exact tensor byte sizes)
BLOCK = {0: (4, 1), 1: (2, 1), 30: (2, 1), 2: (18, 32), 3: (20, 32), 6: (22, 32), 7: (24, 32), 8: (34, 32),
         10: (84, 256), 11: (110, 256), 12: (144, 256), 13: (176, 256), 14: (210, 256), 20: (18, 32),
         23: (136, 256), 21: (110, 256), 22: (82, 256), 18: (66, 256), 16: (66, 256), 17: (74, 256),
         19: (50, 256), 29: (56, 256)}


def nbytes(shape, ty):
    n = 1
    for d in shape: n *= d
    bs, be = BLOCK[ty]
    assert n % be == 0, (shape, ty)
    return n // be * bs


class Reader:
    def __init__(self, path):
        self.path = path; self.f = open(path, 'rb'); self.size = os.path.getsize(path)
        assert self.f.read(4) == b'GGUF'
        self.version = self.rd('I'); self.n_tensors = self.rd('Q'); self.n_kv = self.rd('Q')
        self.kv = []          # (key, type, raw_bytes_of_value)
        for _ in range(self.n_kv):
            key = self.rstr(); ty = self.rd('I'); start = self.f.tell(); self.skip_val(ty); end = self.f.tell()
            self.f.seek(start); raw = self.f.read(end - start); self.kv.append((key, ty, raw))
        self.tensors = []     # (name, shape, type, offset)
        for _ in range(self.n_tensors):
            name = self.rstr(); nd = self.rd('I'); shape = [self.rd('Q') for _ in range(nd)]
            ty = self.rd('I'); off = self.rd('Q'); self.tensors.append((name, shape, ty, off))
        self.data_start = (self.f.tell() + ALIGN - 1) // ALIGN * ALIGN

    def rd(self, fmt): return struct.unpack('<' + fmt, self.f.read(struct.calcsize(fmt)))[0]
    def rstr(self): n = self.rd('Q'); return self.f.read(n).decode('utf-8', 'replace')
    def skip_val(self, ty):
        if ty == 8: n = self.rd('Q'); self.f.seek(n, 1)
        elif ty == 9:
            et = self.rd('I'); n = self.rd('Q')
            for _ in range(n): self.skip_val(et)
        else: self.f.seek(struct.calcsize(T[ty]), 1)

    def read_tensor(self, name, shape, ty, off):
        self.f.seek(self.data_start + off); return self.f.read(nbytes(shape, ty))


def wstr(s):
    b = s.encode('utf-8'); return struct.pack('<Q', len(b)) + b


def patched_kv(kv, n_tensors_out):
    out = []
    for key, ty, raw in kv:
        if key == 'split.no': raw = struct.pack('<H', 0)
        elif key == 'split.count': raw = struct.pack('<H', 1)
        elif key == 'split.tensors.count': raw = struct.pack('<i', n_tensors_out)
        out.append((key, ty, raw))
    return out


def shards_of(first):
    m = re.match(r'(.*)-(\d{5})-of-(\d{5})\.gguf$', first)
    if not m: return [first]
    base, _, count = m.groups()
    return [f'{base}-{i:05d}-of-{count}.gguf' for i in range(1, int(count) + 1)]


def main():
    first, out, pattern = sys.argv[1], sys.argv[2], re.compile(sys.argv[3])
    shards = [Reader(p) for p in shards_of(first)]
    meta = shards[0].kv
    selected = [(r, t) for r in shards for t in r.tensors if pattern.search(t[0])]
    print(f'{len(selected)} tensors selected from {len(shards)} shard(s)')
    kv = patched_kv(meta, len(selected))
    # header
    hdr = b'GGUF' + struct.pack('<IQQ', 3, len(selected), len(kv))
    for key, ty, raw in kv: hdr += wstr(key) + struct.pack('<I', ty) + raw
    infos = b''; off = 0; layout = []
    for r, (name, shape, ty, src_off) in selected:
        sz = nbytes(shape, ty); layout.append((r, name, shape, ty, src_off, off, sz))
        infos += wstr(name) + struct.pack('<I', len(shape)) + b''.join(struct.pack('<Q', d) for d in shape) + struct.pack('<IQ', ty, off)
        off += (sz + ALIGN - 1) // ALIGN * ALIGN
    head = hdr + infos
    pad = (ALIGN - len(head) % ALIGN) % ALIGN
    total = 0
    with open(out, 'wb') as o:
        o.write(head + b'\x00' * pad)
        for r, name, shape, ty, src_off, dst_off, sz in layout:
            r.f.seek(r.data_start + src_off)
            remaining = sz
            while remaining:
                chunk = r.f.read(min(remaining, 64 << 20)); o.write(chunk); remaining -= len(chunk)
            o.write(b'\x00' * ((ALIGN - sz % ALIGN) % ALIGN)); total += sz
            print(f'  {name:40s} type {ty:2d} {sz / 1e6:9.1f} MB', file=sys.stderr)
    print(f'wrote {out}: {total / 1e9:.2f} GB of tensor data')


if __name__ == '__main__':
    main()