#!/usr/bin/env python3 # Generate a tiny random-weight dspark GGUF for CPU tests of the loader's # log-SNR metadata contract (src/models/dspark.cpp load_arch_hparams): # # tests/test-dspark-logsnr-meta.cpp # # The weights are seeded random: the test only exercises model LOADING (the # hparam guard fires before any tensor data is touched), so weight values are # irrelevant. Unlike gen-tiny-qwen35.py this writes no tokenizer: the real # dspark converter calls _set_vocab_none() (the drafter ties to the TARGET # model's vocab), so tokenizer.ggml.model = "none" both matches production # GGUFs. Explicit vocab_size also makes token IDs valid in forward smoke tests. # # usage (single variant): # python3 tests/gen-tiny-dspark.py out.gguf --log-snr-conditioning 1 \ # --min-log-snr -9 --max-log-snr 6 # python3 tests/gen-tiny-dspark.py out.gguf --log-snr-conditioning 1 \ # --min-log-snr nan --max-log-snr 6 # python3 tests/gen-tiny-dspark.py out.gguf --log-snr-conditioning 1 --omit-min # # usage (full suite expected by test-dspark-logsnr-meta): # python3 tests/gen-tiny-dspark.py --suite from __future__ import annotations import argparse import sys from pathlib import Path import numpy as np sys.path.insert(0, str(Path(__file__).resolve().parent.parent / "gguf-py")) from gguf import GGUFWriter, GGUFValueType # noqa: E402 # smallest shape the dspark loader accepts: every tensor in # load_arch_tensors is present, with the optional heads (markov_rank = 0, # confidence_head off) and the tied-output fallback left to defaults. N_LAYER = 2 N_EMBD = 64 N_FF = 96 N_HEAD = 2 N_HEAD_KV = 1 N_CTX = 512 N_VOCAB = 256 # only exists as token_embd.weight's shape (vocab "none") HEAD_DIM = N_EMBD // N_HEAD # 32 TARGET_LAYERS = [1, 16] # n_capture = 2 (target-model indices) BLOCK_SIZE = 4 MASK_TOKEN_ID = 255 N_FREQ = 128 # LogSnrEmbed.NUM_FREQ_FEATURES (fixed in dspark.cpp) def write_variant(out_path: str, conditioning: bool, min_log_snr: float | None, max_log_snr: float | None, correction: bool = False) -> None: rng = np.random.default_rng(1234) def rand(*shape: int) -> np.ndarray: return rng.normal(0.0, 0.02, size=shape).astype(np.float32) w = GGUFWriter(out_path, "dspark") w.add_name("tiny-random-dspark") w.add_file_type(0) # all f32 w.add_block_count(N_LAYER) w.add_context_length(N_CTX) w.add_embedding_length(N_EMBD) w.add_feed_forward_length(N_FF) w.add_head_count(N_HEAD) w.add_head_count_kv(N_HEAD_KV) w.add_layer_norm_rms_eps(1e-6) w.add_rope_freq_base(10000.0) w.add_uint32("dspark.dspark.block_size", BLOCK_SIZE) w.add_uint32("dspark.dspark.mask_token_id", MASK_TOKEN_ID) w.add_key_value("dspark.dspark.target_layers", TARGET_LAYERS, GGUFValueType.ARRAY, sub_type=GGUFValueType.UINT32) w.add_bool("dspark.dspark.log_snr_conditioning", conditioning) if min_log_snr is not None: w.add_float32("dspark.dspark.min_log_snr", min_log_snr) if max_log_snr is not None: w.add_float32("dspark.dspark.max_log_snr", max_log_snr) w.add_tokenizer_model("none") w.add_vocab_size(N_VOCAB) if correction: w.add_bool("dspark.dspark.hidden_correction", True) w.add_bool("dspark.dspark.markov_deploy_exposure", True) w.add_uint32("dspark.dspark.correction_size", N_EMBD) w.add_uint32("dspark.dspark.markov_rank", 8) w.add_tensor("dspark.correction_hidden_norm.weight", np.ones(N_EMBD, dtype=np.float32)) w.add_tensor("dspark.correction_embed_norm.weight", np.ones(N_EMBD, dtype=np.float32)) w.add_tensor("dspark.correction_gate.weight", rand(N_EMBD, 2 * N_EMBD)) w.add_tensor("dspark.correction_up.weight", rand(N_EMBD, 2 * N_EMBD)) w.add_tensor("dspark.correction_down.weight", rand(N_EMBD, N_EMBD)) w.add_tensor("dspark.markov_head_a.weight", rand(N_VOCAB, 8)) w.add_tensor("dspark.markov_head_b.weight", rand(N_VOCAB, 8)) # note: numpy shapes are the reverse of the ggml ne[] order w.add_tensor("token_embd.weight", rand(N_VOCAB, N_EMBD)) w.add_tensor("output_norm.weight", np.ones(N_EMBD, dtype=np.float32)) # output.weight omitted -> duplicated from token_embd n_embd_cap = len(TARGET_LAYERS) * N_EMBD w.add_tensor("dspark.fc.weight", rand(N_EMBD, n_embd_cap)) w.add_tensor("dspark.hidden_norm.weight", np.ones(N_EMBD, dtype=np.float32)) # the log-SNR MLP is REQUIRED by the loader when conditioning is on, and # must be absent otherwise (unreferenced tensors fail the load) if conditioning: w.add_tensor("dspark.log_snr_fc1.weight", rand(N_EMBD, N_FREQ)) w.add_tensor("dspark.log_snr_fc1.bias", rand(N_EMBD)) w.add_tensor("dspark.log_snr_fc2.weight", rand(N_EMBD, N_EMBD)) w.add_tensor("dspark.log_snr_fc2.bias", rand(N_EMBD)) for il in range(N_LAYER): w.add_tensor(f"blk.{il}.attn_norm.weight", np.ones(N_EMBD, dtype=np.float32)) w.add_tensor(f"blk.{il}.attn_q.weight", rand(HEAD_DIM * N_HEAD, N_EMBD)) w.add_tensor(f"blk.{il}.attn_k.weight", rand(HEAD_DIM * N_HEAD_KV, N_EMBD)) w.add_tensor(f"blk.{il}.attn_v.weight", rand(HEAD_DIM * N_HEAD_KV, N_EMBD)) w.add_tensor(f"blk.{il}.attn_output.weight", rand(N_EMBD, HEAD_DIM * N_HEAD)) w.add_tensor(f"blk.{il}.attn_q_norm.weight", np.ones(HEAD_DIM, dtype=np.float32)) w.add_tensor(f"blk.{il}.attn_k_norm.weight", np.ones(HEAD_DIM, dtype=np.float32)) w.add_tensor(f"blk.{il}.ffn_norm.weight", np.ones(N_EMBD, dtype=np.float32)) w.add_tensor(f"blk.{il}.ffn_gate.weight", rand(N_FF, N_EMBD)) w.add_tensor(f"blk.{il}.ffn_down.weight", rand(N_EMBD, N_FF)) w.add_tensor(f"blk.{il}.ffn_up.weight", rand(N_FF, N_EMBD)) w.write_header_to_file() w.write_kv_data_to_file() w.write_tensors_to_file() w.close() print(f"wrote {out_path} (conditioning={int(conditioning)} " f"min={min_log_snr} max={max_log_snr})") # canonical suite consumed by tests/test-dspark-logsnr-meta.cpp: keep the # file names and settings in sync with EXPECTED_VARIANTS there SUITE = [ # name cond min max ("logsnr-valid.gguf", True, -9.0, 6.0), ("logsnr-omit-min.gguf", True, None, 6.0), ("logsnr-omit-max.gguf", True, -9.0, None), ("logsnr-min-nan.gguf", True, float("nan"), 6.0), ("logsnr-max-inf.gguf", True, -9.0, float("inf")), ("logsnr-equal.gguf", True, 6.0, 6.0), ("logsnr-reversed.gguf", True, 6.0, -9.0), ("logsnr-cond-off.gguf", False, None, None), ] def main() -> None: ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) ap.add_argument("out", nargs="?", help="output GGUF path (single-variant mode)") ap.add_argument("--log-snr-conditioning", type=int, choices=(0, 1), default=0) ap.add_argument("--min-log-snr", type=float, default=None, help="min_log_snr metadata value (accepts nan/inf/-inf)") ap.add_argument("--max-log-snr", type=float, default=None, help="max_log_snr metadata value (accepts nan/inf/-inf)") ap.add_argument("--omit-min", action="store_true", help="omit min_log_snr") ap.add_argument("--omit-max", action="store_true", help="omit max_log_snr") ap.add_argument("--hidden-correction", action="store_true") ap.add_argument("--suite", metavar="DIR", help="write the full variant suite for test-dspark-logsnr-meta") args = ap.parse_args() if args.suite: out_dir = Path(args.suite) out_dir.mkdir(parents=True, exist_ok=True) for name, cond, lo, hi in SUITE: write_variant(str(out_dir / name), cond, lo, hi) return if not args.out: ap.error("either an output path or --suite is required") min_v = None if args.omit_min else args.min_log_snr max_v = None if args.omit_max else args.max_log_snr write_variant(args.out, bool(args.log_snr_conditioning), min_v, max_v, args.hidden_correction) if __name__ == "__main__": main()