"""Code-switch text front end (pipeline step, NOT native model capability). Turns a plain sentence + a code-switch request ("mostly English, Igbo greetings") into language-tagged text (.. ..) that the TTS then speaks. Uses a permissively licensed text LLM: default Qwen/Qwen3-14B (Apache-2.0), thinking disabled. N-ATLaS was not used: its terms cap use at 1,000 active users and require a commercial licence beyond that. Safety rails (MADLAD looped in Flash): max_new_tokens tied to input length, repetition_penalty, no_repeat_ngram, output validation (tags only from allowed set, full coverage, length ratio, n-gram repetition cap), up to 2 retries, then a safe fallback (the input as single-language text) with a warning. """ import re, os, json from collections import Counter LANGS = {"eng": "English", "ibo": "Igbo", "yor": "Yoruba", "hau": "Hausa", "pcm": "Nigerian Pidgin"} SYS = """You rewrite text for a Nigerian text-to-speech system as natural code-switched speech. Follow the user's code-switch request exactly (which languages, how much of each, which parts). Rules: - Keep the meaning. Do not add new information, greetings or sentences that the request does not ask for. - Write Igbo, Yoruba and Hausa in standard orthography (use tone marks/diacritics for Yoruba and Igbo if you can). - Wrap every span in a language tag: ..., ..., ..., ..., .... - Nothing may appear outside tags. Output ONLY the tagged text on one line, no explanations. - Prefer these common, correct phrases when greetings/thanks/closings are requested: Igbo: Ndewo (hello), Ụtụtụ ọma (good morning), Ehihie ọma (good afternoon), Mgbede ọma (good evening), Kedu? (how are you), Daalụ (thank you), Nnọọ (welcome), Ka ọ dị (goodbye) Yoruba: Ẹ káàárọ̀ (good morning), Ẹ káàsán (good afternoon), Ẹ kú ìrọ̀lẹ́ (good evening), Ẹ kú àbọ̀ (welcome), Ẹ ṣé (thank you), Ó dàbọ̀ (goodbye) Hausa: Sannu (hello), Barka da safiya (good morning), Barka da yamma (good evening), Maraba (welcome), Na gode (thank you), Sai an jima (goodbye) Nigerian Pidgin: How far (hello), Abeg (please), Thank you well well, We go see later (goodbye) - If you cannot translate a part reliably, keep that part in English rather than guessing.""" SHOT_U = "Request: mostly English, Igbo greetings and closing\nText: Good morning everyone. The meeting starts at ten, please be on time. Thank you." SHOT_A = "Ụtụtụ ọma, unu niile. The meeting starts at ten, please be on time. Daalụ." TAG = re.compile(r"<(eng|ibo|yor|hau|pcm)>(.*?)", re.S) def requested_langs(req): r = req.lower() out = {k for k, v in LANGS.items() if v.lower().split()[-1] in r or k in r} if "pidgin" in r: out.add("pcm") return out or {"eng"} def validate(out, src, allowed): out = out.strip() spans = TAG.findall(out) if not spans: return False, "no tags" rest = TAG.sub(" ", out).strip() if re.sub(r"[\s.,!?;:'\"-]", "", rest): return False, f"text outside tags: {rest[:40]}" bad = {l for l, _ in spans} - allowed - {"eng"} if bad: return False, f"unrequested language {bad}" words = " ".join(t for _, t in spans).split() n_src = len(src.split()) if not (0.4 * n_src <= len(words) <= 2.5 * n_src + 6): return False, f"length ratio {len(words)}/{n_src}" low = [w.lower().strip(".,!?") for w in words] for k in range(len(low) - 3): if low[k] == low[k + 1] == low[k + 2] == low[k + 3]: return False, "word loop" tri = Counter(tuple(low[i:i + 3]) for i in range(len(low) - 2)) if tri and tri.most_common(1)[0][1] > 2: return False, "trigram loop" if any(not t.strip() for _, t in spans): return False, "empty span" return True, "ok" def chrf(hyp, ref, n=6, beta=2.0): def grams(s, k): s = " ".join(s.lower().split()) return Counter(s[i:i + k] for i in range(len(s) - k + 1)) P, R = [], [] for k in range(1, n + 1): h, r = grams(hyp, k), grams(ref, k) if not h or not r: continue m = sum((h & r).values()) P.append(m / sum(h.values())); R.append(m / sum(r.values())) if not P: return 0.0 p, r = sum(P) / len(P), sum(R) / len(R) return 0.0 if p + r == 0 else 100 * (1 + beta ** 2) * p * r / (beta ** 2 * p + r) class CodeSwitchFrontEnd: def __init__(self, model=os.environ.get("CS_LLM", "Qwen/Qwen3-14B"), device="cuda", cache_dir=None, diac=None): import torch from transformers import AutoTokenizer, AutoModelForCausalLM self.torch = torch self.tok = AutoTokenizer.from_pretrained(model, cache_dir=cache_dir) self.lm = AutoModelForCausalLM.from_pretrained(model, torch_dtype=torch.bfloat16, cache_dir=cache_dir).to(device).eval() self.device = device self.diac = diac # optional flash FrontEnd (ByT5 diacritic restoration) for ibo/yor spans def _gen(self, req, text, temp): msgs = [{"role": "system", "content": SYS}, {"role": "user", "content": SHOT_U}, {"role": "assistant", "content": SHOT_A}, {"role": "user", "content": f"Request: {req}\nText: {text}"}] try: p = self.tok.apply_chat_template(msgs, tokenize=False, add_generation_prompt=True, enable_thinking=False) except TypeError: p = self.tok.apply_chat_template(msgs, tokenize=False, add_generation_prompt=True) ids = self.tok(p, return_tensors="pt").to(self.device) n_in = len(self.tok(text)["input_ids"]) # repetition rails act on the GENERATED tokens only (a prompt-wide no_repeat_ngram/repetition_penalty # forbids copying the input and garbles spelling): hard length cap + loop stop + validation below from transformers import StoppingCriteria, StoppingCriteriaList P0 = ids["input_ids"].shape[1] class LoopStop(StoppingCriteria): def __call__(s_, input_ids, scores, **k): g = input_ids[0, P0:].tolist() if len(g) < 18: return False tail = tuple(g[-6:]) hits = sum(1 for i in range(len(g) - 6) if tuple(g[i:i + 6]) == tail) return hits >= 2 kw = dict(max_new_tokens=int(3 * n_in + 48), pad_token_id=self.tok.eos_token_id, stopping_criteria=StoppingCriteriaList([LoopStop()])) if temp > 0: kw.update(do_sample=True, temperature=temp, top_p=0.9) else: kw.update(do_sample=False, temperature=None, top_p=None, top_k=None) with self.torch.no_grad(): o = self.lm.generate(**ids, **kw) s = self.tok.decode(o[0, ids["input_ids"].shape[1]:], skip_special_tokens=True) s = re.sub(r".*?", "", s, flags=re.S).strip().splitlines() return s[0].strip() if s else "" def back_translate(self, span, lang): msgs = [{"role": "user", "content": f"Translate this {LANGS[lang]} text into English. Output only the translation.\n\n{span}"}] try: p = self.tok.apply_chat_template(msgs, tokenize=False, add_generation_prompt=True, enable_thinking=False) except TypeError: p = self.tok.apply_chat_template(msgs, tokenize=False, add_generation_prompt=True) ids = self.tok(p, return_tensors="pt").to(self.device) with self.torch.no_grad(): o = self.lm.generate(**ids, max_new_tokens=3 * len(span.split()) + 24, do_sample=False, temperature=None, top_p=None, top_k=None, pad_token_id=self.tok.eos_token_id) return self.tok.decode(o[0, ids["input_ids"].shape[1]:], skip_special_tokens=True).strip() def meaning_check(self, out, src, thr=45.0): """round trip: back-translate every non-English span; the English reconstruction must match the input (chrF).""" parts = [t if l == "eng" else self.back_translate(t, l) for l, t in TAG.findall(out)] back = " ".join(parts) score = min(chrf(back, src), chrf(back, src, beta=0.5)) # recall- and precision-weighted (added content fails) return score >= thr, round(score, 1), back def __call__(self, text, req, fallback_lang="eng"): allowed = requested_langs(req) log = [] for k, temp in enumerate((0.0, 0.6, 0.8)): # first try + at most 2 retries out = self._gen(req, text, temp) ok, why = validate(out, text, allowed) rt = None if ok: ok, sc, back = self.meaning_check(out, text) rt = dict(chrf=sc, back=back[:300]) if not ok: why = f"meaning drift (round-trip chrF {sc})" log.append(dict(try_=k, ok=ok, why=why, out=out[:300], roundtrip=rt)) if ok: if self.diac is not None: out = TAG.sub(lambda m: f"<{m.group(1)}>{self.diac.diacritise(m.group(2), m.group(1)) if m.group(1) in ('ibo', 'yor') else m.group(2)}", out) return out, dict(ok=True, tries=log) return f"<{fallback_lang}>{text}", dict(ok=False, tries=log, warning="code-switch front end failed validation; spoke the input as-is") if __name__ == "__main__": import sys cache = os.environ.get("HF_CACHE") fe = CodeSwitchFrontEnd(cache_dir=cache) tests = json.load(open(sys.argv[1])) if len(sys.argv) > 1 else [ ["Good morning everyone, welcome to the market. Prices are good today, come and buy.", "mostly English, Igbo greetings"], ["Please call your mother tonight, she has been waiting to hear from you.", "Igbo sentence with a few English words"], ["Welcome to our shop. Everything is half price this weekend, so tell your friends.", "Yoruba greeting, the rest in English"], ["Good evening. The rain will be heavy tomorrow, so carry an umbrella.", "half Hausa, half English"], ["My brother, the work is too much today, but we will finish it.", "Nigerian Pidgin mixed with Igbo"], ["Thank you for coming. We will start the program in five minutes.", "English with a Yoruba closing"], ] res = [] for t, r in tests: o, info = fe(t, r) res.append(dict(text=t, request=r, out=o, ok=info["ok"], tries=len(info["tries"]), log=info["tries"])) print(json.dumps(res[-1], ensure_ascii=False), flush=True) json.dump(res, open(os.environ.get("CS_OUT", "cs_test.json"), "w"), ensure_ascii=False, indent=1)