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
95 lines
3.9 KiB
Python
95 lines
3.9 KiB
Python
#!/usr/bin/env python3
|
|
"""Convert an official ESPnet OWSM-CTC v4 checkpoint to GGUF."""
|
|
|
|
import argparse
|
|
import json
|
|
import shutil
|
|
import subprocess
|
|
from pathlib import Path
|
|
|
|
import torch
|
|
import yaml
|
|
from safetensors.torch import save_file
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(description=__doc__)
|
|
parser.add_argument("source", type=Path, help="Official Hugging Face snapshot root")
|
|
parser.add_argument("output", type=Path)
|
|
parser.add_argument("--staging-dir", required=True, type=Path)
|
|
parser.add_argument("--converter", type=Path, default=Path("build/debug/bin/audiocpp_gguf"))
|
|
parser.add_argument("--spec", type=Path, default=Path("model_specs/owsm_ctc.json"))
|
|
parser.add_argument("--type", choices=("orig", "q8_0", "q4_k"), default="orig")
|
|
parser.add_argument("--overwrite", action="store_true")
|
|
args = parser.parse_args()
|
|
|
|
configs = list(args.source.glob("exp/s2t_train_*/config.yaml"))
|
|
checkpoints = list(args.source.glob("exp/s2t_train_*/valid.total_count.ave_5best*.pth"))
|
|
if len(configs) != 1 or len(checkpoints) != 1:
|
|
raise ValueError("Expected one official OWSM-CTC v4 config and one ave-5 checkpoint")
|
|
config = yaml.safe_load(configs[0].read_text())
|
|
encoder = config["encoder_conf"]
|
|
if config["encoder"] != "e_branchformer_ctc" or (encoder["output_size"], encoder["num_blocks"]) != (1024, 27):
|
|
raise ValueError("Expected the OWSM-CTC v4 1B encoder")
|
|
variant = "medium"
|
|
if encoder["linear_units"] != encoder["output_size"] * 4:
|
|
raise ValueError("OWSM-CTC v4 requires a 4x feed-forward width")
|
|
|
|
state = torch.load(checkpoints[0], map_location="cpu", weights_only=True)
|
|
tensors = {
|
|
name: tensor.detach().cpu().contiguous()
|
|
for name, tensor in state.items()
|
|
if name != "frontend.logmel.melmat"
|
|
}
|
|
args.staging_dir.mkdir(parents=True, exist_ok=True)
|
|
save_file(tensors, args.staging_dir / "model.safetensors")
|
|
tokenizer = args.source / config["bpemodel"]
|
|
shutil.copyfile(tokenizer, args.staging_dir / "bpe.model")
|
|
runtime_config = {
|
|
"variant": variant,
|
|
"hidden_size": encoder["output_size"],
|
|
"num_heads": encoder["attention_heads"],
|
|
"encoder_layers": encoder["num_blocks"],
|
|
"intermediate_size": encoder["linear_units"],
|
|
"vocabulary_size": len(config["token_list"]),
|
|
"max_audio_samples": 480000,
|
|
"frontend_frames": 3001,
|
|
"encoder_frames": 374,
|
|
}
|
|
prompt = config["promptencoder_conf"]
|
|
runtime_config.update({
|
|
"prompt_hidden_size": prompt["output_size"],
|
|
"prompt_num_heads": prompt["attention_heads"],
|
|
"prompt_layers": prompt["num_blocks"],
|
|
"prompt_intermediate_size": prompt["linear_units"],
|
|
"interctc_layers": encoder["interctc_layer_idx"],
|
|
"cross_attention_layers": [
|
|
i for i, enabled in enumerate(encoder["use_cross_attention"]) if enabled
|
|
],
|
|
})
|
|
(args.staging_dir / "config.json").write_text(json.dumps(runtime_config, indent=2) + "\n")
|
|
|
|
command = [
|
|
str(args.converter.resolve()),
|
|
"--input", str((args.staging_dir / "model.safetensors").resolve()),
|
|
"--root", str(args.staging_dir.resolve()),
|
|
"--output", str(args.output.resolve()),
|
|
"--type", args.type,
|
|
"--family", "owsm_ctc",
|
|
"--model-spec", str(args.spec.resolve()),
|
|
"--sidecar", f"{(args.staging_dir / 'config.json').resolve()}=config.json",
|
|
"--sidecar", f"{(args.staging_dir / 'bpe.model').resolve()}=bpe.model",
|
|
]
|
|
if args.type != "orig":
|
|
for name, tensor in tensors.items():
|
|
if tensor.ndim != 2:
|
|
command.extend(("--keep-type", f"{name}=f32"))
|
|
if args.overwrite:
|
|
command.append("--overwrite")
|
|
args.output.parent.mkdir(parents=True, exist_ok=True)
|
|
subprocess.run(command, check=True)
|
|
print(f"Converted OWSM-CTC v4 {variant}: {len(tensors)} tensors, storage={args.type}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|