diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 43ac40e..40cb2f9 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -4,7 +4,6 @@ on: push: branches: [main] pull_request: - branches: [main] concurrency: group: ci-${{ github.ref }} diff --git a/astro-site/src/content/docs/deployment.md b/astro-site/src/content/docs/deployment.md index 1b05387..85534f0 100644 --- a/astro-site/src/content/docs/deployment.md +++ b/astro-site/src/content/docs/deployment.md @@ -70,6 +70,43 @@ Use `load_codec()` for the common interface across families. The lower-level exa The manifest is the source of truth. The runtime reads it first, then resolves the encoder, optional decoder, and codebook files from the names listed there. +## RVQ Codebook Export and Training Parity + +When exporting `vq.get_weights()`, declare `rvq_num_levels`, `rvq_use_ema`, and +`rvq_kmeans_init` from the trained model's configuration. Plain RVQ has one +codebook matrix per level. EMA RVQ also stores count vectors and embedding-sum +matrices, plus an optional k-means flag; these are training state, not additional +codebooks. Export rejects weights that do not match the declared layout. The +standalone NPZ/header writers accept codebook matrices only; first call +`extract_codebooks(weights, num_levels=..., use_ema=..., kmeans_init=...)` for EMA. + +Saved-run repackaging uses the run's `config.json` and loads positional weight +arrays in numeric order. Preserve original packages and use `golden repackage + --output-dir ` when repairing an existing export. + +With sample inputs, `reference_vectors.npz` contains both deployed-path vectors +and independent `source_latents`, `source_indices`, `source_quantized_latents`, +and `source_reconstructions` from the restored training quantizer. Validation +requires identical indices and quantized latents for identical source latents; +the float encoder and decoder comparisons use `atol=rtol=1e-5`. This isolates codebook +correctness from encoder conversion and precision differences. Sample-data +reconstructions also use the discrete trained path rather than bypassing VQ. +Training references run on CPU with full float32 operations, so GPU TF32 rounding +cannot change the reference or fail the float32 conversion gate. +For INT8-only decoder packages, the training comparison uses the exported float +Keras decoder companion; the selected INT8 decoder is checked against its own +deployed reference vectors. Float tolerances are not applied to INT8 outputs. + +Strict release validation rejects older RVQ references without these source +arrays. Re-export them from the original training state; a self-consistent +deployed reference alone does not establish training parity. Repackaging also +recomputes model-card precision measurements instead of copying historical +reports that may have used different codebooks. + +Correcting the exported level count changes packet interpretation. Keep existing +packets paired with their original bundle and bind new packets to the corrected +bundle revision; replacing tables inside an existing package is not a migration. + ## Encoder Preprocessing and INT8 Parity Both RVQ encoders accept one normalized frame: PPG is `(1, 1, 320, 1)` at 64 Hz and ECG is `(1, 1, 512, 1)` at 256 Hz. Resample and frame the raw single-channel recording first, then normalize each frame independently: diff --git a/compressionkit/experiments/repackage.py b/compressionkit/experiments/repackage.py index a0eddfa..bcacea6 100644 --- a/compressionkit/experiments/repackage.py +++ b/compressionkit/experiments/repackage.py @@ -138,7 +138,9 @@ def repackage_rvq_golden( import keras + from compressionkit.export.codebook import extract_codebooks, load_rvq_weights from compressionkit.export.deploy import export_for_deployment + from compressionkit.export.quantization import refresh_rvq_encoder_precision_reports from compressionkit.export.stimulus import export_stimulus_npz, generate_normalized_stimulus output_dir = output_dir or (run_dir / "deploy") @@ -155,8 +157,18 @@ def repackage_rvq_golden( encoder = keras.models.load_model(encoder_path) decoder = keras.models.load_model(decoder_path) - rvq_npz = np.load(rvq_weights_path) - rvq_weights = [rvq_npz[k] for k in sorted(rvq_npz.files)] + rvq_weights = load_rvq_weights(rvq_weights_path) + # Use the saved run contract, not the current registry/config defaults. + saved_model = json.loads((run_dir / "config.json").read_text())["model"] + rvq_num_levels = saved_model["num_levels"] + rvq_use_ema = saved_model.get("use_ema", False) + rvq_kmeans_init = saved_model.get("kmeans_init", False) + codebooks = extract_codebooks( + rvq_weights, num_levels=rvq_num_levels, use_ema=rvq_use_ema, kmeans_init=rvq_kmeans_init + ) + sizes = saved_model.get("codebook_sizes") or [saved_model["latent_width"]] * rvq_num_levels + if [cb.shape for cb in codebooks] != [(size, saved_model["embedding_dim"]) for size in sizes]: + raise ValueError("Saved RVQ codebook shapes do not match the training config") input_shape = encoder.input_shape if len(input_shape) in (3, 4): @@ -181,20 +193,6 @@ def repackage_rvq_golden( "license": "other", "preprocessing_contract": preprocessing_contract, } - precision_paths = [ - run_dir.parent / "precision-study" / f"{experiment.experiment_id}.json", - run_dir.parent / "rvq_encoder_precision_comparison.json", - ] - for precision_path in precision_paths: - if not precision_path.is_file(): - continue - precision_payload = json.loads(precision_path.read_text()) - variants = precision_payload.get("experiments", {}).get(experiment.experiment_id, {}).get("variants", {}) - if isinstance(variants, dict): - model_card_info["encoder_precision_report"] = { - name: payload.get("report", {}) for name, payload in variants.items() if isinstance(payload, dict) - } - break if scorecard_payload is not None: model_card_info["scorecard_summary"] = scorecard_payload @@ -206,21 +204,19 @@ def repackage_rvq_golden( sample_rate=experiment.sample_rate, epsilon=preprocessing_contract["normalization"]["epsilon"], ) - latents = encoder.predict(sample_inputs, verbose=0) - sample_reconstructions = decoder.predict(latents, verbose=0) - if isinstance(sample_reconstructions, dict): - sample_reconstructions = sample_reconstructions.get("reconstruction", sample_reconstructions.get("output")) artifacts = export_for_deployment( encoder=encoder, decoder=decoder, rvq_weights=rvq_weights, + rvq_num_levels=rvq_num_levels, + rvq_use_ema=rvq_use_ema, + rvq_kmeans_init=rvq_kmeans_init, rep_dataset=rep_dataset, quantization_validation_dataset=quantization_validation_dataset, output_dir=output_dir, sample_inputs=sample_inputs, sample_targets=sample_inputs, - sample_reconstructions=np.asarray(sample_reconstructions, dtype=np.float32), model_name=experiment.run_name, export_decoder_float32=True, export_decoder_int8=export_decoder_int8, @@ -228,6 +224,8 @@ def repackage_rvq_golden( ) quantization_report = json.loads((output_dir / "quantization_report.json").read_text()) report_metrics = quantization_report["metrics"] + # Historical reports may have used different (including faulty) codebooks. + refresh_rvq_encoder_precision_reports(output_dir, quantization_validation_dataset) export_stimulus_npz( modality=experiment.modality, output_path=output_dir / "sample_stimulus.npz", diff --git a/compressionkit/export/codebook.py b/compressionkit/export/codebook.py index 382598f..49ea32f 100644 --- a/compressionkit/export/codebook.py +++ b/compressionkit/export/codebook.py @@ -15,25 +15,77 @@ logger = logging.getLogger(__name__) -def extract_codebooks(rvq_weights: list[np.ndarray]) -> list[np.ndarray]: +def extract_codebooks( + rvq_weights: list[np.ndarray], + *, + num_levels: int | None = None, + use_ema: bool = False, + kmeans_init: bool = False, +) -> list[np.ndarray]: """Extract per-level codebook matrices from RVQ weight list. - The weight list from ``rvq.get_weights()`` contains one embedding matrix - per RVQ level in order, each with shape ``(num_embeddings, embedding_dim)``. + Plain RVQ stores one matrix per level. EMA RVQ stores a codebook, + count vector, and embedding-sum matrix per level, followed by an optional + scalar k-means flag. The layout must be declared; matrix rank alone + cannot distinguish codebooks from EMA accumulators. Args: rvq_weights: List of weight arrays from ``rvq.get_weights()``. + num_levels: Expected number of levels. Required for EMA checkpoints. + use_ema: Whether the source is an EMA RVQ checkpoint. + kmeans_init: Whether the EMA checkpoint includes its warm-start flag. Returns: List of codebook arrays, one per level. + + Raises: + ValueError: The weights do not match the declared layout. """ - codebooks = [] - for w in rvq_weights: - if w.ndim == 2: - codebooks.append(w) + if use_ema and num_levels is None: + raise ValueError("EMA codebook extraction requires num_levels") + if kmeans_init and not use_ema: + raise ValueError("kmeans_init requires an EMA checkpoint") + levels = len(rvq_weights) if num_levels is None else num_levels + if levels < 1: + raise ValueError("At least one RVQ level is required") + stride = 3 if use_ema else 1 + expected = stride * levels + int(kmeans_init) + if len(rvq_weights) != expected: + raise ValueError(f"Expected {expected} RVQ weight arrays for {levels} levels, got {len(rvq_weights)}") + codebooks: list[np.ndarray] = [] + for level in range(levels): + cb = np.asarray(rvq_weights[stride * level]) + if cb.ndim != 2 or min(cb.shape) < 1 or not np.issubdtype(cb.dtype, np.floating): + raise ValueError(f"Level {level}: expected a nonempty floating-point codebook matrix") + if not np.isfinite(cb).all(): + raise ValueError(f"Level {level}: non-finite codebook values") + if codebooks and cb.shape[1] != codebooks[0].shape[1]: + raise ValueError("RVQ codebooks must share an embedding dimension") + if use_ema: + count, total = (np.asarray(w) for w in rvq_weights[stride * level + 1 : stride * level + 3]) + if count.shape != (cb.shape[0],) or total.shape != cb.shape: + raise ValueError(f"Level {level}: EMA count/sum shapes do not match the codebook") + codebooks.append(cb) + if kmeans_init: + flag = np.asarray(rvq_weights[-1]) + if flag.shape != () or flag.item() not in (0, 1): + raise ValueError("Expected a scalar 0/1 kmeans_done flag") return codebooks +def load_rvq_weights(path: str | Path) -> list[np.ndarray]: + """Read a positional RVQ training archive in numeric weight order. + + Reject missing or unexpected keys instead of silently reordering or + skipping state. Lexical sorting would place ``arr_10`` before ``arr_2``. + """ + with np.load(path, allow_pickle=False) as archive: + keys = [f"arr_{i}" for i in range(len(archive.files))] + if set(archive.files) != set(keys): + raise ValueError("RVQ training archive must contain contiguous arr_0 ... arr_N keys") + return [archive[key] for key in keys] + + def _format_c_array( name: str, data: np.ndarray, @@ -48,7 +100,10 @@ def _format_c_array( chunk_size = 8 for i in range(0, len(flat), chunk_size): chunk = flat[i : i + chunk_size] - vals = ", ".join(f"{v:.8g}f" for v in chunk) + # Nine significant digits round-trip float32. C requires a decimal + # point or exponent before the suffix (``0f`` is not a valid literal). + literals = [f"{v:.9g}" for v in chunk] + vals = ", ".join(f"{v if '.' in v or 'e' in v else v + '.0'}f" for v in literals) comma = "," if i + chunk_size < len(flat) else "" rows.append(f" {vals}{comma}") rows.append("};") @@ -62,7 +117,8 @@ def export_codebooks_npz( """Save RVQ codebook tables to a NumPy ``.npz`` archive. Args: - rvq_weights: Weight list from ``rvq.get_weights()``. + rvq_weights: Codebook matrices only; use ``extract_codebooks`` first + for an EMA checkpoint. output_path: Path for the ``.npz`` file. Returns: @@ -87,7 +143,8 @@ def export_codebooks_header( """Export RVQ codebook tables as a C header file. Args: - rvq_weights: Weight list from ``rvq.get_weights()``. + rvq_weights: Codebook matrices only; use ``extract_codebooks`` first + for an EMA checkpoint. output_path: Path for the ``.h`` file. prefix: Prefix for C array names (e.g. ``rvq_codebook_l0``). guard: Include guard macro name; auto-generated if ``None``. @@ -99,6 +156,9 @@ def export_codebooks_header( output_path.parent.mkdir(parents=True, exist_ok=True) codebooks = extract_codebooks(rvq_weights) + if any(cb.shape != codebooks[0].shape for cb in codebooks): + raise ValueError("The C header requires equally sized RVQ codebooks") + if guard is None: guard = output_path.stem.upper().replace("-", "_").replace(".", "_") + "_H" @@ -128,4 +188,5 @@ def export_codebooks_header( "export_codebooks_header", "export_codebooks_npz", "extract_codebooks", + "load_rvq_weights", ] diff --git a/compressionkit/export/deploy.py b/compressionkit/export/deploy.py index d8510de..b56dccf 100644 --- a/compressionkit/export/deploy.py +++ b/compressionkit/export/deploy.py @@ -159,6 +159,9 @@ def export_for_deployment( decoder: keras.Model, rvq_weights: list[np.ndarray], *, + rvq_num_levels: int | None = None, + rvq_use_ema: bool = False, + rvq_kmeans_init: bool = False, rep_dataset: np.ndarray, quantization_validation_dataset: np.ndarray | None = None, output_dir: str | Path, @@ -184,13 +187,17 @@ def export_for_deployment( encoder: Trained Keras encoder model. decoder: Trained Keras decoder model. rvq_weights: Weight list from ``rvq.get_weights()``. + rvq_num_levels: Expected trained RVQ level count; required for EMA. + rvq_use_ema: Whether weights use the EMA codebook/count/sum layout. + rvq_kmeans_init: Whether EMA weights include a k-means warm-start flag. rep_dataset: Representative input dataset for TFLite calibration. quantization_validation_dataset: Disjoint preprocessed frames for post-export INT8-vs-float32 parity validation. Required for INT8. output_dir: Directory to write all deployment artifacts. sample_inputs: Optional sample input frames for validation. sample_targets: Optional sample target frames for validation. - sample_reconstructions: Optional model reconstructions of targets. + sample_reconstructions: Optional reconstructions when inputs are absent. + With sample inputs, recomputed through the trained discrete RVQ path. quantization: Quantization mode for TFLite conversion. io_type: I/O type string for TFLite conversion. codebook_prefix: Prefix for codebook C array names. @@ -207,6 +214,29 @@ def export_for_deployment( ``DeploymentArtifacts`` with paths to all exported files. """ output_dir = Path(output_dir) + codebooks = extract_codebooks( + rvq_weights, num_levels=rvq_num_levels, use_ema=rvq_use_ema, kmeans_init=rvq_kmeans_init + ) + if any(cb.shape != codebooks[0].shape for cb in codebooks): + raise ValueError("Deployment currently requires equally sized RVQ codebooks") + if codebooks[0].shape[1] != encoder.output_shape[-1] or codebooks[0].shape[1] != decoder.input_shape[-1]: + raise ValueError("RVQ embedding dimension does not match encoder output and decoder input") + source_reference = None + if sample_inputs is not None: + from compressionkit.export.rvq_reference import build_rvq_reference + + source_reference = build_rvq_reference( + encoder, + decoder, + rvq_weights, + sample_inputs, + num_levels=len(codebooks), + use_ema=rvq_use_ema, + kmeans_init=rvq_kmeans_init, + ) + # Sample reconstructions must represent the discrete trained path, + # not decoder(encoder(x)) with the quantizer bypassed. + sample_reconstructions = source_reference["source_reconstructions"] if quantization == "INT8" and quantization_validation_dataset is None: raise ValueError("INT8 export requires a disjoint quantization_validation_dataset") output_dir.mkdir(parents=True, exist_ok=True) @@ -324,11 +354,11 @@ def export_for_deployment( # 3. Codebook tables logger.info("Exporting codebook tables...") artifacts.codebook_npz = export_codebooks_npz( - rvq_weights, + codebooks, output_dir / "codebook.npz", ) artifacts.codebook_header = export_codebooks_header( - rvq_weights, + codebooks, output_dir / "codebook.h", prefix=codebook_prefix, ) @@ -363,7 +393,6 @@ def export_for_deployment( artifacts.scorecard = write_scorecard_artifact(output_dir, scorecard_summary) # 5. Canonical runtime hydration spec - codebooks = extract_codebooks(rvq_weights) preprocessing_contract = dict(model_card_info.get("preprocessing_contract", {})) if model_card_info else {} codec_spec = { "family": "rvq", @@ -530,6 +559,7 @@ def export_for_deployment( "indices": encoded_indices, "reconstructions": decoded_frames, } + reference_payload.update(source_reference) if sample_targets is not None: reference_payload["targets"] = np.asarray(sample_targets, dtype=np.float32) np.savez_compressed(artifacts.reference_vectors, **reference_payload) @@ -540,6 +570,16 @@ def export_for_deployment( ) artifacts.reference_vectors = Path() + if artifacts.reference_vectors != Path(): + from compressionkit.export.validate import _validate_reference_vectors + + parity_errors: list[str] = [] + _validate_reference_vectors( + output_dir, "rvq", parity_errors, [], max_vectors=len(sample_inputs), require_training_reference=True + ) + if parity_errors: + raise ValueError("RVQ trained-model parity failed: " + "; ".join(parity_errors)) + # 7b. Compare the deployed INT8 encoder against its float32 companion on # a disjoint, real preprocessed validation partition. if quantization == "INT8": diff --git a/compressionkit/export/quantization.py b/compressionkit/export/quantization.py index 31dadf5..d7653e3 100644 --- a/compressionkit/export/quantization.py +++ b/compressionkit/export/quantization.py @@ -41,6 +41,26 @@ def passed(self) -> bool: return self.encoder_input_saturation_fraction_max <= MAX_INPUT_SATURATION_FRACTION +def refresh_rvq_encoder_precision_reports(deploy_dir: str | Path, frames: np.ndarray) -> None: + """Recompute the model card's precision results against this deploy package. + + Call after trained-quantizer parity succeeds and before final checksums. + Historical reports must not be carried over when codebooks change. + """ + root = Path(deploy_dir) + card_path = root / "model_card.json" + card = json.loads(card_path.read_text()) + card["encoder_precision_report"] = { + name: asdict(evaluate_rvq_encoder_quantization(root, frames, max_frames=None, candidate_encoder_name=filename)) + for name, filename in { + "int8": "encoder.tflite", + "fp16": "encoder_fp16.tflite", + "int16x8": "encoder_int16x8.tflite", + }.items() + } + write_json(card_path, card) + + def _prd_percent(actual: np.ndarray, expected: np.ndarray) -> float: """Return RMS percent difference, with a safe zero-reference fallback.""" err = np.asarray(actual, dtype=np.float64) - np.asarray(expected, dtype=np.float64) diff --git a/compressionkit/export/rvq_reference.py b/compressionkit/export/rvq_reference.py new file mode 100644 index 0000000..c13deaf --- /dev/null +++ b/compressionkit/export/rvq_reference.py @@ -0,0 +1,69 @@ +"""Independent references from the training RVQ layer, before deployment.""" + +from __future__ import annotations + +import keras +import numpy as np +import tensorflow as tf + +from compressionkit.export.keras_reference import decode_keras_reference +from compressionkit.layers.ema_residual_vector_quantizer import EmaResidualVectorQuantizer +from compressionkit.layers.residual_vector_quantizer import ResidualVectorQuantizer + + +def build_rvq_reference( + encoder: keras.Model, + decoder: keras.Model, + weights: list[np.ndarray], + inputs: np.ndarray, + *, + num_levels: int, + use_ema: bool = False, + kmeans_init: bool = False, +) -> dict[str, np.ndarray]: + """Restore the source quantizer and generate discrete-path reference arrays. + + This deliberately uses the training layer's weight restoration and + encode/decode methods, independently of export extraction and runtime VQ. + No training calls or EMA updates occur. Source-latent comparisons isolate + codebook export from changes caused by encoder conversion/quantization. + References run on CPU to avoid GPU TF32 rounding changing an FP32 gate. + + Args: + encoder: Float Keras encoder from the training run. + decoder: Float Keras decoder from the training run. + weights: Original, ordered ``vq.get_weights()`` state. + inputs: Nonempty batch of preprocessed input frames. + num_levels: Trained quantizer level count. + use_ema: Whether to restore the EMA quantizer. + kmeans_init: Whether the EMA checkpoint has a warm-start flag. + + Returns: + Source latents, indices, quantized latents, and decoded waveforms. + """ + with tf.device("/CPU:0"): + stride = 3 if use_ema else 1 + kwargs = { + "num_levels": num_levels, + "num_embeddings": [weights[stride * i].shape[0] for i in range(num_levels)], + "embedding_dim": weights[0].shape[1], + } + vq = ( + EmaResidualVectorQuantizer(**kwargs, kmeans_init=kmeans_init) + if use_ema + else ResidualVectorQuantizer(**kwargs) + ) + vq.build(encoder.output_shape) + vq.set_weights(weights) + latent = np.asarray(encoder(inputs, training=False), dtype=np.float32) + indices = vq.encode(latent) + quantized = keras.ops.convert_to_numpy(vq.decode(indices, latent.shape)) + decoded = decode_keras_reference(decoder, quantized) + return { + "source_latents": latent, + "source_indices": np.stack([keras.ops.convert_to_numpy(i) for i in indices], axis=-1).reshape( + *latent.shape[:-1], num_levels + ), + "source_quantized_latents": np.asarray(quantized, dtype=np.float32), + "source_reconstructions": np.asarray(decoded, dtype=np.float32), + } diff --git a/compressionkit/export/validate.py b/compressionkit/export/validate.py index cac593f..88e335d 100644 --- a/compressionkit/export/validate.py +++ b/compressionkit/export/validate.py @@ -93,6 +93,7 @@ def _validate_reference_vectors( warnings: list[str], *, max_vectors: int, + require_training_reference: bool = False, ) -> None: ref_path = deploy_dir / ArtifactFile.REFERENCE_VECTORS if not ref_path.exists(): @@ -147,9 +148,42 @@ def _validate_reference_vectors( return if family == "rvq": + source_keys = {"source_latents", "source_indices", "source_quantized_latents", "source_reconstructions"} + has_source = source_keys.issubset(blob.files) + if not has_source: + message = "RVQ reference vectors lack independent trained-quantizer references; re-export required" + (errors if require_training_reference else warnings).append(message) frames = np.asarray(blob[SampleArray.INPUT_FRAMES], dtype=np.float32) indices = np.asarray(blob[SampleArray.INDICES], dtype=np.int32) recon = np.asarray(blob[SampleArray.RECONSTRUCTIONS], dtype=np.float32) + if frames.shape[0] == 0: + errors.append("RVQ reference vectors contain no frames") + return + if has_source: + latent_shape = blob["source_latents"].shape + if ( + latent_shape[0] != frames.shape[0] + or blob["source_quantized_latents"].shape != latent_shape + or blob["source_indices"].shape != (*latent_shape[:-1], codec.num_levels) + or blob["source_reconstructions"].shape != recon.shape + ): + errors.append("RVQ trained-quantizer reference shapes do not match the deployment contract") + return + float_decoder = None + if has_source: + decoder_spec = (codec.spec or codec.manifest).get("decoder", {}) + float_tflite = decoder_spec.get("float32_tflite") + if not float_tflite or not (deploy_dir / float_tflite).is_file(): + # The runtime may select an INT8 decoder. Compare training + # references to the exported float Keras companion instead; + # quantized output has separate deployed reference vectors. + from compressionkit.export.keras_reference import decode_keras_reference, load_keras_reference_decoder + + keras_decoder = decoder_spec.get("keras") + if not keras_decoder: + errors.append("RVQ training parity requires a float decoder companion") + return + float_decoder = load_keras_reference_decoder(deploy_dir / keras_decoder) for idx in range(min(max_vectors, frames.shape[0])): sample = frames[idx : idx + 1] actual_indices = codec.encode(sample) @@ -160,6 +194,27 @@ def _validate_reference_vectors( actual_recon = codec.decode(actual_indices) if not np.allclose(actual_recon, recon[idx : idx + 1], atol=1e-5): errors.append(f"RVQ reference reconstruction mismatch at sample {idx}") + if has_source: + source_latent = blob["source_latents"][idx : idx + 1] + deployed_latent = codec.encode_latent(sample) + if deployed_latent.shape != source_latent.shape or not np.allclose( + deployed_latent, source_latent, atol=1e-5, rtol=1e-5 + ): + errors.append(f"RVQ trained-encoder latent mismatch at sample {idx} (atol=rtol=1e-5)") + source_indices = blob["source_indices"][idx : idx + 1] + actual_source_indices = codec.quantize_latent(source_latent) + if not np.array_equal(actual_source_indices, source_indices): + errors.append(f"RVQ trained-quantizer index mismatch at sample {idx}") + continue + actual_latent = codec.dequantize_indices(source_indices) + if not np.array_equal(actual_latent, blob["source_quantized_latents"][idx : idx + 1]): + errors.append(f"RVQ trained-quantizer latent mismatch at sample {idx}") + if float_decoder is None: + source_recon = codec.decode_latent(actual_latent) + else: + source_recon = decode_keras_reference(float_decoder, actual_latent) + if not np.allclose(source_recon, blob["source_reconstructions"][idx : idx + 1], atol=1e-5, rtol=1e-5): + errors.append(f"RVQ trained-decoder reconstruction mismatch at sample {idx} (atol=rtol=1e-5)") return warnings.append(f"No reference-vector validator registered for family {family!r}") @@ -297,7 +352,9 @@ def validate_deploy_package( if check_reference_vectors and not errors: try: - _validate_reference_vectors(root, family, errors, warnings, max_vectors=max_vectors) + _validate_reference_vectors( + root, family, errors, warnings, max_vectors=max_vectors, require_training_reference=strict_release + ) except Exception as exc: errors.append(f"reference-vector validation failed: {exc}") diff --git a/compressionkit/recipes/base_rvq.py b/compressionkit/recipes/base_rvq.py index 8951e74..6bddb79 100644 --- a/compressionkit/recipes/base_rvq.py +++ b/compressionkit/recipes/base_rvq.py @@ -281,6 +281,9 @@ def train(self) -> dict[str, Any]: model.encoder, model.decoder, model.vq.get_weights(), + rvq_num_levels=self.cfg.model.num_levels, + rvq_use_ema=self.cfg.model.use_ema, + rvq_kmeans_init=self.cfg.model.kmeans_init, rep_dataset=rep_dataset, quantization_validation_dataset=quantization_validation_dataset, output_dir=run_dir / "deploy", diff --git a/compressionkit/runtime/codec.py b/compressionkit/runtime/codec.py index 541ce46..646a263 100644 --- a/compressionkit/runtime/codec.py +++ b/compressionkit/runtime/codec.py @@ -132,8 +132,15 @@ def __init__(self, deploy_dir: str | Path) -> None: # Load codebook cb_path = self._deploy_dir / runtime_spec["codebook"]["npz"] - cb_data = np.load(cb_path) - self._codebooks = [cb_data[k] for k in sorted(cb_data.files)] + cb_info = runtime_spec["codebook"] + with np.load(cb_path, allow_pickle=False) as cb_data: + keys = [f"level_{i}" for i in range(cb_info["num_levels"])] + if set(cb_data.files) != set(keys): + raise ValueError("Codebook archive levels do not match the deployment contract") + self._codebooks = [cb_data[k] for k in keys] + expected_shape = (cb_info["num_embeddings"], cb_info["embedding_dim"]) + if not self._codebooks or any(cb.shape != expected_shape for cb in self._codebooks): + raise ValueError("Codebook shapes do not match the deployment contract") self._num_levels = len(self._codebooks) self._num_embeddings = self._codebooks[0].shape[0] self._embedding_dim = self._codebooks[0].shape[1] diff --git a/maintainer-notes/rvq-export-audit.md b/maintainer-notes/rvq-export-audit.md new file mode 100644 index 0000000..01ed000 --- /dev/null +++ b/maintainer-notes/rvq-export-audit.md @@ -0,0 +1,47 @@ +# RVQ export repair audit — issue #74 + +2026-10-05. Source checkout: `7fb3663`; repair branch: `issue-74-rvq-codebook-export`. + +## Affected artifacts + +All 11 distinct registered single-stage RVQ runs contain exported EMA accumulators. Four two-stage registry entries reuse those run directories; their entropy priors were not requalified. The codebook payloads at all 11 published revisions pinned in compressionkit-dev's `experiments/mvp_v02/inventory.json` match the faulty local codebooks exactly. This establishes codebook identity, not identity of every file in those bundles. + +| Published model | Pinned revision | +|---|---| +| `Ambiq/compressionkit-ecg-2x-v1.0` | `4d07900f24978ee3830ee2a6a4e2becaa0880645` | +| `Ambiq/compressionkit-ecg-4x-v1.0` | `f8be1ba5454d0eb5c645393ec3261e74d0bf26a8` | +| `Ambiq/compressionkit-ecg-8x-v1.0` | `548a82529fc10075c31f19ab1ed8c09dfb2e58f4` | +| `Ambiq/compressionkit-ecg-16x-v1.0` | `e1c78d7227fec77887bd6a619d3463e238c3e01d` | +| `Ambiq/compressionkit-ecg-32x-v1.0` | `13c9be7e94b80191385f6be3aad32d2effb59e0b` | +| `Ambiq/compressionkit-ecg-64x-v1.0` | `49c2dd0ddb04b23d6817e584c9865046ecd15424` | +| `Ambiq/compressionkit-ppg-2x-v1.0` | `af00d62fcdd8e8c12a2929bae9d22911ab518966` | +| `Ambiq/compressionkit-ppg-4x-v1.0` | `8e3b465983808d50d99799df6135aaffe68029ff` | +| `Ambiq/compressionkit-ppg-8x-v1.0` | `7e2b0cea808d1b1f168c4131ea5c00ffe50e9cce` | +| `Ambiq/compressionkit-ppg-16x-v1.0` | `928d302ac822712c4513488f6d4e0fef962bde37` | +| `Ambiq/compressionkit-ppg-32x-v1.0` | `bff7e89fbe3c6b35a2db829d7da6022eb26fcf34` | + +## Before and after + +Ten saved reference frames per modality were compared with the saved Keras encoder/decoder and restored training quantizer. These measurements assess deployment fidelity, not physiological accuracy or clean-reference denoising quality. + +| Measurement | ECG 8x | PPG 8x | +|---|---:|---:| +| Exported levels | 4 → 2 | 4 → 2 | +| Codebook value bytes | 65,536 → 32,768 | 65,536 → 32,768 | +| Complete packet bytes per frame | 260 → 132 | 164 → 84 | +| Median reconstruction PRD against trained discrete path (%) | 6.36077 → 0.00002241 | 1.62395 → 0.00001964 | + +Both corrected packages have exact token and quantized-latent parity on identical source latents. Keras versus FP32 LiteRT encoder tokens agree on every audited frame. Maximum encoder/decoder differences are below 2.4e-6 in normalized units. + +Packet checks reused `RVQFrames` from the sibling checkout at `80d774d`, supplied with the compressionkit runtime. All 40 old/new packet round trips passed. Counts include the 4-byte float16 mean/scale header and fixed-width tokens. These checks used FP32 encoders and do not establish bit-exactness of MCU INT8 kernels. + +## Validation and provenance + +- Reproduced the original extractor returning four tables from an actual two-level EMA layer. +- Regression coverage includes EMA/non-EMA, warm-start state, numeric checkpoint and level ordering, malformed layouts, independent training references, and C99 compilation with float32 table parity. +- After review fixes, 50 focused tests pass under normal GPU devcontainer settings. Float references use a local CPU context to avoid TF32 rounding; INT8-only decoders retain both float-companion training checks and quantized deployed-reference checks. A fresh independent review found no remaining actionable findings and verified 19 focused tests. +- All 11 regenerated single-stage candidates pass strict deploy validation on ten independent training-reference frames each, with no errors or warnings. +- Recomputed INT8, FP16, and INT16x8 model-card reports on real disjoint validation partitions of 2,048 frames per run. Historical physiological scorecards remain historical results; no physiological evaluation or retraining was performed. +- Original run/deploy artifacts remain unchanged. Corrected candidates and evidence are in the repair worktree's git-ignored `results/issue-74/`: `audit.json`, `source-checksums.json`, `strict-validation.json`, `environment.txt`, test output, and investigation scripts. + +The repair requires declared EMA layout and independent trained-quantizer references for strict RVQ validation. Existing packets must retain their original bundle identity because corrected level counts change packet interpretation. Corrected model artifacts remain unpublished; this audit does not authorize a release or merge. diff --git a/tests/test_base_rvq_scorecard_sync.py b/tests/test_base_rvq_scorecard_sync.py index 18c11bb..b5a14d0 100644 --- a/tests/test_base_rvq_scorecard_sync.py +++ b/tests/test_base_rvq_scorecard_sync.py @@ -70,7 +70,7 @@ def test_base_rvq_trainer_syncs_built_scorecard_into_deploy(monkeypatch, tmp_pat run_name="ppg_rvq_64hz_04x_golden", model_dump=lambda: {"run_name": "ppg_rvq_64hz_04x_golden"}, data=SimpleNamespace(sampling_rate=64, steps_per_epoch=1, epochs=1), - model=SimpleNamespace(kmeans_init=False), + model=SimpleNamespace(kmeans_init=False, use_ema=False, num_levels=1), evaluation=SimpleNamespace( int8_calibration_frames=1, int8_validation_frames=1, diff --git a/tests/test_codebook_export.py b/tests/test_codebook_export.py new file mode 100644 index 0000000..8fe54a9 --- /dev/null +++ b/tests/test_codebook_export.py @@ -0,0 +1,140 @@ +"""Codebook exports must preserve the trained quantizer, including EMA state.""" + +from __future__ import annotations + +import re + +import keras +import numpy as np +import pytest + +from compressionkit.export.codebook import ( + export_codebooks_header, + export_codebooks_npz, + extract_codebooks, + load_rvq_weights, +) +from compressionkit.layers.ema_residual_vector_quantizer import EmaResidualVectorQuantizer +from compressionkit.layers.residual_vector_quantizer import ResidualVectorQuantizer +from compressionkit.runtime.codec import RVQCodec + + +@pytest.mark.parametrize("use_ema,kmeans_init", [(False, False), (True, False), (True, True)]) +def test_export_preserves_trained_quantizer(tmp_path, use_ema, kmeans_init): + rng = np.random.default_rng(74) + kwargs = {"num_levels": 4, "num_embeddings": 8, "embedding_dim": 3} + layer = ( + EmaResidualVectorQuantizer(**kwargs, kmeans_init=kmeans_init) if use_ema else ResidualVectorQuantizer(**kwargs) + ) + layer.build((None, 1, 16, 3)) + for level, cb in enumerate(layer._codebooks): + cb.assign(rng.normal(size=(8, 3)).astype(np.float32) / (level + 1)) + if use_ema: + for count, total in zip(layer._ema_counts, layer._ema_weights): + count.assign(np.full(8, 13.0, dtype=np.float32)) + total.assign(rng.normal(100, 5, size=(8, 3)).astype(np.float32)) + weights = layer.get_weights() + archive = tmp_path / "rvq_weights.npz" + np.savez(archive, *weights) # Includes arr_10 and arr_11 for EMA. + codebooks = extract_codebooks(load_rvq_weights(archive), num_levels=4, use_ema=use_ema, kmeans_init=kmeans_init) + npz = export_codebooks_npz(codebooks, tmp_path / "codebook.npz") + header = export_codebooks_header(codebooks, tmp_path / "codebook.h") + with np.load(npz) as saved: + assert saved.files == [f"level_{i}" for i in range(4)] + for i, cb in enumerate(layer._codebooks): + np.testing.assert_array_equal(saved[f"level_{i}"], keras.ops.convert_to_numpy(cb)) + + # C literals must preserve the same float32 values as the NPZ tables. + literals = re.findall(r"(?\n#include "codebook.h"\nint main(void) {\n' + "return fwrite(rvq_codebook_l0, sizeof(float), 6, stdout) == 6 ? 0 : 1;\n}\n" + ) + executable = tmp_path / "check" + subprocess.run( + [cc, "-std=c99", "-Wall", "-Werror", str(source), "-o", str(executable)], check=True, capture_output=True + ) + result = subprocess.run([str(executable)], check=True, capture_output=True) + np.testing.assert_array_equal(np.frombuffer(result.stdout, dtype=np.float32), values.ravel()) diff --git a/tests/test_rvq_deploy_smoke.py b/tests/test_rvq_deploy_smoke.py index aeb7a1f..41b480d 100644 --- a/tests/test_rvq_deploy_smoke.py +++ b/tests/test_rvq_deploy_smoke.py @@ -6,6 +6,7 @@ import keras import numpy as np +import pytest from compressionkit.experiments import cli as golden_cli from compressionkit.export import export_for_deployment, validate_deploy_package @@ -39,7 +40,9 @@ def _tiny_decoder() -> keras.Model: return keras.Model(inputs, outputs, name="smoke_decoder") -def _build_smoke_rvq_deploy(tmp_path, *, quantization="FP32", export_decoder_float32=True, export_decoder_int8=True): +def _build_smoke_rvq_deploy( + tmp_path, *, use_ema=False, export_decoder_float32=True, export_decoder_int8=True, quantization="FP32" +): encoder = _tiny_encoder() decoder = _tiny_decoder() @@ -48,12 +51,24 @@ def _build_smoke_rvq_deploy(tmp_path, *, quantization="FP32", export_decoder_flo sample_targets = sample_inputs.copy() sample_reconstructions = decoder.predict(encoder.predict(sample_inputs, verbose=0), verbose=0) rvq_weights = [np.random.randn(NUM_EMBEDDINGS, EMBED_DIM).astype(np.float32)] + if use_ema: + from compressionkit.layers.ema_residual_vector_quantizer import EmaResidualVectorQuantizer + + vq = EmaResidualVectorQuantizer(num_levels=2, num_embeddings=NUM_EMBEDDINGS, embedding_dim=EMBED_DIM) + vq.build(encoder.output_shape) + for cb, count, total in zip(vq._codebooks, vq._ema_counts, vq._ema_weights): + cb.assign(np.random.randn(NUM_EMBEDDINGS, EMBED_DIM).astype(np.float32)) + count.assign(np.full(NUM_EMBEDDINGS, 10.0, np.float32)) + total.assign(np.full((NUM_EMBEDDINGS, EMBED_DIM), 50.0, np.float32)) + rvq_weights = vq.get_weights() deploy_dir = tmp_path / "deploy" artifacts = export_for_deployment( encoder, decoder, rvq_weights, + rvq_num_levels=2 if use_ema else 1, + rvq_use_ema=use_ema, rep_dataset=rep_dataset, output_dir=deploy_dir, sample_inputs=sample_inputs, @@ -82,8 +97,9 @@ def _build_smoke_rvq_deploy(tmp_path, *, quantization="FP32", export_decoder_flo return deploy_dir, artifacts -def test_export_for_deployment_smoke_validates_strict_release(tmp_path) -> None: - deploy_dir, artifacts = _build_smoke_rvq_deploy(tmp_path) +@pytest.mark.parametrize("use_ema", [False, True]) +def test_export_for_deployment_smoke_validates_strict_release(tmp_path, use_ema) -> None: + deploy_dir, artifacts = _build_smoke_rvq_deploy(tmp_path, use_ema=use_ema) assert artifacts.manifest.exists() assert artifacts.codec_spec.exists() @@ -108,6 +124,11 @@ def test_export_for_deployment_smoke_validates_strict_release(tmp_path) -> None: assert "reference_vectors.npz" in result.checked_files manifest = json.loads(artifacts.manifest.read_text()) + assert manifest["codebook"]["num_levels"] == (2 if use_ema else 1) + with np.load(artifacts.reference_vectors) as refs: + assert refs["source_indices"].shape[-1] == (2 if use_ema else 1) + with np.load(artifacts.sample_data_npz) as samples: + np.testing.assert_array_equal(samples["reconstructions"], refs["source_reconstructions"]) assert manifest["encoder"]["default_variant"] == "float32" assert manifest["encoder"]["default_tflite"] == "encoder_float32.tflite" assert manifest["encoder"]["float32_tflite"] == "encoder_float32.tflite" @@ -127,6 +148,33 @@ def test_export_for_deployment_smoke_validates_strict_release(tmp_path) -> None: interpreter.invoke() assert interpreter.get_tensor(output_detail["index"]).dtype == np.float32 + if use_ema: + from compressionkit.export.release import write_checksums + + with np.load(artifacts.reference_vectors) as refs: + payload = {key: refs[key] for key in refs.files} + source_latents = payload["source_latents"].copy() + payload["source_latents"] = source_latents + 0.01 + np.savez_compressed(artifacts.reference_vectors, **payload) + write_checksums(deploy_dir) + result = validate_deploy_package(deploy_dir, strict_release=True, max_vectors=1) + assert any("trained-encoder latent mismatch" in error for error in result.errors) + payload["source_latents"] = source_latents + # The deployed references still agree with themselves. An independent + # training token mismatch must nevertheless reject the package. + payload["source_indices"] = (payload["source_indices"] + 1) % NUM_EMBEDDINGS + np.savez_compressed(artifacts.reference_vectors, **payload) + write_checksums(deploy_dir) + result = validate_deploy_package(deploy_dir, strict_release=True, max_vectors=1) + assert any("trained-quantizer index mismatch" in error for error in result.errors) + for key in list(payload): + if key.startswith("source_"): + payload.pop(key) + np.savez_compressed(artifacts.reference_vectors, **payload) + write_checksums(deploy_dir) + result = validate_deploy_package(deploy_dir, strict_release=True, max_vectors=1) + assert any("lack independent trained-quantizer references" in error for error in result.errors) + def test_validate_deploy_cli_smoke(tmp_path, capsys) -> None: deploy_dir, _artifacts = _build_smoke_rvq_deploy(tmp_path) @@ -140,6 +188,35 @@ def test_validate_deploy_cli_smoke(tmp_path, capsys) -> None: assert "scorecard.json" in out +def test_int8_only_decoder_preserves_independent_float_parity(tmp_path) -> None: + from compressionkit.export.release import write_checksums + + deploy_dir, artifacts = _build_smoke_rvq_deploy(tmp_path, export_decoder_float32=False, quantization="INT8") + assert not (deploy_dir / "decoder_float32.tflite").exists() + assert artifacts.decoder_int8_tflite.exists() + result = validate_deploy_package(deploy_dir, strict_release=True, max_vectors=2) + assert result.ok, result.errors + assert result.warnings == [] + + # Float training references and INT8 deployed references both remain + # checked, rather than weakening the float tolerance for quantization. + with np.load(artifacts.reference_vectors) as refs: + payload = {key: refs[key] for key in refs.files} + original = payload["source_reconstructions"].copy() + payload["source_reconstructions"] += 0.1 + np.savez_compressed(artifacts.reference_vectors, **payload) + write_checksums(deploy_dir) + result = validate_deploy_package(deploy_dir, strict_release=True, max_vectors=2) + assert any("trained-decoder reconstruction mismatch" in error for error in result.errors) + + payload["source_reconstructions"] = original + payload["reconstructions"] += 0.1 + np.savez_compressed(artifacts.reference_vectors, **payload) + write_checksums(deploy_dir) + result = validate_deploy_package(deploy_dir, strict_release=True, max_vectors=2) + assert any("RVQ reference reconstruction mismatch" in error for error in result.errors) + + def test_int8_encoder_export_supports_keras_only_decoder(tmp_path) -> None: from pathlib import Path diff --git a/tests/test_rvq_repackage_weights.py b/tests/test_rvq_repackage_weights.py new file mode 100644 index 0000000..d421ca5 --- /dev/null +++ b/tests/test_rvq_repackage_weights.py @@ -0,0 +1,83 @@ +"""Repackaging must honor the saved training layout rather than registry defaults.""" + +from __future__ import annotations + +import json +from types import SimpleNamespace + +import keras +import numpy as np +import pytest + +from compressionkit.experiments import repackage +from compressionkit.export.codebook import extract_codebooks +from compressionkit.layers.ema_residual_vector_quantizer import EmaResidualVectorQuantizer + + +@pytest.mark.parametrize("bad_config", [False, True]) +def test_repackage_restores_saved_ema_weights(tmp_path, monkeypatch, bad_config): + vq = EmaResidualVectorQuantizer(num_levels=4, num_embeddings=8, embedding_dim=3, kmeans_init=True) + vq.build((None, 1, 4, 3)) + for i, cb in enumerate(vq._codebooks): + cb.assign(np.full((8, 3), (i + 1) / 10, np.float32)) + np.savez(tmp_path / "rvq_weights.npz", *vq.get_weights()) + config = { + "model": { + "num_levels": 2 if bad_config else 4, + "latent_width": 8, + "embedding_dim": 3, + "use_ema": True, + "kmeans_init": True, + } + } + (tmp_path / "config.json").write_text(json.dumps(config)) + for name in ("encoder.keras", "decoder.keras"): + (tmp_path / name).touch() + monkeypatch.setattr(keras.models, "load_model", lambda _: SimpleNamespace(input_shape=(None, 1, 16, 1))) + frames = np.ones((2, 1, 16, 1), np.float32) + monkeypatch.setattr( + repackage, + "collect_release_quantization_frames", + lambda *args: (frames, frames + 1, {"normalization": {"epsilon": 0.001}}), + ) + monkeypatch.setattr(repackage, "finalize_release_metadata", lambda *args: None) + from compressionkit.export import deploy, quantization, stimulus + + monkeypatch.setattr(stimulus, "export_stimulus_npz", lambda **kwargs: None) + monkeypatch.setattr(stimulus, "generate_normalized_stimulus", lambda **kwargs: frames.copy()) + monkeypatch.setattr(quantization, "refresh_rvq_encoder_precision_reports", lambda *args: None) + calls = [] + + def export(**kwargs): + calls.append(kwargs) + books = extract_codebooks( + kwargs["rvq_weights"], + num_levels=kwargs["rvq_num_levels"], + use_ema=kwargs["rvq_use_ema"], + kmeans_init=kwargs["rvq_kmeans_init"], + ) + for cb, expected in zip(books, vq._codebooks): + np.testing.assert_array_equal(cb, keras.ops.convert_to_numpy(expected)) + output = kwargs["output_dir"] + output.mkdir() + (output / "quantization_report.json").write_text( + json.dumps( + { + "passed": True, + "metrics": {"reconstruction_prd_percent_p90": 0.0, "encoder_input_saturation_fraction_max": 0.0}, + } + ) + ) + return SimpleNamespace(as_dict=lambda: {}) + + monkeypatch.setattr(deploy, "export_for_deployment", export) + experiment = SimpleNamespace( + method="rvq", experiment_id="test", run_name="test", modality="ecg", sample_rate=256, compression_ratio=8 + ) + if bad_config: + with pytest.raises(ValueError, match=r"Expected .* RVQ weight arrays"): + repackage.repackage_rvq_golden(experiment, run_dir=tmp_path) + assert calls == [] + else: + repackage.repackage_rvq_golden(experiment, run_dir=tmp_path) + assert len(calls) == 1