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
68 lines
2.5 KiB
Python
68 lines
2.5 KiB
Python
#!/usr/bin/env python3
|
|
"""Package the official MIT AST AudioSet Safetensors checkpoint as GGUF."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import json
|
|
import subprocess
|
|
from pathlib import Path
|
|
|
|
|
|
SOURCE_REPO = "MIT/ast-finetuned-audioset-10-10-0.4593"
|
|
|
|
|
|
def validate_source(source: Path) -> None:
|
|
config = json.loads((source / "config.json").read_text(encoding="utf-8"))
|
|
processor = json.loads((source / "preprocessor_config.json").read_text(encoding="utf-8"))
|
|
expected = {
|
|
"hidden_size": 768,
|
|
"num_hidden_layers": 12,
|
|
"num_attention_heads": 12,
|
|
"intermediate_size": 3072,
|
|
"patch_size": 16,
|
|
"frequency_stride": 10,
|
|
"time_stride": 10,
|
|
"max_length": 1024,
|
|
"num_mel_bins": 128,
|
|
}
|
|
mismatches = [f"{key}={config.get(key)!r}, expected {value!r}" for key, value in expected.items()
|
|
if config.get(key) != value]
|
|
if len(config.get("id2label", {})) != 527:
|
|
mismatches.append("id2label must contain 527 classes")
|
|
if processor.get("sampling_rate") != 16000 or processor.get("feature_extractor_type") != "ASTFeatureExtractor":
|
|
mismatches.append("unexpected AST feature extractor")
|
|
if mismatches:
|
|
raise ValueError("not the expected MIT AST AudioSet checkpoint: " + "; ".join(mismatches))
|
|
if not (source / "model.safetensors").is_file():
|
|
raise FileNotFoundError(source / "model.safetensors")
|
|
|
|
|
|
def main() -> int:
|
|
parser = argparse.ArgumentParser(description=__doc__)
|
|
parser.add_argument("--source", type=Path, required=True)
|
|
parser.add_argument("--converter", type=Path, default=Path("build/debug/bin/audiocpp_gguf"))
|
|
parser.add_argument("--output", type=Path, required=True)
|
|
parser.add_argument("--type", choices=["orig", "f16", "q8_0", "q4_0", "q4_k"], default="orig")
|
|
args = parser.parse_args()
|
|
source = args.source.resolve()
|
|
validate_source(source)
|
|
args.output.resolve().parent.mkdir(parents=True, exist_ok=True)
|
|
command = [
|
|
str(args.converter.resolve()),
|
|
"--input", str((source / "model.safetensors").resolve()),
|
|
"--root", str(source),
|
|
"--family", "ast_audioset",
|
|
"--model-spec", str(Path("model_specs/ast_audioset.json").resolve()),
|
|
"--type", args.type,
|
|
"--output", str(args.output.resolve()),
|
|
"--overwrite",
|
|
]
|
|
print("source_repo=" + SOURCE_REPO)
|
|
print("+", " ".join(command))
|
|
subprocess.run(command, check=True)
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|