| """ |
| V6 Tokenizer — char-level for Bulgarian TTS with MioCodec |
| ========================================================== |
| Same character set as V5, but adapted for: |
| - MioCodec single codebook (no interleaving) |
| - Speaker embedding (no speaker tokens in encoder input) |
| """ |
|
|
| import re |
| import torch |
| from typing import Optional |
|
|
| from config import ( |
| TEXT_CHARS, TEXT_OFFSET, AUDIO_OFFSET, |
| SPECIAL_TOKENS, NUM_SPECIAL_TOKENS, CODEC_CODEBOOK_SIZE, |
| TOTAL_VOCAB_SIZE, |
| PAD_TOKEN_ID, START_OF_TEXT_TOKEN_ID, END_OF_TEXT_TOKEN_ID, |
| START_OF_SPEECH_TOKEN_ID, END_OF_SPEECH_TOKEN_ID, |
| is_audio_token, is_special_token, is_text_token, |
| ) |
|
|
|
|
| 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)) |
|
|
| |
|
|
| 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 = [START_OF_TEXT_TOKEN_ID] + text_ids + [END_OF_TEXT_TOKEN_ID] |
| 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) -> Optional[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 |
|
|
