| |
| """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") |
|
|
|
|
| |
| 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)''') |
|
|
| |
| 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}")''') |
|
|
| |
| 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''') |
|
|
| |
| 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") |
|
|
