| """Check native layer extraction and the indexed position lookup on real media.""" |
|
|
| import argparse |
| import json |
| from pathlib import Path |
|
|
| from mlx_vlm.models.gemma4.vision import VisionPatchEmbedder |
|
|
| from gemma_rlcd import State |
| from gemma_rlcd.core import ScoringRequest |
| from gemma_rlcd.head_backend import DecisionHeadBackend |
|
|
|
|
| def main(): |
| parser = argparse.ArgumentParser() |
| parser.add_argument("--model", required=True) |
| parser.add_argument("--media", required=True, type=Path) |
| parser.add_argument("--report", required=True, type=Path) |
| parser.add_argument("--state-layers", type=int, default=35) |
| parser.add_argument("--dtype", choices=["float16", "float32", "bfloat16"], default="float32") |
| args = parser.parse_args() |
| backend = DecisionHeadBackend( |
| args.model, state_layers=args.state_layers, compute_dtype=args.dtype |
| ) |
| mx = backend.mx |
| media = args.media.resolve() |
| records = [] |
| for name, state in [ |
| ("text", State(text="A cat sits on a blue chair.")), |
| ("image", State(images=(str(media / "red.png"),))), |
| ("speech", State(audio=(str(media / "dog.wav"),))), |
| ("video_speech", State(videos=(str(media / "red-blue-speech.mp4"),))), |
| ]: |
| prepared = backend.prepare(state, [ScoringRequest("Read the state.", ("A", "B"))]) |
| actual = backend._state_features(prepared.inputs) |
| features = backend.model.get_input_embeddings(**prepared.inputs) |
| captured = [] |
| final = backend.model.language_model.model( |
| inputs=prepared.inputs["input_ids"], |
| inputs_embeds=features.inputs_embeds, |
| per_layer_inputs=features.per_layer_inputs, |
| mm_token_type_ids=prepared.inputs.get("mm_token_type_ids"), |
| hidden_sink=captured, |
| capture_layer_ids=[backend.state_layers - 1], |
| logits_to_keep=1, |
| ) |
| mx.eval(actual, captured, final) |
| record = { |
| "state": name, |
| "native_layer_max_abs_delta": float(mx.max(mx.abs(actual - captured[0])).item()), |
| } |
| indexed = backend.model.vision_tower.patch_embedder |
| dense = VisionPatchEmbedder(backend.model.config.vision_config) |
| dense.input_proj = indexed.input_proj |
| dense.position_embedding_table = indexed.position_embedding_table |
| backend.model.vision_tower.patch_embedder = dense |
| reference = backend.model.get_input_embeddings(**prepared.inputs).inputs_embeds |
| mx.eval(reference) |
| backend.model.vision_tower.patch_embedder = indexed |
| record["dense_vs_indexed_embedding_max_abs_delta"] = float( |
| mx.max(mx.abs(reference - features.inputs_embeds)).item() |
| ) |
| records.append(record) |
| print(json.dumps(record), flush=True) |
| args.report.write_text( |
| json.dumps( |
| {"compute_dtype": args.dtype, "state_layers": args.state_layers, "results": records}, |
| indent=2, |
| ) |
| + "\n" |
| ) |
| if any( |
| row["native_layer_max_abs_delta"] != 0 |
| or (args.dtype != "float32" and row["dense_vs_indexed_embedding_max_abs_delta"] != 0) |
| for row in records |
| ): |
| raise SystemExit("Encoder equivalence check failed; inspect recorded differences") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
