File size: 8,419 Bytes
b11e73e | 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
"""Phoneme-input adapter for Fish S2 Pro (run from the fish-speech root, idempotent).
1. tokenizer.py : when EGY_PHON=1, the content of <|phoneme_start|>...<|phoneme_end|> is read as space-separated symbols of a
fixed 42-symbol Egyptian inventory and mapped to atomic ids PHON_BASE+i (never BPE pieces).
2. llama.py : BaseTransformer gets `phon_embeddings` (64 x dim, input-only); embed()/forward_generate() look phoneme ids up
there; after loading base weights each symbol row is initialised to the mean of its BPE-piece embeddings.
3. lora.py : setup_lora() leaves phon_embeddings trainable.
4. lit_module.py: LoRA-only checkpoints also keep phon_embeddings."""
import re
from pathlib import Path
INVENTORY = ["sˤ", "dˤ", "tˤ", "zˤ", "aː", "eː", "iː", "oː", "uː",
"b", "t", "d", "k", "g", "q", "ʔ", "ʕ", "ħ", "x", "ɣ", "ʃ", "ʒ", "s", "z", "f", "v", "p", "m", "n", "l", "r", "w", "y", "h",
"a", "e", "i", "o", "u", "_"]
def patch(path, old, new, must=True):
p = Path(path)
s = p.read_text(encoding="utf-8")
if new in s:
print(f" {path}: already patched")
return
if old not in s:
if must:
raise SystemExit(f"anchor not found in {path}: {old[:70]!r}")
print(f" {path}: anchor missing, skipped")
return
p.write_text(s.replace(old, new, 1), encoding="utf-8")
print(f" {path}: patched")
# ---------------------------------------------------------------- 1. tokenizer
tok_old = ''' def encode(
self, text: str, add_special_tokens: bool = False, **kwargs
) -> List[int]:
# [FIX] Force Qwen/Tiktoken backends to parse special tokens inline
import inspect
'''
tok_new = ''' # [egy-phon] atomic Egyptian phoneme symbols -> input-only ids (PHON_BASE + index); active when EGY_PHON=1
PHON_BASE = 200000
PHON_INVENTORY = %r
PHON_SYM2ID = {s: i for i, s in enumerate(PHON_INVENTORY)}
def encode(
self, text: str, add_special_tokens: bool = False, **kwargs
) -> List[int]:
import os as _os, re as _re
if _os.environ.get("EGY_PHON") == "1" and "<|phoneme_start|>" in text and "<|phoneme_end|>" in text:
ids, pos = [], 0
for m in _re.finditer(r"<\\|phoneme_start\\|>(.*?)<\\|phoneme_end\\|>", text, flags=_re.S):
if m.start() > pos:
ids += self._encode_plain(text[pos:m.start()], add_special_tokens, **kwargs)
ids.append(self.get_token_id("<|phoneme_start|>"))
ids += [self.PHON_BASE + self.PHON_SYM2ID[s] for s in m.group(1).split() if s in self.PHON_SYM2ID]
ids.append(self.get_token_id("<|phoneme_end|>"))
pos = m.end()
if pos < len(text):
ids += self._encode_plain(text[pos:], add_special_tokens, **kwargs)
return ids
return self._encode_plain(text, add_special_tokens, **kwargs)
def _encode_plain(
self, text: str, add_special_tokens: bool = False, **kwargs
) -> List[int]:
# [FIX] Force Qwen/Tiktoken backends to parse special tokens inline
import inspect
''' % (INVENTORY,)
patch("fish_speech/tokenizer.py", tok_old, tok_new)
patch("fish_speech/tokenizer.py",
''' def decode(self, tokens: Union[List[int], int], **kwargs) -> str:
return self._tokenizer.decode(tokens, **kwargs)''',
''' def decode(self, tokens: Union[List[int], int], **kwargs) -> str:
if isinstance(tokens, list): # [egy-phon] phoneme ids are input-only: render them as their symbols
if any(isinstance(t, int) and t >= self.PHON_BASE for t in tokens):
out, buf = [], []
for t in tokens:
if isinstance(t, int) and t >= self.PHON_BASE:
buf.append(self.PHON_INVENTORY[t - self.PHON_BASE] if t - self.PHON_BASE < len(self.PHON_INVENTORY) else "?")
else:
if buf:
out.append(" ".join(buf)); buf = []
out.append(self._tokenizer.decode([t], **kwargs))
if buf:
out.append(" ".join(buf))
return "".join(out)
return self._tokenizer.decode(tokens, **kwargs)''')
# ---------------------------------------------------------------- 2. model
llama = "fish_speech/models/text2semantic/llama.py"
patch(llama,
''' self.codebook_embeddings = nn.Embedding(
config.codebook_size * config.num_codebooks,
config.dim,
)
self.layers = nn.ModuleList(''',
''' self.codebook_embeddings = nn.Embedding(
config.codebook_size * config.num_codebooks,
config.dim,
)
# [egy-phon] input-only atomic phoneme embeddings (ids >= PHON_BASE in the text row); never predicted
self.phon_embeddings = nn.Embedding(64, config.dim)
self.layers = nn.ModuleList(''')
patch(llama,
''' def embed(self, inp: Tensor) -> Tensor:
embeds = []
''',
''' PHON_BASE = 200000
def _tok_embed(self, tok: Tensor) -> Tensor:
"""[egy-phon] text-row lookup: base vocabulary or the phoneme table"""
is_phon = tok >= self.PHON_BASE
if not bool(is_phon.any()):
return self.embeddings(tok)
base = self.embeddings(torch.where(is_phon, torch.zeros_like(tok), tok))
ph = self.phon_embeddings((tok - self.PHON_BASE).clamp(min=0, max=63))
return torch.where(is_phon.unsqueeze(-1), ph.to(base.dtype), base)
@torch.no_grad()
def init_phon_embeddings(self, tokenizer_path: str, inventory: list) -> None:
"""[egy-phon] initialise each symbol row to the mean of its BPE-piece embeddings (Hebrew S2-Pro recipe)"""
from transformers import AutoTokenizer
tk = AutoTokenizer.from_pretrained(tokenizer_path)
w = self.embeddings.weight
for i, sym in enumerate(inventory):
pieces = tk.encode(" " if sym == "_" else sym, add_special_tokens=False)
if pieces:
self.phon_embeddings.weight[i] = w[torch.tensor(pieces, device=w.device)].float().mean(0).to(w.dtype)
def embed(self, inp: Tensor) -> Tensor:
embeds = []
''')
patch(llama,
''' vq_embeds_sum[~is_semantic] = 0
x = self.embeddings(inp[:, 0]) + vq_embeds_sum
''',
''' vq_embeds_sum[~is_semantic] = 0
x = self._tok_embed(inp[:, 0]) + vq_embeds_sum
''')
patch(llama,
''' vq_embeds_sum[~vq_masks] = 0
x = self.embeddings(inp[:, 0]) + vq_embeds_sum
''',
''' vq_embeds_sum[~vq_masks] = 0
x = self._tok_embed(inp[:, 0]) + vq_embeds_sum
''')
patch(llama,
''' err = model.load_state_dict(weights, strict=False, assign=True)''',
''' err = model.load_state_dict(weights, strict=False, assign=True)
try: # [egy-phon] symbol rows start as the mean of their BPE pieces; a phoneme checkpoint loaded later overrides them
from fish_speech.tokenizer import FishTokenizer as _FT
model.init_phon_embeddings(str(path), _FT.PHON_INVENTORY)
logger.info("[egy-phon] phoneme embeddings initialised from BPE-piece means")
except Exception as _e:
logger.warning(f"[egy-phon] phoneme embedding init skipped: {_e}")''')
# ---------------------------------------------------------------- 3. lora
patch("fish_speech/models/text2semantic/lora.py",
''' lora.mark_only_lora_as_trainable(model, bias="none")''',
''' lora.mark_only_lora_as_trainable(model, bias="none")
for _n, _p in model.named_parameters(): # [egy-phon] the phoneme table trains alongside the adapters
if "phon_embeddings" in _n:
_p.requires_grad = True''')
# ---------------------------------------------------------------- 4. checkpoint filter
patch("fish_speech/models/text2semantic/lit_module.py",
''' for name in list(state_dict.keys()):
if "lora" not in name:
state_dict.pop(name)''',
''' for name in list(state_dict.keys()):
if "lora" not in name and "phon_embeddings" not in name: # [egy-phon]
state_dict.pop(name)''')
print("PHON_PATCH_DONE")
|