Dynamic Accent Flash: add Flash-DPO (CosyVoice3) as --version flash (default), diacritic front-end, Flash results, samples
e4ae2bb verified | """Dynamic Accent Flash text front-end. | |
| from frontend import FrontEnd | |
| fe = FrontEnd(diac_model="models/diacritics", mt=True) | |
| fe("<eng>The meeting is at 3:30 on 21/10/2026.</eng>", mode="codeswitch") # tags kept, numbers verbalised | |
| fe("Ojo naa dara pupo", mode="native", lang="yor") # adds Yorùbá tone marks/underdots | |
| fe("Please come home early today.", mode="translate", lang="ibo") # MADLAD-400 eng/pcm -> ibo | |
| fe("I go call you tomorrow", mode="accent", lang="ibo") # English/Pidgin words, Igbo accent | |
| Modes | |
| codeswitch : input already tagged <ibo>…</ibo> <eng>…</eng> (or untagged = one language `lang`); each segment is | |
| normalised in its own language; Yorùbá/Igbo segments that carry no marks are diacritised. | |
| native : the whole text is one language (`lang`); normalise + diacritise. | |
| translate : `lang` is the target (ibo|yor|hau); the input (English or Pidgin) is translated first with | |
| MADLAD-400-3B-MT (Apache-2.0), then normalised + diacritised (MADLAD Yorùbá is often unmarked). | |
| accent : keep the words as written but read them with the `lang` accent (whole text tagged <lang>), i.e. | |
| "English words with a native accent". Numbers are verbalised in English. | |
| Output: a tagged string for infer.py / flash_lm.py. | |
| """ | |
| import re | |
| import unicodedata | |
| try: | |
| from .normalize import normalize | |
| except ImportError: # run as a script | |
| from normalize import normalize | |
| SEG = re.compile(r"<(\w{3})>(.*?)</\1>", re.S) | |
| def _strip(s): | |
| return "".join(c for c in unicodedata.normalize("NFD", s) if unicodedata.category(c) != "Mn") | |
| def unmarked(text, thresh=0.05): | |
| w = re.findall(r"\w+", unicodedata.normalize("NFC", text)) | |
| if not w: | |
| return False | |
| return sum(_strip(x) != unicodedata.normalize("NFC", x) for x in w) / len(w) < thresh | |
| class FrontEnd: | |
| def __init__(self, diac_model=None, mt=False, device="cuda"): | |
| self.diac_dir, self.device = diac_model, device | |
| self._d = self._mt = None | |
| self.use_mt = mt | |
| # ---- diacritics (ByT5-small fine-tune, word-guarded) | |
| def diacritise(self, text, lang): | |
| if lang not in ("yor", "ibo") or not self.diac_dir: | |
| return text | |
| import torch | |
| from transformers import AutoTokenizer, T5ForConditionalGeneration | |
| if self._d is None: | |
| self._d = (AutoTokenizer.from_pretrained(self.diac_dir), | |
| T5ForConditionalGeneration.from_pretrained(self.diac_dir).to(self.device).eval()) | |
| tok, m = self._d | |
| src = unicodedata.normalize("NFC", _strip(text)) | |
| x = tok([f"{lang}: {src}"], return_tensors="pt").to(self.device) | |
| with torch.no_grad(): | |
| y = m.generate(**x, max_length=int(x.input_ids.shape[1] * 1.6) + 16) | |
| hyp = unicodedata.normalize("NFC", tok.decode(y[0], skip_special_tokens=True)) | |
| sw, hw = src.split(), hyp.split() | |
| if len(sw) != len(hw): | |
| return text | |
| return " ".join(h if unicodedata.normalize("NFC", _strip(h)) == s else s for s, h in zip(sw, hw)) | |
| # ---- MT (MADLAD-400-3B-MT) | |
| def translate(self, text, tgt): | |
| import torch | |
| from transformers import T5ForConditionalGeneration, T5Tokenizer | |
| if self._mt is None: | |
| self._mt = (T5Tokenizer.from_pretrained("google/madlad400-3b-mt"), | |
| T5ForConditionalGeneration.from_pretrained("google/madlad400-3b-mt", torch_dtype=torch.bfloat16).to(self.device).eval()) | |
| tok, m = self._mt | |
| code = {"ibo": "ig", "yor": "yo", "hau": "ha", "eng": "en"}[tgt] | |
| out = [] | |
| for sent in re.split(r"(?<=[.!?])\s+", text.strip()): | |
| if not sent: | |
| continue | |
| x = tok([f"<2{code}> {sent}"], return_tensors="pt").to(self.device) | |
| with torch.no_grad(): | |
| y = m.generate(**x, max_new_tokens=256, num_beams=4) | |
| out.append(tok.decode(y[0], skip_special_tokens=True)) | |
| return " ".join(out) | |
| def _seg(self, lang, s): | |
| s = normalize(s, lang) | |
| if lang in ("yor", "ibo") and unmarked(s): | |
| s = self.diacritise(s, lang) | |
| return f"<{lang}>{s}</{lang}>" | |
| def __call__(self, text, mode="codeswitch", lang="eng"): | |
| text = unicodedata.normalize("NFC", text).strip() | |
| if mode == "translate": | |
| if lang not in ("ibo", "yor", "hau"): | |
| raise ValueError("translate mode needs lang in ibo|yor|hau") | |
| plain = re.sub(r"</?\w{3}>", "", text) | |
| return self._seg(lang, self.translate(normalize(plain, "eng"), lang)) | |
| if mode == "accent": | |
| plain = re.sub(r"</?\w{3}>", "", text) | |
| return f"<{lang}>{normalize(plain, 'eng')}</{lang}>" | |
| if mode == "native" or not SEG.search(text): | |
| return self._seg(lang, re.sub(r"</?\w{3}>", "", text)) | |
| parts, pos = [], 0 | |
| for m in SEG.finditer(text): | |
| if text[pos:m.start()].strip(): | |
| parts.append(self._seg(lang, text[pos:m.start()])) | |
| parts.append(self._seg(m.group(1), m.group(2))) | |
| pos = m.end() | |
| if text[pos:].strip(): | |
| parts.append(self._seg(lang, text[pos:])) | |
| return " ".join(parts) | |
| if __name__ == "__main__": | |
| import argparse | |
| ap = argparse.ArgumentParser() | |
| ap.add_argument("text"); ap.add_argument("--mode", default="codeswitch"); ap.add_argument("--lang", default="eng") | |
| ap.add_argument("--diac_model", default=None); ap.add_argument("--cpu", action="store_true") | |
| a = ap.parse_args() | |
| print(FrontEnd(a.diac_model, device="cpu" if a.cpu else "cuda")(a.text, a.mode, a.lang)) | |