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
63 lines
2.9 KiB
Python
63 lines
2.9 KiB
Python
#!/usr/bin/env python3
|
|
"""Package an official mispeech CED checkpoint with complete AudioSet labels."""
|
|
|
|
import argparse
|
|
import csv
|
|
import json
|
|
from pathlib import Path
|
|
import shutil
|
|
import subprocess
|
|
import tempfile
|
|
|
|
import torchaudio
|
|
from safetensors.torch import load_file, save_file
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(description=__doc__)
|
|
parser.add_argument("--source", type=Path, required=True)
|
|
parser.add_argument("--labels-csv", type=Path, required=True)
|
|
parser.add_argument("--output", type=Path, required=True)
|
|
parser.add_argument("--converter", type=Path, default=Path("build/debug/bin/audiocpp_gguf"))
|
|
parser.add_argument("--type", default="orig")
|
|
parser.add_argument("--overwrite", action="store_true")
|
|
args = parser.parse_args()
|
|
source = args.source.resolve()
|
|
config = json.loads((source / "config.json").read_text())
|
|
with args.labels_csv.open(newline="", encoding="utf-8") as file:
|
|
labels = {row["index"]: row["display_name"] for row in csv.DictReader(file)}
|
|
if set(labels) != {str(i) for i in range(config["outputdim"])}:
|
|
raise ValueError("AudioSet labels must cover every classifier output")
|
|
# HF's CedConfig replaces the serialized labels at load time. Some of the
|
|
# serialized names are truncated at commas; preserve the actual CSV labels.
|
|
config["id2label"] = labels
|
|
config["label2id"] = {label: int(index) for index, label in labels.items()}
|
|
processor = json.loads((source / "preprocessor_config.json").read_text())
|
|
frontend = torchaudio.transforms.MelSpectrogram(
|
|
sample_rate=processor["sampling_rate"], n_fft=processor["n_fft"],
|
|
win_length=processor["win_size"], hop_length=processor["hop_size"],
|
|
f_min=processor["f_min"], f_max=processor["f_max"],
|
|
n_mels=processor["feature_size"], center=processor["center"],
|
|
)
|
|
tensors = load_file(source / "model.safetensors")
|
|
tensors["frontend.window"] = frontend.spectrogram.window.contiguous()
|
|
tensors["frontend.mel_filterbank"] = frontend.mel_scale.fb.T.contiguous()
|
|
args.output.parent.mkdir(parents=True, exist_ok=True)
|
|
with tempfile.TemporaryDirectory(prefix="ced-gguf-") as directory:
|
|
root = Path(directory)
|
|
(root / "config.json").write_text(json.dumps(config, indent=2) + "\n")
|
|
shutil.copyfile(source / "preprocessor_config.json", root / "preprocessor_config.json")
|
|
save_file(tensors, root / "model.safetensors")
|
|
command = [
|
|
str(args.converter.resolve()), "--input", str(root / "model.safetensors"),
|
|
"--root", str(root), "--family", "ced", "--model-spec",
|
|
str(Path(__file__).resolve().parents[2] / "model_specs/ced.json"),
|
|
"--type", args.type, "--output", str(args.output.resolve()),
|
|
]
|
|
if args.overwrite:
|
|
command.append("--overwrite")
|
|
subprocess.run(command, check=True)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|