#!/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")