| |
| """Export Paraformer to GGUF with experimental Q5_0/Q4_0 support.""" |
|
|
| import argparse |
| import glob |
| import json |
| import os |
| import re |
|
|
| import gguf |
| import numpy as np |
| import torch |
|
|
|
|
| def parse_mvn(path): |
| blocks = [ |
| np.array([float(x) for x in block.split()], np.float32) |
| for block in re.findall(r"\[([^\]]*)\]", open(path).read()) |
| ] |
| vectors = [block for block in blocks if block.size > 1] |
| return vectors[0], vectors[1] |
|
|
|
|
| def main(): |
| parser = argparse.ArgumentParser() |
| parser.add_argument("--model_pt", required=True) |
| parser.add_argument("--mvn", required=True) |
| parser.add_argument("--out", required=True) |
| parser.add_argument( |
| "--wtype", |
| default="q5_0", |
| choices=["f32", "f16", "q8_0", "q5_0", "q4_0"], |
| ) |
| parser.add_argument("--tokens", default=None) |
| args = parser.parse_args() |
|
|
| state_dict = torch.load(args.model_pt, map_location="cpu") |
| state_dict = state_dict.get("state_dict", state_dict) |
| writer = gguf.GGUFWriter(args.out, "paraformer") |
| writer.add_uint32("pf.enc.output_size", 512) |
| writer.add_uint32("pf.enc.attention_heads", 4) |
| writer.add_uint32("pf.enc.num_blocks", 50) |
| writer.add_uint32("pf.enc.kernel_size", 11) |
| writer.add_uint32("pf.dec.num_blocks", 16) |
| writer.add_uint32("pf.dec.att_layer_num", 16) |
| writer.add_uint32("pf.dec.decoders3", 1) |
| writer.add_uint32("pf.dec.attention_heads", 4) |
| writer.add_uint32("pf.dec.kernel_size", 11) |
| writer.add_uint32("pf.vocab_size", 8404) |
|
|
| token_path = args.tokens or ( |
| glob.glob(os.path.join(os.path.dirname(args.model_pt), "tokens.json")) + [None] |
| )[0] |
| if token_path and os.path.exists(token_path): |
| with open(token_path, encoding="utf-8") as token_file: |
| tokens = json.load(token_file) |
| writer.add_array("pf.vocab", tokens) |
| print(f"embedded pf.vocab ({len(tokens)} tokens) from {token_path}") |
| else: |
| print("WARNING: tokens.json not found - GGUF will have no vocabulary") |
|
|
| writer.add_float32("pf.predictor.tail_threshold", 0.45) |
| writer.add_float32("pf.predictor.threshold", 1.0) |
| shift, scale = parse_mvn(args.mvn) |
| writer.add_tensor("cmvn.shift", shift) |
| writer.add_tensor("cmvn.scale", scale) |
|
|
| from gguf import GGMLQuantizationType as QuantType |
| from gguf import quants |
|
|
| quant_types = { |
| "q8_0": QuantType.Q8_0, |
| "q5_0": QuantType.Q5_0, |
| "q4_0": QuantType.Q4_0, |
| } |
| quant_type = quant_types.get(args.wtype) |
| quant_block_size = ( |
| quants._type_traits[quant_type].block_size if quant_type is not None else 1 |
| ) |
| tensor_count = 0 |
| quantized_count = 0 |
|
|
| for name, value in state_dict.items(): |
| if not name.startswith(("encoder.", "decoder.", "predictor.")): |
| continue |
| if name == "decoder.embed.0.weight": |
| continue |
|
|
| array = value.detach().to(torch.float32).contiguous().numpy() |
| if name.endswith("fsmn_block.weight") and array.ndim == 3: |
| array = np.ascontiguousarray(array[:, 0, :].T) |
| elif ( |
| args.wtype == "f16" |
| and array.ndim == 2 |
| and "norm" not in name |
| and "cif_output" not in name |
| ): |
| array = array.astype(np.float16) |
|
|
| can_quantize = ( |
| quant_type is not None |
| and array.ndim == 2 |
| and "norm" not in name |
| and "fsmn_block" not in name |
| and "predictor" not in name |
| and array.shape[1] % quant_block_size == 0 |
| ) |
| if can_quantize: |
| writer.add_tensor( |
| name, |
| quants.quantize(array, quant_type), |
| raw_dtype=quant_type, |
| ) |
| quantized_count += 1 |
| else: |
| writer.add_tensor(name, array) |
| tensor_count += 1 |
|
|
| print( |
| f"writing {tensor_count} tensors (+cmvn), {quantized_count} quantized " |
| f"as {args.wtype}, to {args.out}" |
| ) |
| writer.write_header_to_file() |
| writer.write_kv_data_to_file() |
| writer.write_tensors_to_file() |
| writer.close() |
| print(f"done: {args.out} ({os.path.getsize(args.out) / 1e6:.1f} MB)") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
