File size: 5,754 Bytes
e4ae2bb | 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 | """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))
|