cs_frontend.py

"""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
(<eng>..</eng> <ibo>..</ibo>) 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: <eng>...</eng>, <ibo>...</ibo>, <yor>...</yor>, <hau>...</hau>, <pcm>...</pcm>.
- 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 = "<ibo>Ụtụtụ ọma, unu niile.</ibo> <eng>The meeting starts at ten, please be on time.</eng> <ibo>Daalụ.</ibo>"

TAG = re.compile(r"<(eng|ibo|yor|hau|pcm)>(.*?)</\1>", 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"<think>.*?</think>", "", 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)}</{m.group(1)}>", out)
                return out, dict(ok=True, tries=log)
        return f"<{fallback_lang}>{text}</{fallback_lang}>", 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)