feat(v4): pipeline tools — corpus prep, GGUF embed/head extractor, chunked dump driver

tools/extract_embed_head.py

@@ -0,0 +1,61 @@
+#!/usr/bin/env python3
+"""Extract token_embd (F16) + output head (Q8_0 -> dequant f16) from a V4 GGUF
+as the frozen draft init (DeepSpec copies target embed/lm_head into the draft)."""
+import argparse, struct
+import numpy as np
+from safetensors.numpy import save_file
+
+ap = argparse.ArgumentParser()
+ap.add_argument("--gguf", required=True)
+ap.add_argument("--out", required=True)
+args = ap.parse_args()
+
+f = open(args.gguf, "rb")
+f.read(8)
+n_t = struct.unpack("<Q", f.read(8))[0]
+n_kv = struct.unpack("<Q", f.read(8))[0]
+def rd_str():
+    n = struct.unpack("<Q", f.read(8))[0]; return f.read(n)
+sizes = {0:1,1:1,2:2,3:2,4:4,5:4,6:4,7:1,10:8,11:8,12:8}
+align = 32
+for _ in range(n_kv):
+    k = rd_str().decode(); t = struct.unpack("<I", f.read(4))[0]
+    if t == 8: rd_str()
+    elif t == 9:
+        et = struct.unpack("<I", f.read(4))[0]; n = struct.unpack("<Q", f.read(8))[0]
+        if et == 8:
+            for _ in range(n): rd_str()
+        else: f.seek(sizes[et]*n, 1)
+    else:
+        v = f.read(sizes[t])
+        if k == "general.alignment": align = struct.unpack("<I", v)[0]
+infos = {}
+for _ in range(n_t):
+    name = rd_str().decode()
+    nd = struct.unpack("<I", f.read(4))[0]
+    dims = struct.unpack(f"<{nd}Q", f.read(8*nd))
+    ttype = struct.unpack("<I", f.read(4))[0]
+    off = struct.unpack("<Q", f.read(8))[0]
+    if name in ("token_embd.weight", "output.weight"):
+        infos[name] = (dims, ttype, off)
+data_start = (f.tell() + align - 1) // align * align
+out = {}
+for name, (dims, ttype, off) in infos.items():
+    ne = 1
+    for d in dims: ne *= d
+    f.seek(data_start + off)
+    if ttype == 1:  # F16
+        w = np.frombuffer(f.read(ne*2), dtype=np.float16).reshape(dims[::-1])
+    elif ttype == 8:  # Q8_0: 32-blocks of f16 scale + 32 i8
+        nb = ne // 32
+        blk = np.frombuffer(f.read(nb*34), dtype=np.uint8).reshape(nb, 34)
+        d = blk[:, :2].copy().view(np.float16).astype(np.float32)
+        q = blk[:, 2:].copy().view(np.int8).astype(np.float32)
+        w = (q * d).astype(np.float16).reshape(dims[::-1])
+    else:
+        raise SystemExit(f"unexpected ggml type {ttype} for {name}")
+    out[name] = w
+    print(name, w.shape, w.dtype)
+save_file({"embed_tokens.weight": out["token_embd.weight"],
+           "lm_head.weight": out["output.weight"]}, args.out)
+print("saved ->", args.out)

tools/prepare_corpus.py

@@ -0,0 +1,37 @@
+#!/usr/bin/env python3
+"""open-perfectblend parquet -> JSONL for hanzo-engine v4_cache_dump.
+
+Each line: {"text": human + "\n\n" + gpt, "prompt_frac": len(human)/len(text), "source": ...}
+prompt_frac lets the converter zero loss_mask over the prompt region by ratio.
+"""
+import argparse, glob, json
+import pyarrow.parquet as pq
+
+ap = argparse.ArgumentParser()
+ap.add_argument("--dataset-dir", required=True)
+ap.add_argument("--out", required=True)
+ap.add_argument("--min-chars", type=int, default=200)
+ap.add_argument("--max-chars", type=int, default=8000)
+args = ap.parse_args()
+
+fs = sorted(glob.glob(f"{args.dataset_dir}/**/*.parquet", recursive=True))
+out = open(args.out, "w")
+n = kept = 0
+for fp in fs:
+    for batch in pq.ParquetFile(fp).iter_batches(batch_size=2000):
+        for row in batch.to_pylist():
+            n += 1
+            conv = row["conversations"]
+            h = next((t["value"] for t in conv if t["from"] == "human"), None)
+            g = next((t["value"] for t in conv if t["from"] == "gpt"), None)
+            if not h or not g:
+                continue
+            text = h.strip() + "\n\n" + g.strip()
+            if not (args.min_chars <= len(text) <= args.max_chars):
+                continue
+            kept += 1
+            out.write(json.dumps({"text": text, "prompt_frac": round(len(h) / len(text), 4),
+                                  "source": row.get("source", "")}) + "\n")
+    print(f"{fp.split('/')[-1]}: total={n} kept={kept}", flush=True)
+out.close()
+print(f"DONE: {kept}/{n} -> {args.out}")

tools/run_dump_chunks.sh

@@ -0,0 +1,23 @@
+#!/bin/bash
+# Chunked V4 training-cache dump: 5000-sample chunks, resumable, nohup-friendly.
+# Usage: run_dump_chunks.sh <corpus.jsonl> <out_root> <start_chunk> <n_chunks> <engine_dir>
+set -u
+CORPUS=$1; ROOT=$2; START=${3:-0}; N=${4:-4}; ENG=${5:-/home/z/work/hanzo/engine-v4bench}
+CH=5000
+for ((i=START; i<START+N; i++)); do
+  DIR=$ROOT/chunk_$(printf %03d $i)
+  [ -f "$DIR/.done" ] && echo "chunk $i done, skip" && continue
+  mkdir -p "$DIR"
+  FROM=$((i*CH+1)); TO=$(((i+1)*CH))
+  sed -n "${FROM},${TO}p" "$CORPUS" > "$DIR/input.jsonl"
+  echo "════ chunk $i: lines $FROM-$TO ════"
+  (cd "$ENG" && IGPU_MEMORY_FRACTION=0.92 RUST_LOG=warn CUDA_HOME=/usr/local/cuda \
+    LD_LIBRARY_PATH=/usr/local/cuda/lib64 \
+    ./target/release/examples/v4_cache_dump "$DIR/input.jsonl" "$DIR" $CH 2048) \
+    >> "$ROOT/dump.log" 2>&1
+  # exit 139 teardown-after-write is benign; judge by the index
+  LINES=$(wc -l < "$DIR/index.jsonl" 2>/dev/null || echo 0)
+  echo "chunk $i: $LINES samples indexed"
+  [ "$LINES" -gt $((CH*9/10)) ] && touch "$DIR/.done" || { echo "chunk $i INCOMPLETE, stopping"; exit 1; }
+done
+echo ALL_CHUNKS_DONE