ATOM (atom-ai): code snapshot (100fae0 + batched-train), README with measured numbers, curated checkpoints (part 4)
b9301f1 verified | #!/usr/bin/env python3 | |
| """Idempotent speech-lock patcher for src/atom_native.py. | |
| python3 scripts/apply_speech_lock.py | |
| python3 -m pytest test/test_speech_lock.py -q | |
| python3 scripts/speech_chat.py --prompt Bonjour | |
| """ | |
| from __future__ import annotations | |
| import ast | |
| from pathlib import Path | |
| ROOT = Path(__file__).resolve().parents[1] | |
| TARGET = ROOT / "src" / "atom_native.py" | |
| IMPORT = "from .speech_lock import format_dialogue_prompt, speech_ok\n" | |
| GENERATE = ''' def generate_packets( | |
| self, | |
| prompt: str, | |
| max_packets: int = 8, | |
| temperature: float = 1.0, | |
| top_k: int | None = None, | |
| deterministic: bool = True, | |
| max_length: int | None = None, | |
| prefer_printable: bool = True, | |
| reset: bool = True, | |
| merge_enabled: bool = False, | |
| payload_copy: bool = False, | |
| dialogue_wrap: bool = True, | |
| speech_gate: bool = True, | |
| ) -> bytes: | |
| """Prime on a prompt and generate bounded atom surface packets.""" | |
| self.eval() | |
| if self.field_obligatory_hard: | |
| self.surface.set_obligatory_hard(True) | |
| prefer_printable = True | |
| prev_payload_flag = bool(getattr(self.surface, "payload_enabled", True)) | |
| self.surface.payload_enabled = bool(payload_copy) | |
| try: | |
| if reset: | |
| self.reset_state(reset_atomizer=True) | |
| if dialogue_wrap: | |
| prompt = format_dialogue_prompt(prompt) | |
| prompt_packets = self.atomizer.encode(prompt, reset=reset) | |
| if not prompt_packets: | |
| raise ValueError("prompt must produce at least one atom packet") | |
| for packet in prompt_packets[:-1]: | |
| self.forward_packet(packet, merge_enabled=merge_enabled) | |
| current = prompt_packets[-1] | |
| generated = bytearray() | |
| for _ in range(max_packets): | |
| payload = self.predict_packet( | |
| current, | |
| temperature=temperature, | |
| top_k=top_k, | |
| deterministic=deterministic, | |
| prefer_printable=prefer_printable, | |
| merge_enabled=merge_enabled, | |
| ) | |
| if not payload: | |
| break | |
| if speech_gate and not speech_ok(payload): | |
| break | |
| generated.extend(payload) | |
| current = self.atomizer.packet_from_payload(payload) | |
| if max_length is not None and len(generated) >= max_length: | |
| break | |
| if generated.endswith(b"\\n\\n") and len(generated) > 2: | |
| break | |
| if max_length is not None: | |
| return bytes(generated[:max_length]) | |
| return bytes(generated) | |
| finally: | |
| self.surface.payload_enabled = prev_payload_flag | |
| ''' | |
| def main() -> None: | |
| text = TARGET.read_text(encoding="utf-8") | |
| changed = False | |
| if "from .speech_lock import" not in text: | |
| needle = "from .io.atomizer import AtomPacket, Atomizer\n" | |
| if needle not in text: | |
| raise SystemExit("atomizer import not found") | |
| text = text.replace(needle, needle + IMPORT, 1) | |
| changed = True | |
| print("[apply] import speech_lock") | |
| old_copy = ( | |
| " copy = F.softplus(self.payload_copy_scale) * 4.0\n" | |
| " # Per-atom copy-bias: winner atom's own hist / positions only.\n" | |
| " for i, (payload, hist_i) in enumerate(zip(kept_payloads, local_hists)):\n" | |
| " w = copy_w[i]\n" | |
| " mass_i = hist_i.sum().clamp_min(1.0)\n" | |
| " byte_logits = byte_logits + copy * w * (hist_i / mass_i).unsqueeze(0)\n" | |
| " for pos, b in enumerate(payload[:L]):\n" | |
| " byte_logits[pos, int(b)] = byte_logits[pos, int(b)] + copy * w\n" | |
| ) | |
| new_copy = ( | |
| " copy = F.softplus(self.payload_copy_scale) * 4.0\n" | |
| " # Position-aligned copy only. No unigram broadcast across span slots.\n" | |
| " for i, payload in enumerate(kept_payloads):\n" | |
| " w = copy_w[i]\n" | |
| " for pos, b in enumerate(payload[:L]):\n" | |
| " byte_logits[pos, int(b)] = byte_logits[pos, int(b)] + copy * w\n" | |
| ) | |
| if old_copy in text: | |
| text = text.replace(old_copy, new_copy, 1) | |
| changed = True | |
| print("[apply] position-aligned copy") | |
| elif "No unigram broadcast" in text: | |
| print("[apply] copy already aligned") | |
| if "self.payload_enabled = True" not in text: | |
| old_init = ( | |
| " self.payload_copy_scale = nn.Parameter(torch.tensor(1.0))\n" | |
| " self._init_payload_production()\n" | |
| ) | |
| new_init = ( | |
| " self.payload_copy_scale = nn.Parameter(torch.tensor(1.0))\n" | |
| " self.payload_enabled = True\n" | |
| " self._init_payload_production()\n" | |
| ) | |
| if old_init in text: | |
| text = text.replace(old_init, new_init, 1) | |
| changed = True | |
| print("[apply] payload_enabled flag") | |
| old_pay = ( | |
| " pay_byte, pay_len = self.payload_produce(\n" | |
| " atom_payloads,\n" | |
| " a_hat,\n" | |
| " a_only,\n" | |
| " atom_energies=atom_energies,\n" | |
| " )\n" | |
| ) | |
| new_pay = ( | |
| " if getattr(self, \"payload_enabled\", True):\n" | |
| " pay_byte, pay_len = self.payload_produce(\n" | |
| " atom_payloads,\n" | |
| " a_hat,\n" | |
| " a_only,\n" | |
| " atom_energies=atom_energies,\n" | |
| " )\n" | |
| " else:\n" | |
| " pay_byte = torch.zeros_like(frozen_byte)\n" | |
| " pay_len = torch.zeros_like(frozen_len)\n" | |
| ) | |
| if old_pay in text: | |
| text = text.replace(old_pay, new_pay, 1) | |
| changed = True | |
| print("[apply] gated payload_produce") | |
| start = text.find(" def generate_packets(") | |
| end = text.find(" def _checkpoint(self)") | |
| if start < 0 or end < 0: | |
| raise SystemExit("generate_packets / _checkpoint anchors not found") | |
| if text[start:end] != GENERATE: | |
| text = text[:start] + GENERATE + "\n" + text[end:] | |
| changed = True | |
| print("[apply] generate_packets rewritten") | |
| ast.parse(text) | |
| if changed: | |
| TARGET.write_text(text, encoding="utf-8") | |
| print(f"[apply] wrote {TARGET}") | |
| else: | |
| print("[apply] already applied") | |
| if __name__ == "__main__": | |
| main() | |