tokenizer.py

import re
from typing import List
from vachana_g2p import th2ipa
from pythainlp.tokenize import word_tokenize
from pythainlp.util import normalize as pythai_normalize
    
PAD = "_"
BOS = "^"
EOS = "$"
SPACE = " "
UNK = "?"

_ARABIC_DIGITS = list("0123456789")
_PUNCT = list(" .,!?;:()\"'-…")
_IPA_THAI = ['a', 'b', 'd', 'e', 'f', 'h', 'i', 'j', 'k', 'l', 'm', 'n', 'o', 'p', 'r', 's', 't', 'u', 'w', 'ŋ', 'ɔ', 'ɕ', 'ə', 'ɛ', 'ɯ', 'ʔ', 'ʰ', 'ː']
_IPA_TONE = list('̀'+'́'+'̂'+'̌')

SYMBOLS = (
    [PAD, BOS, EOS, UNK]
    + _IPA_THAI
    + _IPA_TONE
    + _ARABIC_DIGITS
    + _PUNCT
    + ["|"] 
)

SYMBOLS = list(dict.fromkeys(SYMBOLS))

_SYM2ID = {s: i for i, s in enumerate(SYMBOLS)}
_ID2SYM = {i: s for i, s in enumerate(SYMBOLS)}

VOCAB_SIZE = len(SYMBOLS)

def chunk_text(text, max_char=1000):
    words = word_tokenize(text)

    chunks = []
    current = ""

    for word in words:
        if len(current) + len(word) <= max_char:
            current += word
        else:
            chunks.append(current)
            current = word

    if current:
        chunks.append(current)

    return chunks
    
def _normalize_text(text: str) -> str:
    text = text.strip()
    text = pythai_normalize(text)
    text = re.sub(r"\s+", " ", text)
    return text

def text_to_words(text: str) -> List[str]:
    text = _normalize_text(text)
    text = th2ipa(text) + "."
    return [text]

def tokenize(text: str, add_bos_eos: bool = True) -> List[int]:
    """Text -> list of symbol ids, with '|' inserted at word boundaries."""
    words = text_to_words(text)
    ids: List[int] = []
    if add_bos_eos:
        ids.append(_SYM2ID[BOS])
    for wi, w in enumerate(words):
        for ch in w:
            ids.append(_SYM2ID.get(ch, _SYM2ID[UNK]))
        if wi != len(words) - 1:
            ids.append(_SYM2ID["|"])
    if add_bos_eos:
        ids.append(_SYM2ID[EOS])
    return ids

def ids_to_text(ids: List[int]) -> str:
    return "".join(_ID2SYM.get(i, UNK) for i in ids if i not in
                    (_SYM2ID[PAD], _SYM2ID[BOS], _SYM2ID[EOS]))