File size: 2,264 Bytes
fc55559 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 | """Convert g2p_en 2.1.0 checkpoint20.npz to naptrap's safetensors layout.
Usage: python convert.py path/to/g2p_en/checkpoint20.npz OUTPUT_DIR
Arrays are copied unchanged (float32); only their names change.
"""
import hashlib
import json
import sys
from pathlib import Path
import numpy as np
import torch
from safetensors.torch import save_file
SOURCE_SHA256 = "b8af35e4596d8dd5836dfd3fe9b2ba4f97b9c311efe8879544cbcfcbd566d8c6"
NAMES = {
"enc_emb": "enc_emb", "dec_emb": "dec_emb", "fc_w": "fc_w", "fc_b": "fc_b",
**{f"{side}_{p}": f"{'encoder' if side == 'enc' else 'decoder'}.{p}"
for side in ("enc", "dec") for p in ("w_ih", "w_hh", "b_ih", "b_hh")},
}
def main(source, output):
source, output = Path(source), Path(output)
if hashlib.sha256(source.read_bytes()).hexdigest() != SOURCE_SHA256:
raise SystemExit("Unexpected checkpoint20.npz; expected the g2p_en 2.1.0 release")
arrays = np.load(source, allow_pickle=False)
if set(arrays.files) != set(NAMES):
raise SystemExit(f"Unexpected arrays: {sorted(arrays.files)}")
state = {NAMES[k]: torch.from_numpy(np.ascontiguousarray(arrays[k], dtype=np.float32)) for k in arrays.files}
output.mkdir(parents=True, exist_ok=True)
weights = output / "model.safetensors"
save_file(state, str(weights), metadata={"format": "naptrap-g2p-en-decoder-v1"})
config = {
"format": "naptrap-g2p-en-decoder-v1",
"hidden_size": int(arrays["dec_emb"].shape[1]),
"graphemes": int(arrays["enc_emb"].shape[0]),
"phonemes": int(arrays["dec_emb"].shape[0]),
"max_decoder_steps": 20,
"output": "greedy decoder hidden states per word, separated by the <pad> phoneme embedding",
"source": {"package": "g2p_en", "version": "2.1.0", "file": "g2p_en/checkpoint20.npz",
"sha256": SOURCE_SHA256, "license": "Apache-2.0"},
"sha256": {"model.safetensors": hashlib.sha256(weights.read_bytes()).hexdigest()},
}
(output / "config.json").write_text(json.dumps(config, indent=2) + "\n")
for k in arrays.files:
if not np.array_equal(arrays[k], state[NAMES[k]].numpy()):
raise SystemExit("Conversion changed " + k)
if __name__ == "__main__":
main(*sys.argv[1:])
|