ATOM (atom-ai): code snapshot (100fae0 + batched-train), README with measured numbers, curated checkpoints (part 4)
b9301f1 verified | #!/usr/bin/env python3 | |
| """Apply --field-obligatory-hard (hard-v2 scaffold) to atom_native.py (idempotent). | |
| Note: hard-v2 forward (frozen+α_proj) lives in src/atom_native.py; this script | |
| only scaffolds the hard flag / freeze helpers if missing. | |
| """ | |
| from __future__ import annotations | |
| from pathlib import Path | |
| ROOT = Path(__file__).resolve().parents[1] | |
| TARGET = ROOT / "src" / "atom_native.py" | |
| INIT_OLD = ( | |
| " self.field_obligatory_readout = False # enabled via AtomNativeModel flag\n" | |
| ) | |
| INIT_NEW = ( | |
| " self.field_obligatory_readout = False # enabled via AtomNativeModel flag\n" | |
| " self.field_obligatory_hard = False # mix floor=1.0 + freeze non-α bypass\n" | |
| " self._soft_obl_mix_floor = 0.35\n" | |
| ) | |
| METHODS = ''' | |
| def set_obligatory_hard(self, enabled: bool) -> None: | |
| """Hard obligatory: mix floor 1.0; zero/freeze trainable non-α bypass.""" | |
| enabled = bool(enabled) | |
| self.field_obligatory_hard = enabled | |
| if enabled: | |
| self.field_obligatory_readout = True | |
| self.obl_mix_floor = 1.0 | |
| self._freeze_non_alpha_bypass(True) | |
| else: | |
| self.obl_mix_floor = float(getattr(self, "_soft_obl_mix_floor", 0.35)) | |
| self._freeze_non_alpha_bypass(False) | |
| def _freeze_non_alpha_bypass(self, freeze: bool) -> None: | |
| """Zero + freeze decoder/skip/gates that can bypass α on the CE path.""" | |
| modules = ( | |
| self.field_to_state, | |
| self.field_byte_skip, | |
| self.field_length_skip, | |
| self.byte_decoder, | |
| self.length_decoder, | |
| ) | |
| scalars = (self.field_gate, self.skip_gate, self.obl_mix) | |
| with torch.no_grad(): | |
| if freeze: | |
| for module in modules: | |
| for param in module.parameters(): | |
| param.zero_() | |
| param.requires_grad_(False) | |
| for param in scalars: | |
| param.zero_() | |
| param.requires_grad_(False) | |
| else: | |
| for module in modules: | |
| for param in module.parameters(): | |
| param.requires_grad_(True) | |
| for param in scalars: | |
| param.requires_grad_(True) | |
| ''' | |
| FORWARD_NEEDLE = ( | |
| " train_byte = self.alpha_byte_proj(a_hat).view(self.max_payload_bytes, 256)\n" | |
| " train_len = self.alpha_length_proj(a_hat)\n" | |
| " mix = self.obligatory_mix_weight()\n" | |
| ) | |
| FORWARD_INSERT = ( | |
| " if getattr(self, \"field_obligatory_hard\", False):\n" | |
| " # Hard: mix≡1 and non-α bypass residual forced off.\n" | |
| " return {\"byte_logits\": frozen_byte, \"length_logits\": frozen_len}\n" | |
| " train_byte = self.alpha_byte_proj(a_hat).view(self.max_payload_bytes, 256)\n" | |
| " train_len = self.alpha_length_proj(a_hat)\n" | |
| " mix = self.obligatory_mix_weight()\n" | |
| ) | |
| MODEL_INIT_OLD = ( | |
| " field_obligatory_readout: bool = False,\n" | |
| " slow_rms_rel_tol: float = 0.15,\n" | |
| ) | |
| MODEL_INIT_NEW = ( | |
| " field_obligatory_readout: bool = False,\n" | |
| " field_obligatory_hard: bool = False,\n" | |
| " slow_rms_rel_tol: float = 0.15,\n" | |
| ) | |
| MODEL_ASSIGN_OLD = ( | |
| " self.field_obligatory_readout = bool(field_obligatory_readout)\n" | |
| " self.surface.field_obligatory_readout = self.field_obligatory_readout\n" | |
| ) | |
| MODEL_ASSIGN_NEW = ( | |
| " self.field_obligatory_hard = bool(field_obligatory_hard)\n" | |
| " self.field_obligatory_readout = bool(field_obligatory_readout) or self.field_obligatory_hard\n" | |
| " self.surface.field_obligatory_readout = self.field_obligatory_readout\n" | |
| " if self.field_obligatory_hard:\n" | |
| " self.surface.set_obligatory_hard(True)\n" | |
| " else:\n" | |
| " self.surface.field_obligatory_hard = False\n" | |
| ) | |
| def main() -> None: | |
| text = TARGET.read_text(encoding="utf-8") | |
| if "def set_obligatory_hard" in text and "field_obligatory_hard: bool" in text: | |
| print(f"already applied: {TARGET}") | |
| return | |
| if INIT_OLD not in text: | |
| raise SystemExit("INIT_OLD anchor not found") | |
| text = text.replace(INIT_OLD, INIT_NEW, 1) | |
| if "def set_obligatory_hard" not in text: | |
| anchor = " def field_features(" | |
| if anchor not in text: | |
| raise SystemExit("field_features anchor not found") | |
| text = text.replace(anchor, METHODS + anchor, 1) | |
| if "field_obligatory_hard" in text and "return {\"byte_logits\": frozen_byte" in text: | |
| pass | |
| elif FORWARD_NEEDLE not in text: | |
| raise SystemExit("FORWARD_NEEDLE not found") | |
| else: | |
| text = text.replace(FORWARD_NEEDLE, FORWARD_INSERT, 1) | |
| if MODEL_INIT_OLD not in text: | |
| raise SystemExit("MODEL_INIT_OLD not found") | |
| text = text.replace(MODEL_INIT_OLD, MODEL_INIT_NEW, 1) | |
| if MODEL_ASSIGN_OLD not in text: | |
| raise SystemExit("MODEL_ASSIGN_OLD not found") | |
| text = text.replace(MODEL_ASSIGN_OLD, MODEL_ASSIGN_NEW, 1) | |
| TARGET.write_text(text, encoding="utf-8") | |
| print(f"applied hard obligatory to {TARGET}") | |
| if __name__ == "__main__": | |
| main() | |