#!/usr/bin/env python3
"""Patch fish_speech lora.py: add slow_* target names that do NOT imply fast_*.

Upstream semantics: "attention" targets slow AND fast (backwards compat).
After patch: "slow_attention"/"slow_mlp"/"slow_embeddings"/"slow_output"
target ONLY the slow (text) transformer, leaving the RL-aligned audio decoder
(fast transformer) frozen.
"""

from pathlib import Path

P = Path("/opt/work/fish-speech/fish_speech/models/text2semantic/lora.py")
src = P.read_text(encoding="utf-8")

old = '''    # Slow transformer: targeted by unprefixed names (e.g. "attention")
    slow_attention = "attention" in targets
    slow_mlp = "mlp" in targets
    slow_embeddings = "embeddings" in targets
    slow_output = "output" in targets'''
new = '''    # Slow transformer: targeted by unprefixed names (e.g. "attention")
    # or explicit "slow_*" names (which do NOT imply fast_*)
    slow_attention = "attention" in targets or "slow_attention" in targets
    slow_mlp = "mlp" in targets or "slow_mlp" in targets
    slow_embeddings = "embeddings" in targets or "slow_embeddings" in targets
    slow_output = "output" in targets or "slow_output" in targets'''

old2 = '''    fast_attention = slow_attention or "fast_attention" in targets
    fast_mlp = slow_mlp or "fast_mlp" in targets
    fast_embeddings = slow_embeddings or "fast_embeddings" in targets
    fast_output = slow_output or "fast_output" in targets'''
new2 = '''    fast_attention = "attention" in targets or "fast_attention" in targets
    fast_mlp = "mlp" in targets or "fast_mlp" in targets
    fast_embeddings = "embeddings" in targets or "fast_embeddings" in targets
    fast_output = "output" in targets or "fast_output" in targets'''

assert old in src, "anchor 1 not found"
assert old2 in src, "anchor 2 not found"
src = src.replace(old, new).replace(old2, new2)
P.write_text(src, encoding="utf-8")
print("lora.py patched: slow_* names available")

# verify semantics quickly
import re
assert 'slow_attention = "attention" in targets or "slow_attention"' in src
assert 'fast_attention = "attention" in targets or "fast_attention"' in src
print("verified")