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
126 lines
4.2 KiB
Python
126 lines
4.2 KiB
Python
#!/usr/bin/env python3
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import json
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
import torch
|
|
from safetensors.torch import load_file, save_file
|
|
|
|
|
|
REPO_ROOT = Path(__file__).resolve().parents[2]
|
|
|
|
|
|
def parse_args() -> argparse.Namespace:
|
|
parser = argparse.ArgumentParser(description="Convert OpenAI Whisper medium.pt to safetensors for audio.cpp.")
|
|
parser.add_argument(
|
|
"--ckpt",
|
|
default=str(Path.home() / ".cache" / "whisper" / "medium.pt"),
|
|
help="Path to OpenAI Whisper medium.pt.",
|
|
)
|
|
parser.add_argument(
|
|
"--output-dir",
|
|
default="models/whisper-medium",
|
|
help="Directory to write model.safetensors and config.json.",
|
|
)
|
|
return parser.parse_args()
|
|
|
|
|
|
def resolve_path(path: str) -> Path:
|
|
candidate = Path(path).expanduser()
|
|
if candidate.is_absolute():
|
|
return candidate
|
|
return REPO_ROOT / candidate
|
|
|
|
|
|
def require_dims(payload: dict[str, Any]) -> dict[str, int]:
|
|
raw_dims = payload.get("dims")
|
|
if not isinstance(raw_dims, dict):
|
|
raise RuntimeError("Whisper checkpoint is missing dims")
|
|
dims: dict[str, int] = {}
|
|
for key, value in raw_dims.items():
|
|
if not isinstance(key, str) or not isinstance(value, int):
|
|
raise RuntimeError(f"invalid Whisper dim entry: {key!r}={value!r}")
|
|
dims[key] = value
|
|
return dims
|
|
|
|
|
|
def require_state_dict(payload: dict[str, Any]) -> dict[str, torch.Tensor]:
|
|
raw_state = payload.get("model_state_dict")
|
|
if not isinstance(raw_state, dict):
|
|
raise RuntimeError("Whisper checkpoint is missing model_state_dict")
|
|
state: dict[str, torch.Tensor] = {}
|
|
for key, value in raw_state.items():
|
|
if not isinstance(key, str):
|
|
raise RuntimeError(f"invalid non-string Whisper tensor key: {key!r}")
|
|
if not isinstance(value, torch.Tensor):
|
|
raise RuntimeError(f"invalid non-tensor Whisper state entry: {key}")
|
|
state[key] = value.detach().cpu().contiguous()
|
|
return state
|
|
|
|
|
|
def export_config(dims: dict[str, int], tensor_count: int) -> dict[str, Any]:
|
|
return {
|
|
"model_type": "whisper",
|
|
"variant": "medium",
|
|
"source_format": "openai_whisper_pt",
|
|
"weight_file": "model.safetensors",
|
|
"tensor_count": tensor_count,
|
|
"dims": dims,
|
|
"audio_encoder": {
|
|
"n_mels": dims["n_mels"],
|
|
"n_audio_ctx": dims["n_audio_ctx"],
|
|
"n_audio_state": dims["n_audio_state"],
|
|
"n_audio_head": dims["n_audio_head"],
|
|
"n_audio_layer": dims["n_audio_layer"],
|
|
},
|
|
"text_decoder": {
|
|
"n_vocab": dims["n_vocab"],
|
|
"n_text_ctx": dims["n_text_ctx"],
|
|
"n_text_state": dims["n_text_state"],
|
|
"n_text_head": dims["n_text_head"],
|
|
"n_text_layer": dims["n_text_layer"],
|
|
},
|
|
}
|
|
|
|
|
|
def main() -> int:
|
|
args = parse_args()
|
|
ckpt_path = resolve_path(args.ckpt)
|
|
output_dir = resolve_path(args.output_dir)
|
|
if not ckpt_path.is_file():
|
|
raise RuntimeError(f"Whisper checkpoint does not exist: {ckpt_path}")
|
|
|
|
payload = torch.load(ckpt_path, map_location="cpu")
|
|
if not isinstance(payload, dict):
|
|
raise RuntimeError(f"unexpected Whisper checkpoint payload type: {type(payload)}")
|
|
|
|
dims = require_dims(payload)
|
|
state = require_state_dict(payload)
|
|
|
|
output_dir.mkdir(parents=True, exist_ok=True)
|
|
model_path = output_dir / "model.safetensors"
|
|
save_file(state, str(model_path))
|
|
|
|
reloaded = load_file(str(model_path))
|
|
if set(reloaded.keys()) != set(state.keys()):
|
|
raise RuntimeError("saved Whisper safetensors key set does not match source checkpoint")
|
|
for key, value in state.items():
|
|
if not torch.equal(value, reloaded[key]):
|
|
raise RuntimeError(f"saved Whisper safetensors changed tensor: {key}")
|
|
|
|
(output_dir / "config.json").write_text(
|
|
json.dumps(export_config(dims, len(state)), indent=2, ensure_ascii=False) + "\n",
|
|
encoding="utf-8",
|
|
)
|
|
print(f"ckpt={ckpt_path}")
|
|
print(f"output_dir={output_dir}")
|
|
print(f"tensor_count={len(state)}")
|
|
print(f"model_safetensors={model_path}")
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|