File size: 3,911 Bytes
ad7a3dd | 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 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 | """Map the pinned Google checkpoint to native PyTorch state; no training."""
import argparse
import hashlib
import json
from pathlib import Path
import numpy as np
import torch
from naptrap.google import FORMAT, PARENT_SHA256, GoogleSpeechEmbedding
def digest(path):
return hashlib.sha256(path.read_bytes()).hexdigest()
def main():
parser = argparse.ArgumentParser()
parser.add_argument("--reference", type=Path, required=True)
parser.add_argument("--output", type=Path, required=True)
args = parser.parse_args()
report = json.loads((args.reference / "report.json").read_text())
if report["stage"] != "reference_captured" or report["source"]["archive_sha256"] != PARENT_SHA256:
raise ValueError("Authoritative Google reference required")
for file, field in (("variables.npz", "variable_npz_sha256"), ("frontend.npz", "frontend_npz_sha256")):
if digest(args.reference / file) != report[field]:
raise ValueError("Conversion input changed")
args.output.mkdir(parents=True, exist_ok=False)
variables = np.load(args.reference / "variables.npz", allow_pickle=False)
frontend = np.load(args.reference / "frontend.npz", allow_pickle=False)
model = GoogleSpeechEmbedding()
state = model.state_dict()
prefix = "embeddings/state/MudpuppyLite/"
used = {prefix + "ZeroMeanConv2d/weights"}
state["window"] = torch.from_numpy(frontend["window"].copy())
state["mel_matrix"] = torch.from_numpy(frontend["mel_matrix"].copy())
state["initial.weight"] = torch.from_numpy(frontend["initial_centered_hwio"].transpose(3, 2, 0, 1).copy())
for index in range(19):
name = prefix + ("Conv" if index == 0 else "Conv_" + str(index)) + "/weights"
target = f"convs.{index}.weight" if index < 18 else "output.weight"
state[target] = torch.from_numpy(variables[name].transpose(3, 2, 0, 1).copy())
used.add(name)
for index in range(19):
scope = prefix + ("BatchNorm" if index == 0 else "BatchNorm_" + str(index))
for old, new in (("gamma", "weight"), ("beta", "bias"), ("moving_mean", "running_mean"), ("moving_variance", "running_var")):
name = scope + "/" + old
state[f"norms.{index}.{new}"] = torch.from_numpy(variables[name].copy())
used.add(name)
if used != set(variables.files):
raise ValueError(f"Unconsumed or missing Google variables: {used.symmetric_difference(variables.files)}")
model.load_state_dict(state, strict=True)
path = args.output / "model.pt"
torch.save({"format": FORMAT, "parent_archive_sha256": PARENT_SHA256, "model": model.state_dict()}, path)
manifest = {"format": FORMAT, "parent_archive_sha256": PARENT_SHA256, "model_sha256": digest(path),
"license": "Apache-2.0", "publisher": "Google", "original": "google/speech_embedding/1",
"source": "https://www.kaggle.com/models/google/speech-embedding/TensorFlow1/speech-embedding/1",
"modifications": ["independent PyTorch graph implementation", "HWIO kernels transposed to OIHW",
"initial spatial-mean-subtracted kernel materialized exactly as original TF inference", "original frontend constants retained"],
"source_variable_count": len(used), "reference_report_sha256": digest(args.reference / "report.json"),
"parameters": sum(p.numel() for p in model.parameters()), "model_fitting": False,
"preprocessing": {"sample_rate": 16000, "leading_zeros": 11120, "frames": "floor(samples/1280)",
"stft": "400-sample Hann, hop160, FFT512, no centering, magnitude spectrum", "mel": "32 bins, natural log(mel+0.001)"}}
(args.output / "manifest.json").write_text(json.dumps(manifest, indent=2) + "\n")
print(json.dumps(manifest, indent=2))
if __name__ == "__main__":
main()
|