File size: 6,882 Bytes
b9301f1 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 | #!/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()
|