Buckets:
| import json | |
| import tempfile | |
| import unittest | |
| from pathlib import Path | |
| import torch | |
| from zebra_sft.analyze_ablations import analyze_suite | |
| def _row(sample_id, *, condition, correct, t2, masked=False): | |
| text1 = {"type": "text", "event": "image_gen", "token_ids": [10, 11]} | |
| text2 = {"type": "text", "event": "im_end", "token_ids": [t2]} | |
| image = {"type": "image", "index": 1, "grid": [2, 2]} | |
| step = {"step": 1, "generated_grid": [2, 2]} | |
| row = { | |
| "sample_id": sample_id, | |
| "config": f"config-{sample_id}", | |
| "context_limit": 32768, | |
| "stop_reason": "im_end", | |
| "generated_images": 1, | |
| "generated_grids": [[2, 2]], | |
| "sequence_ids": [10, 11, 20, 21, 99, 99, 99, 99, 22, t2], | |
| "events": [ | |
| text1, | |
| { | |
| "type": "scaffold", | |
| "event": "vision_start", | |
| "token_ids": [20, 21], | |
| "grid": [2, 2], | |
| }, | |
| image, | |
| text2, | |
| ], | |
| "visual_steps": [step], | |
| "score": {"correct": correct, "prediction": str(t2)}, | |
| "relaxed_complete": True, | |
| "has_final_answer": True, | |
| "balanced_think_tags": True, | |
| } | |
| if condition is not None: | |
| row["ablation"] = condition | |
| image.update( | |
| { | |
| "ablation": condition, | |
| "serialized_positions": 0 if condition == "delete_visual" else 4, | |
| "attention_value": 1, | |
| } | |
| ) | |
| step.clear() | |
| step.update( | |
| { | |
| "step": 1, | |
| "declared_grid": [2, 2], | |
| "serialized_positions": image["serialized_positions"], | |
| } | |
| ) | |
| text1["masked_key_spans"] = [] | |
| text2["masked_key_spans"] = [[0, 1]] if masked else [] | |
| row["masked_thought_spans"] = [[0, 1]] if masked else [] | |
| row["sequence_attention_mask"] = [1] * len(row["sequence_ids"]) | |
| return row | |
| def _write_sample(root, row, codes): | |
| sample = root / "samples" / row["sample_id"] | |
| sample.mkdir(parents=True) | |
| (sample / "result.json").write_text(json.dumps(row), encoding="utf-8") | |
| torch.save(codes, sample / "generated_codes.pt") | |
| class AblationAnalysisTest(unittest.TestCase): | |
| def test_reports_paired_outcomes_t2_and_expected_intervention(self): | |
| with tempfile.TemporaryDirectory() as directory: | |
| root = Path(directory) | |
| control = root / "control" | |
| suite = root / "suite" | |
| for sample_id, correct in (("a", False), ("b", True)): | |
| codes = [torch.full((1, 2), ord(sample_id), dtype=torch.uint8)] | |
| _write_sample( | |
| control, | |
| _row(sample_id, condition=None, correct=correct, t2=30), | |
| codes, | |
| ) | |
| _write_sample( | |
| suite / "noise", | |
| _row( | |
| sample_id, | |
| condition="noise", | |
| correct=not correct, | |
| t2=31 if sample_id == "a" else 30, | |
| ), | |
| [codes[0] + 1], | |
| ) | |
| result = analyze_suite( | |
| suite, control, conditions=["noise"], expected_samples=2 | |
| )["ablations"]["noise"] | |
| self.assertEqual( | |
| result["paired_control"]["outcomes"], | |
| {"W→W": 0, "W→R": 1, "R→W": 1, "R→R": 0}, | |
| ) | |
| self.assertEqual(result["paired_control"]["t2"]["exact_token_ids"], 1) | |
| self.assertEqual( | |
| result["paired_control"]["first_state_or_output_divergence"], | |
| {"I1": 2}, | |
| ) | |
| self.assertTrue(result["expected_intervention"]["result"]["all_passed"]) | |
| self.assertTrue( | |
| all( | |
| check["all_passed"] | |
| for check in result["protocol_invariants"].values() | |
| ) | |
| ) | |
| def test_mask_audits_first_bsq_and_span_history(self): | |
| with tempfile.TemporaryDirectory() as directory: | |
| root = Path(directory) | |
| control = root / "control" | |
| suite = root / "suite" | |
| codes = [torch.arange(2, dtype=torch.uint8).reshape(1, 2)] | |
| _write_sample( | |
| control, | |
| _row("a", condition=None, correct=False, t2=30), | |
| codes, | |
| ) | |
| _write_sample( | |
| suite / "mask_prior_thought", | |
| _row( | |
| "a", | |
| condition="mask_prior_thought", | |
| correct=False, | |
| t2=31, | |
| masked=True, | |
| ), | |
| [codes[0].clone()], | |
| ) | |
| result = analyze_suite( | |
| suite, | |
| control, | |
| conditions=["mask_prior_thought"], | |
| expected_samples=1, | |
| )["ablations"]["mask_prior_thought"] | |
| self.assertTrue(result["expected_intervention"]["result"]["all_passed"]) | |
| self.assertEqual( | |
| result["paired_control"]["first_state_or_output_divergence"], | |
| {"T2": 1}, | |
| ) | |
| self.assertTrue( | |
| all(check["all_passed"] for check in result["mask_span_audit"].values()) | |
| ) | |
| if __name__ == "__main__": | |
| unittest.main() | |
Xet Storage Details
- Size:
- 5.42 kB
- Xet hash:
- c369a933dd3f60004731034fa93ddd35a16eea802d2cf9a1baa2619e6360dc27
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.