Fabio
c36510c21e
Some checks failed
Linux Build / Linux cpu (push) Has been cancelled
Linux Build / Linux cpu-clang (push) Has been cancelled
Linux Build / Linux vulkan (push) Has been cancelled
mac-build / macOS CPU (push) Has been cancelled
Nix Build / Nix (cpu) (push) Has been cancelled
Nix Build / Nix (vulkan) (push) Has been cancelled
windows-build / Windows CPU (push) Has been cancelled
103 lines
4.7 KiB
Python
103 lines
4.7 KiB
Python
#!/usr/bin/env python3
|
|
"""Build AuK component GGUFs from upstream safetensors."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import json
|
|
import os
|
|
import shutil
|
|
import subprocess
|
|
import tempfile
|
|
from pathlib import Path
|
|
|
|
|
|
REPO = Path(__file__).resolve().parents[2]
|
|
CONVERTER = REPO / "build/debug/bin/audiocpp_gguf"
|
|
|
|
|
|
def main() -> None:
|
|
variants = (
|
|
("qwen-bf16", "qwen2.5-omni-3b-bf16.gguf", "bf16", 923),
|
|
("qwen-q8", "qwen2.5-omni-3b-q8_0.gguf", "q8_0", 923),
|
|
("base-f32", "auk-base-f32.gguf", "f32", 420),
|
|
("base-f16", "auk-base-f16.gguf", "f16", 420),
|
|
("base-q8", "auk-base-q8_0.gguf", "q8_0", 420),
|
|
("flash-f32", "auk-flash-f32.gguf", "f32", 420),
|
|
("flash-f16", "auk-flash-f16.gguf", "f16", 420),
|
|
("flash-q8", "auk-flash-q8_0.gguf", "q8_0", 420),
|
|
("vae-f32", "auk-vae-f32.gguf", "f32", 1137),
|
|
)
|
|
parser = argparse.ArgumentParser(description=__doc__)
|
|
parser.add_argument("--base-dir", type=Path, required=True, help="AuK Base checkpoint directory")
|
|
parser.add_argument("--flash-dir", type=Path, required=True, help="AuK-Flash checkpoint directory")
|
|
parser.add_argument("--qwen-dir", type=Path, required=True, help="Qwen2.5-Omni-3B checkpoint directory")
|
|
parser.add_argument("--output-dir", type=Path, required=True, help="destination for component GGUFs and sidecars")
|
|
parser.add_argument("--only", choices=[item[0] for item in variants], action="append")
|
|
args = parser.parse_args()
|
|
selected = set(args.only or [item[0] for item in variants])
|
|
base_dir = args.base_dir.resolve()
|
|
flash_dir = args.flash_dir.resolve()
|
|
qwen_dir = args.qwen_dir.resolve()
|
|
output_dir = args.output_dir.resolve()
|
|
subprocess.run(
|
|
["cmake", "--build", str(REPO / "build/debug"), "--target", "audiocpp_gguf", "-j", "8"],
|
|
check=True,
|
|
)
|
|
output_dir.mkdir(parents=True, exist_ok=True)
|
|
(output_dir / "config").mkdir(exist_ok=True)
|
|
(output_dir / "tokenizer").mkdir(exist_ok=True)
|
|
shutil.copyfile(base_dir / "config.yaml", output_dir / "config/auk-base.yaml")
|
|
shutil.copyfile(flash_dir / "config.yaml", output_dir / "config/auk-flash.yaml")
|
|
for name in ("tokenizer.json", "tokenizer_config.json"):
|
|
shutil.copyfile(qwen_dir / name, output_dir / "tokenizer" / name)
|
|
|
|
with tempfile.TemporaryDirectory(prefix="auk-qwen-index-") as workspace:
|
|
shards = Path(workspace)
|
|
index = json.loads((qwen_dir / "model.safetensors.index.json").read_text())
|
|
weight_map = {
|
|
name: shard for name, shard in index["weight_map"].items()
|
|
if name.startswith(("thinker.model.", "thinker.audio_tower."))
|
|
}
|
|
if len(weight_map) != 923:
|
|
raise RuntimeError(f"expected 923 AuK Qwen tensors, found {len(weight_map)}")
|
|
(shards / "model.safetensors.index.json").write_text(json.dumps({"weight_map": weight_map}, sort_keys=True))
|
|
for name in set(weight_map.values()):
|
|
(shards / name).symlink_to(qwen_dir / name)
|
|
|
|
sources = {
|
|
"qwen": shards / "model.safetensors.index.json",
|
|
"base": base_dir / "auk_base.safetensors",
|
|
"flash": flash_dir / "auk_flash.safetensors",
|
|
"vae": base_dir / "vae.safetensors",
|
|
}
|
|
for variant, filename, weight_type, tensor_count in variants:
|
|
if variant not in selected:
|
|
continue
|
|
source = sources[variant.split("-")[0]]
|
|
fd, temporary = tempfile.mkstemp(prefix=variant + "-", suffix=".gguf", dir=output_dir)
|
|
os.close(fd)
|
|
temporary = Path(temporary)
|
|
try:
|
|
subprocess.run([
|
|
str(CONVERTER), "--input", str(source), "--root", str(output_dir),
|
|
"--output", str(temporary), "--type", weight_type,
|
|
"--family", "auk", "--model-spec", str(REPO / "model_specs/auk.json"),
|
|
"--no-sidecars", "--allow-missing-model-spec", "--overwrite",
|
|
], check=True)
|
|
inspection = subprocess.run(
|
|
[str(CONVERTER), "--inspect", str(temporary)],
|
|
check=True, capture_output=True, text=True,
|
|
).stdout
|
|
if f"tensors={tensor_count}" not in inspection.splitlines() or \
|
|
"embedded_sidecar_count=0" not in inspection.splitlines():
|
|
raise RuntimeError(f"AuK component inspection failed:\n{inspection}")
|
|
output = output_dir / filename
|
|
os.replace(temporary, output)
|
|
print(f"Rebuilt {output}\n{inspection}", flush=True)
|
|
finally:
|
|
temporary.unlink(missing_ok=True)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|