Update tokenizer.py

tokenizer.py

@@ -1,161 +1,80 @@
 import re
-import torch
-
-CODEC_CODEBOOK_SIZE = 12800
-
-TH_CHARS = 'กขฃคฅฆงจฉชซฌญฎฏฐฑฒณดตถทธนบปผฝพฟภมยรลวศษสหฬอฮฤฦะัาำิีึืุูเแโใไๅํ็่้๊๋ฯฺๆ์ํ๎๏๚๛๐๑๒๓๔๕๖๗๘๙฿'
-EN_LOWER  = "abcdefghijklmnopqrstuvwxyz"
-EN_UPPER  = "ABCDEFGHIJKLMNOPQRSTUVWXYZ"
-DIGITS    = "0123456789"
-PUNCT     = '.,!?;:-–—…"\'()[]{}«»„"" '
-EXTRA     = "\n\t"
-
-_ALL_CHARS: list[str] = []
-_seen: set[str] = set()
-for _src in [TH_CHARS, EN_LOWER, EN_UPPER, DIGITS, PUNCT, EXTRA]:
-    for _ch in _src:
-        if _ch not in _seen:
-            _ALL_CHARS.append(_ch)
-            _seen.add(_ch)
-
-SPECIAL_TOKENS = {
-    "<pad>":             0,
-    "<start_of_text>":   1,
-    "<end_of_text>":     2,
-    "<start_of_speech>": 3,
-    "<end_of_speech>":   4,
-    "<spk_0>":           5,  # kept for compatibility, but speaker embedding is primary
-    "<spk_1>":           6,
-    "<spk_2>":           7,
-    "<spk_3>":           8,
-}
-NUM_SPECIAL_TOKENS = len(SPECIAL_TOKENS)     # 9
-
-# ── Vocab offsets ───────────────────────────────────────────────
-TEXT_CHARS       = _ALL_CHARS
-TEXT_VOCAB_SIZE  = len(TEXT_CHARS)             # ~146
-TEXT_OFFSET      = NUM_SPECIAL_TOKENS         # 9
-AUDIO_OFFSET     = TEXT_OFFSET + TEXT_VOCAB_SIZE  # 155
-NUM_AUDIO_TOKENS = CODEC_CODEBOOK_SIZE            # 12,800
-TOTAL_VOCAB_SIZE = AUDIO_OFFSET + NUM_AUDIO_TOKENS  # 12,955
-
-# Encoder needs only text vocab; decoder needs full vocab
-ENCODER_VOCAB_SIZE = AUDIO_OFFSET      # 155 (special + text)
-DECODER_VOCAB_SIZE = TOTAL_VOCAB_SIZE  # 12,955 (full)
-
-# ── Convenience IDs ─────────────────────────────────────────────
-PAD_TOKEN_ID             = SPECIAL_TOKENS["<pad>"]
-START_OF_TEXT_TOKEN_ID   = SPECIAL_TOKENS["<start_of_text>"]
-END_OF_TEXT_TOKEN_ID     = SPECIAL_TOKENS["<end_of_text>"]
-START_OF_SPEECH_TOKEN_ID = SPECIAL_TOKENS["<start_of_speech>"]
-END_OF_SPEECH_TOKEN_ID   = SPECIAL_TOKENS["<end_of_speech>"]
-SPK_0_TOKEN_ID           = SPECIAL_TOKENS["<spk_0>"]
-SPK_1_TOKEN_ID           = SPECIAL_TOKENS["<spk_1>"]
-
-# ── Helper functions ────────────────────────────────────────────
-def audio_token_id(code: int) -> int:
-    """MioCodec code → global token ID."""
-    return AUDIO_OFFSET + code
-
-def decode_audio_token(token_id: int) -> int:
-    """Global token ID → MioCodec code."""
-    return token_id - AUDIO_OFFSET
-
-def is_audio_token(token_id: int) -> bool:
-    return AUDIO_OFFSET <= token_id < AUDIO_OFFSET + NUM_AUDIO_TOKENS
-
-def is_special_token(token_id: int) -> bool:
-    return 0 <= token_id < NUM_SPECIAL_TOKENS
-
-def is_text_token(token_id: int) -> bool:
-    return TEXT_OFFSET <= token_id < AUDIO_OFFSET
-
-class TTSTokenizer:
-    def __init__(self):
-        self.char2id: dict[str, int] = {}
-        self.id2char: dict[int, str] = {}
-        for i, ch in enumerate(TEXT_CHARS):
-            tid = TEXT_OFFSET + i
-            self.char2id[ch] = tid
-            self.id2char[tid] = ch
-
-        self._special_id_to_name = {v: k for k, v in SPECIAL_TOKENS.items()}
-        self.vocab_size = TOTAL_VOCAB_SIZE
-        self.text_vocab_size = len(TEXT_CHARS)
-
-    def normalize_text(self, text: str) -> str:
-        text = re.sub(r'\s+', ' ', text).strip()
-        text = re.sub(r'[–—]', '-', text)
-        text = re.sub(r'[«»„""]', '"', text)
-        return text
-
-    def encode_text(self, text: str) -> list[int]:
-        text = self.normalize_text(text)
-        return [self.char2id[ch] for ch in text if ch in self.char2id]
-
-    def decode_text(self, ids: list[int]) -> str:
-        return "".join(self.id2char.get(t, "") for t in ids if is_text_token(t))
-
-    # ── Encoder-Decoder methods ──────────────────────────────
-
-    def build_encoder_input(self, text: str) -> torch.Tensor:
-        """
-        Encoder input: <sot> text_chars <eot>
-        No speaker token — speaker info comes from embedding.
-        """
-        text_ids = self.encode_text(text)
-        seq = text_ids
-        return torch.tensor(seq, dtype=torch.long)
-
-    def build_decoder_input(self, audio_codes: torch.Tensor) -> torch.Tensor:
-        """
-        Decoder input: <sos> [audio_codes + AUDIO_OFFSET] <eos>
-        audio_codes: raw MioCodec codes in [0, 12799]
-        """
-        seq = (
-            [START_OF_SPEECH_TOKEN_ID]
-            + (audio_codes + AUDIO_OFFSET).tolist()
-            + [END_OF_SPEECH_TOKEN_ID]
-        )
-        return torch.tensor(seq, dtype=torch.long)
-
-    def build_decoder_prefix(self) -> torch.Tensor:
-        """For inference: just <sos> to start generation."""
-        return torch.tensor([START_OF_SPEECH_TOKEN_ID], dtype=torch.long)
-
-    def extract_audio_codes(self, sequence: torch.Tensor):
-        """Extract raw MioCodec codes from a token sequence."""
-        mask = torch.tensor([is_audio_token(t.item()) for t in sequence])
-        if not mask.any():
-            return None
-        return sequence[mask] - AUDIO_OFFSET
-
-    def describe(self, seq: torch.Tensor, max_tok: int = 30) -> str:
-        parts = []
-        for t in seq[:max_tok]:
-            tid = t.item()
-            if is_special_token(tid):
-                parts.append(self._special_id_to_name.get(tid, f"<sp_{tid}>"))
-            elif is_text_token(tid):
-                ch = self.id2char.get(tid, "?")
-                parts.append(ch if ch != " " else "·")
-            elif is_audio_token(tid):
-                code = tid - AUDIO_OFFSET
-                parts.append(f"♪{code}")
-            else:
-                parts.append(f"?{tid}")
-        r = " ".join(parts)
-        if len(seq) > max_tok:
-            r += f" ... [{len(seq) - max_tok} more]"
-        return r
-
-if __name__ == "__main__":
-    tokens = TTSTokenizer()
-    text = """สวัสดีค่ะ วันนี้อยากเล่าเรื่องหนึ่งที่เราไม่เคยคิดว่าจะเปลี่ยนชีวิตเราได้ขนาดนี้ มันเป็นวันที่ธรรมดา"""
-    print(f"Text Len: {len(text)}")
-    encode = tokens.encode_text(text)
-    build_encode = tokens.build_encoder_input(text)
-    print(build_encode)
-    print(encode)
-    print(f"Encode Len: {len(encode)}")
-    print(f"Build Encode Len: {len(build_encode)}")
+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]))