| """One-time TensorFlow reference/constant extraction for the PyTorch port.""" |
|
|
| import argparse |
| import hashlib |
| import json |
| import os |
| import sqlite3 |
| from pathlib import Path |
|
|
| os.environ.setdefault("CUDA_VISIBLE_DEVICES", "") |
| os.environ.setdefault("TF_ENABLE_ONEDNN_OPTS", "0") |
| os.environ.setdefault("TF_CPP_MIN_LOG_LEVEL", "2") |
|
|
| import numpy as np |
| import tensorflow as tf |
| from tensorflow.core.framework import graph_pb2 |
| from tensorflow.core.protobuf import saved_model_pb2 |
|
|
|
|
| def constant_value(graph_def, name): |
| """Evaluate the original constant-only subgraph, without map control edges.""" |
| nodes = {n.name: n for n in graph_def.node} |
| selected = set() |
| def add(key): |
| key = key.split(":")[0] |
| if key in selected: |
| return |
| node = nodes[key] |
| if node.op in ("Placeholder", "ReadVariableOp", "VarHandleOp"): |
| raise ValueError("Requested frontend constant depends on runtime state: " + key) |
| selected.add(key) |
| for parent in node.input: |
| if not parent.startswith("^"): |
| add(parent) |
| add(name) |
| subset = graph_pb2.GraphDef() |
| for node in graph_def.node: |
| if node.name in selected: |
| copied = subset.node.add() |
| copied.CopyFrom(node) |
| inputs = [v for v in copied.input if not v.startswith("^")] |
| del copied.input[:] |
| copied.input.extend(inputs) |
| with tf.Graph().as_default() as graph: |
| tf.import_graph_def(subset, name="") |
| with tf.compat.v1.Session(graph=graph) as session: |
| return session.run(graph.get_tensor_by_name(name + ":0")) |
|
|
|
|
| def main(): |
| parser = argparse.ArgumentParser() |
| parser.add_argument("--project", type=Path, required=True) |
| parser.add_argument("--output", type=Path, required=True) |
| args = parser.parse_args() |
| root = args.project.resolve() |
| output = args.output.resolve() |
| output.mkdir(parents=True, exist_ok=False) |
| source = root / "data/artifacts/google-speech-embedding" |
| identity = json.loads((source / ".source.json").read_text()) |
| for name, expected in identity["files"].items(): |
| if hashlib.sha256((source / name).read_bytes()).hexdigest() != expected: |
| raise ValueError("Original artifact changed") |
| saved = saved_model_pb2.SavedModel() |
| saved.ParseFromString((source / "saved_model.pb").read_bytes()) |
| meta = next(m for m in saved.meta_graphs if not m.meta_info_def.tags) |
| window = constant_value(meta.graph_def, "map/while/stft/hann_window/sub_2") |
| mel = constant_value(meta.graph_def, "map/while/linear_to_mel_weight_matrix") |
| reader = tf.train.load_checkpoint(str(source / "variables/variables")) |
| weights = {k: reader.get_tensor(k) for k in sorted(reader.get_variable_to_shape_map())} |
| key = "embeddings/state/MudpuppyLite/ZeroMeanConv2d/weights" |
| initial = tf.convert_to_tensor(weights[key]) |
| centered = (initial - tf.reduce_mean(initial, axis=[0, 1], keepdims=True)).numpy() |
| np.savez(output / "frontend.npz", window=window, mel_matrix=mel, initial_centered_hwio=centered) |
| np.savez(output / "variables.npz", **weights) |
| manifest = json.loads((root / "data/derived/development-clean460-v1/manifest.json").read_text()) |
| crops = sorted((r for r in manifest["query_crops"].values() if r["split"] == "dev-select"), key=lambda r: (r["num_samples"], r["crop_id"])) |
| positions = np.unique(np.linspace(0, len(crops) - 1, 32, dtype=int)) |
| db = sqlite3.connect((root / "data/cache/keyword-clean460-google-v1/features.sqlite3").as_uri() + "?mode=ro&immutable=1", uri=True) |
| examples = [] |
| for i in positions: |
| crop = crops[i] |
| row = db.execute("SELECT pcm16,data,valid_length,sha256 FROM features WHERE crop_id=?", (crop["crop_id"],)).fetchone() |
| if hashlib.sha256(row[0]).hexdigest() != crop["audio_sha256"] or hashlib.sha256(row[1]).hexdigest() != row[3]: |
| raise ValueError("Retained feature/crop identity changed") |
| wave = np.frombuffer(row[0], dtype="<i2").astype(np.float32) / 32768 |
| cached = np.frombuffer(row[1], dtype="<f4").copy().reshape(row[2], 96) |
| examples.append(("speech-" + crop["crop_id"], wave, cached)) |
| db.close() |
| rng = np.random.default_rng(20261011) |
| for length in (8000, 8001, 8960, 16000, 31999, 32000): |
| examples.append((f"silence-{length}", np.zeros(length, dtype=np.float32), None)) |
| for length in (8000, 16000, 32000): |
| noise = (rng.standard_normal(length) * .08).clip(-1, 1).astype(np.float32) |
| tone = (.15 * np.sin(np.arange(length, dtype=np.float32) * np.float32(2 * np.pi * 440 / 16000))).astype(np.float32) |
| examples += [(f"noise-{length}", noise, None), (f"tone-{length}", tone, None)] |
| impulse = np.zeros(16000, dtype=np.float32) |
| impulse[0], impulse[7999], impulse[-1] = .9, -.9, .9 |
| examples.append(("impulses-16000", impulse, None)) |
| configuration = tf.compat.v1.ConfigProto(intra_op_parallelism_threads=4, inter_op_parallelism_threads=1, device_count={"GPU": 0}) |
| records = [] |
| with tf.Graph().as_default() as graph: |
| with tf.compat.v1.Session(graph=graph, config=configuration) as session: |
| tf.compat.v1.saved_model.loader.load(session, [], str(source)) |
| inputs = graph.get_tensor_by_name(meta.signature_def["default"].inputs["default"].name) |
| endpoint = graph.get_tensor_by_name(meta.signature_def["default"].outputs["default"].name) |
| logmel = graph.get_tensor_by_name("Squeeze:0") |
| for name, wave, cached in examples: |
| features, mel_values = session.run([endpoint, logmel], {inputs: np.pad(wave, (11120, 0))[None]}) |
| features = features[0, :, 0, :] |
| if features.shape != (len(wave) // 1280, 96) or not np.isfinite(features).all(): |
| raise ValueError("Reference feature shape/finite failure") |
| arrays = {"waveform": wave, "features": features, "log_mel": mel_values[0]} |
| if cached is not None: |
| arrays["cached_source_features"] = cached |
| np.savez(output / (name + ".npz"), **arrays) |
| records.append({"name": name, "samples": len(wave), "frames": len(features), |
| "cached_source_max_abs": float(np.max(np.abs(cached - features))) if cached is not None else None}) |
| report = {"stage": "reference_captured", "tensorflow": tf.__version__, "numpy": np.__version__, "source": identity, |
| "role": "32 metadata-selected dev-select query crops plus deterministic engineering stress fixtures", |
| "examples": records, "variables": len(weights), "window_shape": list(window.shape), "mel_shape": list(mel.shape), |
| "variable_npz_sha256": hashlib.sha256((output / "variables.npz").read_bytes()).hexdigest(), |
| "frontend_npz_sha256": hashlib.sha256((output / "frontend.npz").read_bytes()).hexdigest(), |
| "model_fitting": False, "final_evaluation": False} |
| (output / "report.json").write_text(json.dumps(report, indent=2) + "\n") |
| print(json.dumps({k: v for k, v in report.items() if k not in ("examples", "source")}, indent=2)) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
