| """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() |
|
|
