scripts/frontend_bopomofo.py

"""zh-TW/en unified frontend for the Inflect-Nano retrain.
zh chars -> bopomofo (g2pw, Taiwan readings) -> zhuyin symbol units + tone (1-5);
en words -> arpabet (g2p_en) + stress; one sequence, per-phone language id (ZH/EN).
"""
from __future__ import annotations
import re
from g2pw import G2PWConverter
from g2p_en import G2p

# 37 standard zhuyin symbols (U+3105..U+3129)
ZHUYIN = [chr(c) for c in range(0x3105, 0x312A)]
ARPABET = ['AA','AE','AH','AO','AW','AY','B','CH','D','DH','EH','ER','EY','F','G','HH',
           'IH','IY','JH','K','L','M','N','NG','OW','OY','P','R','S','SH','T','TH',
           'UH','UW','V','W','Y','Z','ZH']
PUNCT = [',', '.', '?', '!', '…', '-', "'"]
SPECIAL = ['_blank', '_pad', 'UNK', 'SP']            # SP = inter-word/space pause
SYMBOLS = SPECIAL + ZHUYIN + ARPABET + PUNCT
SYM2ID = {s: i for i, s in enumerate(SYMBOLS)}
LANG = {'ZH': 0, 'EN': 1}                            # per-phone language id

_g2pw = None
_g2pen = None
_zh_num = None

_ZH_DIGIT = {"0":"零","1":"一","2":"二","3":"三","4":"四","5":"五","6":"六","7":"七","8":"八","9":"九"}

def _make_zh_normalizer():
    import cn2an
    def norm(text):
        # long digit runs (>=5, e.g. phone/order numbers) -> digit-by-digit zh; else cardinal
        def repl(m):
            d = m.group(0)
            if len(d) >= 5:
                return "".join(_ZH_DIGIT[c] for c in d)
            try:
                return cn2an.an2cn(d, "low")
            except Exception:
                return "".join(_ZH_DIGIT[c] for c in d)
        text = re.sub(r"\d+", repl, text)
        text = text.replace("，", ",").replace("。", ".").replace("？", "?").replace("！", "!")
        return text
    return norm

def _lazy():
    global _g2pw, _g2pen, _zh_num
    if _g2pw is None:
        _g2pw = G2PWConverter()
        _g2pen = G2p()
        _zh_num = _make_zh_normalizer()

def _split_syllable(syl: str):
    """'ㄓㄨㄢ3' -> (['ㄓ','ㄨ','ㄢ'], tone 3)."""
    tone = 0
    if syl and syl[-1].isdigit():
        tone = int(syl[-1]); syl = syl[:-1]
    units = [c for c in syl if c in SYM2ID]
    return units, tone

def text_to_phones(text: str):
    _lazy()
    text = _zh_num(text)                              # numbers->zh words, normalize punct
    bopo = _g2pw(text)[0]                             # per-char bopomofo or None
    chars = list(text)
    phones, tones, langs = [], [], []
    i = 0
    while i < len(chars):
        b = bopo[i] if i < len(bopo) else None
        ch = chars[i]
        if b is not None:                             # zh char
            units, tone = _split_syllable(b)
            for u in units:
                phones.append(u); tones.append(min(tone, 5)); langs.append(LANG['ZH'])
            i += 1
        elif re.match(r'[A-Za-z]', ch):               # English run -> g2p_en
            j = i
            while j < len(chars) and re.match(r"[A-Za-z']", chars[j]):
                j += 1
            word = ''.join(chars[i:j])
            for p in _g2pen(word):
                p = p.strip()
                if not p:
                    continue
                stress = 0
                if p[-1].isdigit():
                    stress = int(p[-1]); p = p[:-1]
                if p in SYM2ID:
                    phones.append(p); tones.append(stress); langs.append(LANG['EN'])
            phones.append('SP'); tones.append(0); langs.append(LANG['EN'])
            i = j
        else:                                         # punctuation / space / other
            if ch in PUNCT:
                phones.append(ch); tones.append(0); langs.append(LANG['ZH'])
            elif ch.strip() == '':
                if phones and phones[-1] != 'SP':
                    phones.append('SP'); tones.append(0); langs.append(LANG['ZH'])
            i += 1
    return phones, tones, langs

def text_to_ids(text: str):
    phones, tones, langs = text_to_phones(text)
    ids = [SYM2ID.get(p, SYM2ID['UNK']) for p in phones]
    return {"phones": phones, "phone_ids": ids, "tone_ids": tones, "lang_ids": langs}