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
452 lines
17 KiB
Python
452 lines
17 KiB
Python
#!/usr/bin/env python3
|
|
"""Convert a free Kroko Zipformer2 package into audio.cpp tensor assets.
|
|
|
|
Kroko community ``.data`` files contain a small JSON header followed by four
|
|
little-endian length-prefixed blobs: encoder ONNX, decoder ONNX, joiner ONNX,
|
|
and tokens.txt. The ONNX graphs published by Kroko are dynamically quantized.
|
|
This converter recovers named floating-point weights from those graphs so the
|
|
native audio.cpp runtime and the normal GGUF converter can consume them.
|
|
|
|
Commercial/encrypted Kroko packages are intentionally unsupported.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import json
|
|
import struct
|
|
from collections import defaultdict
|
|
from pathlib import Path
|
|
from tempfile import TemporaryDirectory
|
|
from typing import Iterable
|
|
|
|
import numpy as np
|
|
import onnx
|
|
from onnx import numpy_helper
|
|
from safetensors.numpy import save_file
|
|
|
|
|
|
LAYER_PREFIXES = [
|
|
*(f"encoder.encoders.0.layers.{i}" for i in range(2)),
|
|
*(f"encoder.encoders.1.encoder.layers.{i}" for i in range(2)),
|
|
*(f"encoder.encoders.2.encoder.layers.{i}" for i in range(4)),
|
|
*(f"encoder.encoders.3.encoder.layers.{i}" for i in range(5)),
|
|
*(f"encoder.encoders.4.encoder.layers.{i}" for i in range(4)),
|
|
*(f"encoder.encoders.5.encoder.layers.{i}" for i in range(2)),
|
|
]
|
|
|
|
|
|
def read_u32(data: bytes, offset: int) -> tuple[int, int]:
|
|
if offset + 4 > len(data):
|
|
raise ValueError("truncated Kroko package")
|
|
return struct.unpack_from("<I", data, offset)[0], offset + 4
|
|
|
|
|
|
def unpack_package(source: Path, destination: Path) -> dict:
|
|
data = source.read_bytes()
|
|
header_size, offset = read_u32(data, 0)
|
|
header_end = offset + header_size
|
|
if header_end > len(data):
|
|
raise ValueError("truncated Kroko JSON header")
|
|
header = json.loads(data[offset:header_end].decode("utf-8"))
|
|
if not header.get("free", False):
|
|
raise ValueError(
|
|
"commercial/encrypted Kroko packages are not supported; "
|
|
"use a community package with free=true"
|
|
)
|
|
offset = header_end
|
|
names = ("encoder.onnx", "decoder.onnx", "joiner.onnx", "tokens.txt")
|
|
destination.mkdir(parents=True, exist_ok=True)
|
|
for name in names:
|
|
length, offset = read_u32(data, offset)
|
|
end = offset + length
|
|
if end > len(data):
|
|
raise ValueError(f"truncated Kroko package block: {name}")
|
|
(destination / name).write_bytes(data[offset:end])
|
|
offset = end
|
|
if offset != len(data):
|
|
raise ValueError(f"unexpected {len(data) - offset} trailing package bytes")
|
|
return header
|
|
|
|
|
|
def tensor_map(model: onnx.ModelProto) -> dict[str, np.ndarray]:
|
|
return {
|
|
tensor.name: np.asarray(numpy_helper.to_array(tensor))
|
|
for tensor in model.graph.initializer
|
|
}
|
|
|
|
|
|
def consumers(model: onnx.ModelProto) -> dict[str, list[onnx.NodeProto]]:
|
|
result: dict[str, list[onnx.NodeProto]] = defaultdict(list)
|
|
for node in model.graph.node:
|
|
for name in node.input:
|
|
result[name].append(node)
|
|
return result
|
|
|
|
|
|
def downstream_bias(
|
|
node: onnx.NodeProto,
|
|
initializers: dict[str, np.ndarray],
|
|
by_input: dict[str, list[onnx.NodeProto]],
|
|
) -> str | None:
|
|
frontier = list(node.output)
|
|
seen: set[int] = set()
|
|
for _ in range(5):
|
|
next_frontier: list[str] = []
|
|
for value in frontier:
|
|
for consumer in by_input.get(value, ()):
|
|
identity = id(consumer)
|
|
if identity in seen:
|
|
continue
|
|
seen.add(identity)
|
|
if consumer.op_type == "Add":
|
|
for input_name in consumer.input:
|
|
if input_name in initializers:
|
|
return input_name
|
|
next_frontier.extend(consumer.output)
|
|
frontier = next_frontier
|
|
return None
|
|
|
|
|
|
def dequantized_matmul_weights(
|
|
model: onnx.ModelProto,
|
|
*,
|
|
linear_pos_prefixes: Iterable[str] = (),
|
|
) -> dict[str, np.ndarray]:
|
|
initializers = tensor_map(model)
|
|
by_input = consumers(model)
|
|
linear_pos = iter(linear_pos_prefixes)
|
|
output: dict[str, np.ndarray] = {}
|
|
for node in model.graph.node:
|
|
if node.op_type != "MatMulInteger":
|
|
continue
|
|
if len(node.input) < 4:
|
|
raise ValueError("Kroko MatMulInteger is missing quantization inputs")
|
|
weight_name = node.input[1]
|
|
zero_name = node.input[3]
|
|
scale_name = weight_name.removesuffix("_quantized") + "_scale"
|
|
if weight_name not in initializers or zero_name not in initializers or scale_name not in initializers:
|
|
raise ValueError(f"missing quantization tensors for {weight_name}")
|
|
quantized = initializers[weight_name].astype(np.float32)
|
|
zero = initializers[zero_name].astype(np.float32)
|
|
scale = initializers[scale_name].astype(np.float32)
|
|
# ONNX MatMul stores [in, out], whereas audio.cpp Linear weights follow
|
|
# the PyTorch [out, in] convention.
|
|
weight = ((quantized - zero) * scale).T.copy()
|
|
bias_name = downstream_bias(node, initializers, by_input)
|
|
if bias_name is not None:
|
|
semantic_name = bias_name.removesuffix(".bias") + ".weight"
|
|
else:
|
|
try:
|
|
semantic_name = next(linear_pos) + ".self_attn_weights.linear_pos.weight"
|
|
except StopIteration as exc:
|
|
raise ValueError(f"cannot name bias-free MatMul weight {weight_name}") from exc
|
|
if semantic_name in output:
|
|
raise ValueError(f"duplicate recovered weight name: {semantic_name}")
|
|
output[semantic_name] = weight
|
|
try:
|
|
unexpected = next(linear_pos)
|
|
except StopIteration:
|
|
unexpected = None
|
|
if unexpected is not None:
|
|
raise ValueError(f"not all expected positional weights were recovered; first missing layer: {unexpected}")
|
|
return output
|
|
|
|
|
|
def recover_zipformer_constants(
|
|
model: onnx.ModelProto,
|
|
) -> dict[str, np.ndarray]:
|
|
"""Give semantic names to parameters folded by the ONNX exporter.
|
|
|
|
``SimpleDownsample.bias.softmax()`` is exported as a constant weight, and
|
|
the two learned chunkwise-convolution edge scales are anonymous
|
|
``onnx::Concat_*`` initializers. They are model parameters, not disposable
|
|
graph helpers, so preserve them under stable audio.cpp names.
|
|
"""
|
|
initializers = tensor_map(model)
|
|
recovered: dict[str, np.ndarray] = {}
|
|
|
|
downsample_weights: list[np.ndarray] = []
|
|
for node in model.graph.node:
|
|
if (
|
|
node.op_type == "Mul"
|
|
and node.name.startswith("/downsample")
|
|
and len(node.input) >= 2
|
|
and node.input[1] in initializers
|
|
):
|
|
value = initializers[node.input[1]]
|
|
if value.ndim == 3 and value.shape[1:] == (1, 1):
|
|
downsample_weights.append(
|
|
np.ascontiguousarray(value.reshape(-1).astype(np.float32))
|
|
)
|
|
if len(downsample_weights) != 6:
|
|
raise ValueError(
|
|
"expected five Zipformer stack downsamplers and one output "
|
|
f"downsampler, found {len(downsample_weights)}"
|
|
)
|
|
for stack, value in enumerate(downsample_weights[:-1], start=1):
|
|
recovered[
|
|
f"encoder.encoders.{stack}.downsample.weights"
|
|
] = value
|
|
recovered["encoder.downsample_output.weights"] = downsample_weights[-1]
|
|
|
|
nodes = list(model.graph.node)
|
|
chunk_scales = 0
|
|
for index, node in enumerate(nodes):
|
|
if (
|
|
node.op_type != "Conv"
|
|
or len(node.input) < 2
|
|
or not node.input[1].endswith(
|
|
".depthwise_conv.chunkwise_conv.weight"
|
|
)
|
|
):
|
|
continue
|
|
semantic_base = node.input[1].removesuffix(
|
|
".chunkwise_conv.weight"
|
|
)
|
|
conv_weight = initializers[node.input[1]]
|
|
if conv_weight.ndim != 3:
|
|
raise ValueError(
|
|
f"unexpected chunkwise convolution weight shape for {semantic_base}"
|
|
)
|
|
scale_shape = (conv_weight.shape[0], conv_weight.shape[2])
|
|
candidates: list[str] = []
|
|
# Exporters use either a nearby Concat (128-frame packages) or
|
|
# dynamic Slice nodes (64-frame packages) to select the learned left
|
|
# and right edge scales. In both layouts the two source
|
|
# initializers have the unambiguous [channels, kernel] shape.
|
|
for follower in nodes[index + 1 : index + 13]:
|
|
for input_name in follower.input:
|
|
value = initializers.get(input_name)
|
|
if value is not None and value.shape == scale_shape:
|
|
candidates.append(input_name)
|
|
candidates = list(dict.fromkeys(candidates))
|
|
if len(candidates) != 2:
|
|
raise ValueError(
|
|
"cannot recover two chunkwise edge scales for "
|
|
f"{semantic_base}; found {candidates}"
|
|
)
|
|
recovered[f"{semantic_base}.chunk_scale_left"] = np.ascontiguousarray(
|
|
initializers[candidates[0]].astype(np.float32)
|
|
)
|
|
recovered[f"{semantic_base}.chunk_scale_right"] = np.ascontiguousarray(
|
|
initializers[candidates[1]].astype(np.float32)
|
|
)
|
|
chunk_scales += 1
|
|
if chunk_scales != len(LAYER_PREFIXES) * 2:
|
|
raise ValueError(
|
|
"expected two chunkwise convolution scales per Zipformer layer; "
|
|
f"recovered {chunk_scales}"
|
|
)
|
|
return recovered
|
|
|
|
|
|
def convert_graph(
|
|
path: Path,
|
|
namespace: str,
|
|
*,
|
|
linear_pos_prefixes: Iterable[str] = (),
|
|
) -> dict[str, np.ndarray]:
|
|
model = onnx.load(path, load_external_data=True)
|
|
initializers = tensor_map(model)
|
|
recovered = dequantized_matmul_weights(
|
|
model, linear_pos_prefixes=linear_pos_prefixes
|
|
)
|
|
if namespace == "encoder":
|
|
recovered.update(recover_zipformer_constants(model))
|
|
producers = {
|
|
output_name: node
|
|
for node in model.graph.node
|
|
for output_name in node.output
|
|
}
|
|
for node in model.graph.node:
|
|
if node.op_type != "DequantizeLinear" or len(node.input) < 3:
|
|
continue
|
|
quant_name, scale_name, zero_name = node.input[:3]
|
|
if quant_name not in initializers:
|
|
producer = producers.get(quant_name)
|
|
if (
|
|
producer is not None
|
|
and producer.op_type == "Gather"
|
|
and producer.input
|
|
and producer.input[0] in initializers
|
|
):
|
|
quant_name = producer.input[0]
|
|
if (
|
|
quant_name not in initializers
|
|
or scale_name not in initializers
|
|
or zero_name not in initializers
|
|
):
|
|
continue
|
|
semantic_name = quant_name.removesuffix("_quantized")
|
|
recovered[semantic_name] = np.ascontiguousarray(
|
|
(initializers[quant_name].astype(np.float32)
|
|
- initializers[zero_name].astype(np.float32))
|
|
* initializers[scale_name].astype(np.float32)
|
|
)
|
|
# Skip only tensors that belong to an ONNX quantization tuple. A blanket
|
|
# ``*_scale`` filter would incorrectly remove real Zipformer parameters
|
|
# such as ``bypass_scale`` and ``norm.log_scale``.
|
|
quantization_helpers: set[str] = set()
|
|
for node in model.graph.node:
|
|
if node.op_type == "MatMulInteger" and len(node.input) >= 4:
|
|
weight_name = node.input[1]
|
|
quantization_helpers.update(
|
|
{
|
|
weight_name,
|
|
node.input[3],
|
|
weight_name.removesuffix("_quantized") + "_scale",
|
|
}
|
|
)
|
|
elif node.op_type == "DequantizeLinear" and len(node.input) >= 3:
|
|
quantization_helpers.update(node.input[:3])
|
|
quant_name = node.input[0]
|
|
producer = producers.get(quant_name)
|
|
if (
|
|
producer is not None
|
|
and producer.op_type == "Gather"
|
|
and producer.input
|
|
and producer.input[0] in initializers
|
|
):
|
|
quantization_helpers.add(producer.input[0])
|
|
result: dict[str, np.ndarray] = {}
|
|
for name, value in initializers.items():
|
|
if name in quantization_helpers or name.startswith("onnx::MatMul_"):
|
|
continue
|
|
# Safetensors represents rank-0 values as a one-element tensor so the
|
|
# same output can also be packed into GGUF.
|
|
if value.ndim == 0:
|
|
value = value.reshape(1)
|
|
result[f"{namespace}.{name}"] = np.ascontiguousarray(value)
|
|
for name, value in recovered.items():
|
|
result[f"{namespace}.{name}"] = np.ascontiguousarray(value)
|
|
return result
|
|
|
|
|
|
def metadata_from_encoder(path: Path) -> dict[str, object]:
|
|
model = onnx.load(path, load_external_data=False)
|
|
values = {item.key: item.value for item in model.metadata_props}
|
|
|
|
def ints(key: str) -> list[int]:
|
|
return [int(part) for part in values[key].split(",")]
|
|
|
|
left_context_len = ints("left_context_len")
|
|
base_left_context = left_context_len[0]
|
|
if any(
|
|
length <= 0 or base_left_context % length != 0
|
|
for length in left_context_len
|
|
):
|
|
raise ValueError(
|
|
"Kroko Zipformer left-context lengths do not describe integral "
|
|
"encoder downsampling factors"
|
|
)
|
|
|
|
return {
|
|
"model_type": values.get("model_type", "zipformer2"),
|
|
"sample_rate": 16000,
|
|
"feature_dim": 80,
|
|
"frame_length_ms": 25.0,
|
|
"frame_shift_ms": 10.0,
|
|
"dither": 0.0,
|
|
"snip_edges": False,
|
|
"normalize_samples": True,
|
|
"chunk_size": int(values["T"]),
|
|
"chunk_shift": int(values["decode_chunk_len"]),
|
|
"subsampling_factor": 4,
|
|
"encoder_dims": ints("encoder_dims"),
|
|
"query_head_dims": ints("query_head_dims"),
|
|
"value_head_dims": ints("value_head_dims"),
|
|
"num_heads": ints("num_heads"),
|
|
"num_encoder_layers": ints("num_encoder_layers"),
|
|
"cnn_module_kernels": ints("cnn_module_kernels"),
|
|
"left_context_len": left_context_len,
|
|
"downsampling_factors": [
|
|
base_left_context // length for length in left_context_len
|
|
],
|
|
}
|
|
|
|
|
|
def decoder_metadata(path: Path) -> dict[str, int]:
|
|
model = onnx.load(path, load_external_data=False)
|
|
values = {item.key: item.value for item in model.metadata_props}
|
|
return {
|
|
"vocab_size": int(values["vocab_size"]),
|
|
"context_size": int(values["context_size"]),
|
|
"blank_id": 0,
|
|
"unk_id": 2,
|
|
}
|
|
|
|
|
|
def main() -> None:
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("input", type=Path, help="Free Kroko .data package or extracted directory")
|
|
parser.add_argument("output", type=Path, help="Output audio.cpp model directory")
|
|
parser.add_argument("--overwrite", action="store_true")
|
|
args = parser.parse_args()
|
|
|
|
output = args.output.resolve()
|
|
weights_path = output / "model.safetensors"
|
|
if weights_path.exists() and not args.overwrite:
|
|
raise SystemExit(f"output exists (pass --overwrite): {weights_path}")
|
|
|
|
temporary: TemporaryDirectory[str] | None = None
|
|
if args.input.is_dir():
|
|
extracted = args.input.resolve()
|
|
header = {
|
|
"type": "zipformer2",
|
|
"free": True,
|
|
"language": {"name": "Unknown", "iso": "auto"},
|
|
}
|
|
else:
|
|
temporary = TemporaryDirectory(prefix="audiocpp-kroko-")
|
|
extracted = Path(temporary.name)
|
|
header = unpack_package(args.input.resolve(), extracted)
|
|
|
|
required = [extracted / name for name in ("encoder.onnx", "decoder.onnx", "joiner.onnx", "tokens.txt")]
|
|
missing = [str(path) for path in required if not path.is_file()]
|
|
if missing:
|
|
raise SystemExit("missing Kroko resources: " + ", ".join(missing))
|
|
|
|
tensors: dict[str, np.ndarray] = {}
|
|
tensors.update(
|
|
convert_graph(
|
|
extracted / "encoder.onnx",
|
|
"encoder",
|
|
linear_pos_prefixes=LAYER_PREFIXES,
|
|
)
|
|
)
|
|
tensors.update(convert_graph(extracted / "decoder.onnx", "decoder"))
|
|
tensors.update(convert_graph(extracted / "joiner.onnx", "joiner"))
|
|
|
|
config = metadata_from_encoder(extracted / "encoder.onnx")
|
|
config.update(decoder_metadata(extracted / "decoder.onnx"))
|
|
config["audiocpp_family"] = "kroko_asr"
|
|
config["variant"] = "Kroko-Community-Streaming"
|
|
config["language"] = header.get("language", {})
|
|
config["package_id"] = header.get("id", "")
|
|
config["source_format"] = "kroko-free-data"
|
|
|
|
output.mkdir(parents=True, exist_ok=True)
|
|
save_file(
|
|
tensors,
|
|
str(weights_path),
|
|
metadata={
|
|
"format": "pt",
|
|
"source": "Banafo/Kroko-ASR",
|
|
"audiocpp_family": "kroko_asr",
|
|
},
|
|
)
|
|
(output / "config.json").write_text(
|
|
json.dumps(config, indent=2, ensure_ascii=False) + "\n",
|
|
encoding="utf-8",
|
|
)
|
|
(output / "tokens.txt").write_bytes((extracted / "tokens.txt").read_bytes())
|
|
if temporary is not None:
|
|
temporary.cleanup()
|
|
print(f"wrote {weights_path} ({len(tensors)} tensors)")
|
|
print(f"wrote {output / 'config.json'}")
|
|
print(f"wrote {output / 'tokens.txt'}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|