wehub-resource-sync
78 lines
3.4 KiB
Python
78 lines
3.4 KiB
Python
#!/usr/bin/env python3
|
|
"""Export SenseVoiceSmall (encoder + CTC head + query embeddings + CMVN) to GGUF
|
|
for the ggml C++ runtime. The encoder is the same SAN-M architecture as
|
|
Fun-ASR-Nano, so the C++ forward is shared.
|
|
"""
|
|
import argparse, os, re
|
|
import numpy as np, torch, gguf
|
|
|
|
|
|
def parse_mvn(path):
|
|
"""am.mvn (kaldi nnet): two `<LearnRateCoef> 0 [ ... ]` blocks -> shift, scale.
|
|
apply: out = (in + shift) * scale, per-dim (560)."""
|
|
txt = open(path).read()
|
|
blocks = re.findall(r"\[([^\]]*)\]", txt)
|
|
shift = np.array([float(x) for x in blocks[0].split()], dtype=np.float32)
|
|
scale = np.array([float(x) for x in blocks[1].split()], dtype=np.float32)
|
|
return shift, scale
|
|
|
|
|
|
def main():
|
|
ap = argparse.ArgumentParser()
|
|
ap.add_argument("--model_pt", required=True)
|
|
ap.add_argument("--mvn", required=True)
|
|
ap.add_argument("--out", required=True)
|
|
ap.add_argument("--wtype", default="f32", choices=["f32", "f16", "q8_0"])
|
|
ap.add_argument("--spm", default=None, help="sentencepiece .bpe.model; default: next to model_pt")
|
|
args = ap.parse_args()
|
|
|
|
sd = torch.load(args.model_pt, map_location="cpu"); sd = sd.get("state_dict", sd)
|
|
w = gguf.GGUFWriter(args.out, "sensevoice-small")
|
|
w.add_uint32("sv.input_size", 560)
|
|
w.add_uint32("sv.output_size", 512)
|
|
w.add_uint32("sv.attention_heads", 4)
|
|
w.add_uint32("sv.num_blocks", 50)
|
|
w.add_uint32("sv.tp_blocks", 20)
|
|
w.add_uint32("sv.kernel_size", 11)
|
|
w.add_uint32("sv.vocab_size", 25055)
|
|
w.add_uint32("sv.blank_id", 0)
|
|
# query token embed indices used at inference: [lid(auto=0), 1, 2, textnorm(woitn=15)]
|
|
w.add_array("sv.query_tokens", [0, 1, 2, 14]) # 14=withitn (use_itn=True), matches authoritative
|
|
import glob
|
|
spm_path = args.spm or (glob.glob(os.path.join(os.path.dirname(args.model_pt), "*.bpe.model")) + [None])[0]
|
|
if spm_path and os.path.exists(spm_path):
|
|
import sentencepiece as spm
|
|
sp = spm.SentencePieceProcessor(model_file=spm_path)
|
|
pieces = [sp.id_to_piece(i) for i in range(sp.get_piece_size())]
|
|
w.add_array("sv.vocab", pieces)
|
|
print(f"embedded sv.vocab ({len(pieces)} pieces) from {spm_path}")
|
|
else:
|
|
print("WARNING: *.bpe.model not found - gguf will have no vocab (binary falls back to ids)")
|
|
|
|
shift, scale = parse_mvn(args.mvn)
|
|
w.add_tensor("cmvn.shift", shift) # (560,)
|
|
w.add_tensor("cmvn.scale", scale) # (560,)
|
|
|
|
n = 0
|
|
for k, v in sd.items():
|
|
if not (k.startswith("encoder.") or k.startswith("ctc.") or k == "embed.weight"):
|
|
continue
|
|
arr = v.detach().to(torch.float32).contiguous().numpy()
|
|
if k.endswith("fsmn_block.weight"): # (D,1,K) -> (K,D)
|
|
arr = np.ascontiguousarray(arr[:, 0, :].T)
|
|
elif args.wtype == "f16" and arr.ndim == 2 and "norm" not in k:
|
|
arr = arr.astype(np.float16)
|
|
if args.wtype == "q8_0" and arr.ndim == 2 and "norm" not in k and "fsmn_block" not in k and k != "embed.weight" and arr.shape[1] % 32 == 0:
|
|
from gguf import quants as _q, GGMLQuantizationType as _QT
|
|
w.add_tensor(k, _q.quantize(arr, _QT.Q8_0), raw_dtype=_QT.Q8_0)
|
|
else:
|
|
w.add_tensor(k, arr)
|
|
n += 1
|
|
print(f"writing {n} tensors (+cmvn) to {args.out}")
|
|
w.write_header_to_file(); w.write_kv_data_to_file(); w.write_tensors_to_file(); w.close()
|
|
print(f"done: {args.out} ({os.path.getsize(args.out)/1e6:.1f} MB)")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|