From b3aaaa05231615b4a1561bac6152c46494b54d09 Mon Sep 17 00:00:00 2001 From: Wannaphong Date: Mon, 26 Oct 2026 16:51:13 +0700 Subject: [PATCH] Add custom dict --- README.md | 43 ++++ demo.py | 7 + pythaitts/__init__.py | 116 +++++++++- pythaitts/preprocess.py | 113 ++++++++- pythaitts/pretrained/fastthaig2p/__init__.py | 2 + pythaitts/pretrained/fastthaig2p/g2p.py | 102 ++++++++- .../pretrained/fastthaig2p/normalizer.py | 35 ++- pythaitts/pretrained/fastthaig2p/tokenizer.py | 57 ++++- pythaitts/pretrained/fastthaig2p/tts.py | 56 ++++- pythaitts/realtime.py | 11 +- tests/test_custom_dict.py | 214 ++++++++++++++++++ 11 files changed, 715 insertions(+), 41 deletions(-) create mode 100644 tests/test_custom_dict.py diff --git a/README.md b/README.md index 56a0237..1be57a2 100644 --- a/README.md +++ b/README.md @@ -142,6 +142,49 @@ print(eng("project")) # Output: /proː˧.t͡ɕek̚˨˩ ``` +### Custom Word Pronunciation Dictionary + +PyThaiTTS supports custom word dictionaries to define custom Thai pronunciation (as Thai text reading or IPA phonemes) for specific words, acronyms, or names. + +Dictionaries can be passed as Python dicts, JSON files, or delimited text files (TSV, CSV, or colon-separated): + +```python +from pythaitts import TTS + +# 1. Provide custom Thai pronunciation (hyphenated or plain text) +custom_dict = { + "เพชรบุรี": "เพ็ด-ชะ-บุ-รี", + "สถิติ": "สะ-ถิ-ติ", + "ChatGPT": "แชตจีพีที", +} + +tts = TTS(custom_dict=custom_dict) +tts.tts("ยินดีต้อนรับสู่เพชรบุรี", filename="custom.wav") + +# 2. Or pass per-call or with IPA phonemes directly +tts.tts( + "ผมไปเพชรบุรี", + custom_dict={"เพชรบุรี": "/pʰet̚˨˩.t͡ɕʰa˦˥.bu˨˩.riː˧/"}, + filename="ipa.wav", +) + +# 3. Add or update custom words dynamically +tts.add_custom_dict({"ก.พ.": "กอ-พอ"}) + +# 4. Load from a JSON or TSV file +tts.set_custom_dict("my_words.json") +``` + +You can also use custom dictionaries with G2P directly: + +```python +from pythaitts.pretrained.fastthaig2p import G2P + +g2p = G2P(custom_dict={"เพชรบุรี": "เพ็ด-ชะ-บุ-รี"}) +print(g2p.convert("ผมไปเพชรบุรี")) +``` + + ### Text Preprocessing PyThaiTTS includes automatic text preprocessing to improve TTS quality: diff --git a/demo.py b/demo.py index bf202c6..8c16971 100644 --- a/demo.py +++ b/demo.py @@ -41,6 +41,13 @@ def main(): print(f"✓ Waveform generated successfully!") print(f"Waveform shape: {waveform.shape}") print() + + # Demonstrate custom word pronunciation dictionary + print("Demonstrating custom word pronunciation dictionary...") + custom_dict = {"พายไทยทีทีเอส": "พาย-ไทย-ที-ที-เอส"} + waveform_custom = tts.tts(text, custom_dict=custom_dict, return_type="waveform") + print(f"✓ Custom pronunciation waveform generated! Shape: {waveform_custom.shape}") + print() print("=" * 60) print("Demo completed successfully!") diff --git a/pythaitts/__init__.py b/pythaitts/__init__.py index 5710e7f..4b838fd 100644 --- a/pythaitts/__init__.py +++ b/pythaitts/__init__.py @@ -4,9 +4,15 @@ """ __version__ = "0.8.0" -from typing import Union, Iterable, Iterator +from typing import Union, Iterable, Iterator, Optional, Dict, Any +from pathlib import Path -from pythaitts.preprocess import preprocess_text, num_to_thai, expand_maiyamok +from pythaitts.preprocess import ( + preprocess_text, + num_to_thai, + expand_maiyamok, + load_custom_dict, +) from pythaitts.pretrained.fastthaig2p.en_g2p import ( EnglishG2P, en_g2p, @@ -24,12 +30,21 @@ class TTS: - def __init__(self, pretrained="fastthaig2p", mode="last_checkpoint", version="1.0", device:str="cpu", **kwargs) -> None: + def __init__( + self, + pretrained="fastthaig2p", + mode="last_checkpoint", + version="1.0", + device: str = "cpu", + custom_dict: Optional[Union[Dict[str, str], str, Path, Iterable[Any]]] = None, + **kwargs, + ) -> None: """ :param str pretrained: TTS pretrained (fastthaig2p, lunarlist_onnx, khanomtan, lunarlist, vachana) :param str mode: pretrained mode (fastthaig2p, lunarlist_onnx, and vachana don't support) :param str version: model version (default is 1.0 or 1.1) :param str device: device for running model. (fastthaig2p, lunarlist_onnx, and vachana support CPU only.) + :param custom_dict: Optional custom word dictionary for custom Thai pronunciation in words **Options for mode** * *last_checkpoint* (default) - last checkpoint of model @@ -54,6 +69,9 @@ def __init__(self, pretrained="fastthaig2p", mode="last_checkpoint", version="1. self.pretrained = pretrained self.mode = mode self.device = device + self.custom_dict = custom_dict + if custom_dict is not None: + kwargs["custom_dict"] = custom_dict self.load_pretrained(version=version, **kwargs) def load_pretrained(self, version="1.0", **kwargs): @@ -80,7 +98,40 @@ def load_pretrained(self, version="1.0", **kwargs): "PyThaiTTS doesn't support %s pretrained." % self.pretrained ) - def tts(self, text: str, speaker_idx: str = "thai_som", language_idx: str = "th-th", return_type: str = "file", filename: str = None, preprocess: bool = True, **kwargs): + def add_custom_dict( + self, custom_dict: Union[Dict[str, str], str, Path, Iterable[Any]] + ) -> None: + """Add custom word pronunciations to the TTS model.""" + loaded = load_custom_dict(custom_dict) + if self.custom_dict is None: + self.custom_dict = {} + elif isinstance(self.custom_dict, (str, Path)): + self.custom_dict = load_custom_dict(self.custom_dict) + elif not isinstance(self.custom_dict, dict): + self.custom_dict = load_custom_dict(self.custom_dict) + self.custom_dict.update(loaded) + if hasattr(self.model, "add_custom_dict"): + self.model.add_custom_dict(loaded) + + def set_custom_dict( + self, custom_dict: Union[Dict[str, str], str, Path, Iterable[Any]] + ) -> None: + """Replace custom word pronunciations with new dictionary.""" + self.custom_dict = custom_dict + if hasattr(self.model, "set_custom_dict"): + self.model.set_custom_dict(custom_dict) + + def tts( + self, + text: str, + speaker_idx: str = "thai_som", + language_idx: str = "th-th", + return_type: str = "file", + filename: str = None, + preprocess: bool = True, + custom_dict: Optional[Union[Dict[str, str], str, Path, Iterable[Any]]] = None, + **kwargs, + ): """ speech synthesis @@ -90,23 +141,56 @@ def tts(self, text: str, speaker_idx: str = "thai_som", language_idx: str = "th- :param str return_type: return type (default is file) :param str filename: path filename for save wav file if return_type is file. :param bool preprocess: whether to preprocess text (convert numbers to Thai text and expand ๆ). Default is True. + :param custom_dict: Optional custom word dictionary for custom Thai pronunciation in words :param kwargs: Additional parameters passed to the underlying model. """ + effective_custom_dict = ( + custom_dict if custom_dict is not None else getattr(self, "custom_dict", None) + ) + # Preprocess text if requested if preprocess: from pythaitts.preprocess import preprocess_text - text = preprocess_text(text) + + if self.pretrained not in ("fastthaig2p", "FastThaiG2P") and effective_custom_dict: + text = preprocess_text(text, custom_dict=effective_custom_dict) + else: + text = preprocess_text(text) + elif self.pretrained not in ("fastthaig2p", "FastThaiG2P") and effective_custom_dict: + from pythaitts.preprocess import preprocess_text + + text = preprocess_text( + text, + expand_numbers=False, + expand_maiyamok_char=False, + custom_dict=effective_custom_dict, + ) if self.pretrained == "lunarlist" or self.pretrained == "lunarlist_onnx": - return self.model(text=text,return_type=return_type,filename=filename) + return self.model(text=text, return_type=return_type, filename=filename) elif self.pretrained == "vachana": if speaker_idx in ("thai_som", "Linda", None): speaker_idx = "th_f_1" - return self.model(text=text,speaker_idx=speaker_idx,return_type=return_type,filename=filename, **kwargs) + return self.model( + text=text, + speaker_idx=speaker_idx, + return_type=return_type, + filename=filename, + **kwargs, + ) elif self.pretrained in ("fastthaig2p", "FastThaiG2P"): if speaker_idx in ("Linda", None): speaker_idx = "thai_som" - return self.model(text=text,speaker_idx=speaker_idx,return_type=return_type,filename=filename, **kwargs) + call_kwargs = dict(kwargs) + if effective_custom_dict is not None: + call_kwargs["custom_dict"] = effective_custom_dict + return self.model( + text=text, + speaker_idx=speaker_idx, + return_type=return_type, + filename=filename, + **call_kwargs, + ) if speaker_idx in ("thai_som", None): speaker_idx = "Linda" return self.model( @@ -114,7 +198,7 @@ def tts(self, text: str, speaker_idx: str = "thai_som", language_idx: str = "th- speaker_idx=speaker_idx, language_idx=language_idx, return_type=return_type, - filename=filename + filename=filename, ) def stream( @@ -125,6 +209,7 @@ def stream( play: bool = False, preprocess: bool = True, max_phonemes: int = 400, + custom_dict: Optional[Union[Dict[str, str], str, Path, Iterable[Any]]] = None, **kwargs, ): """ @@ -136,9 +221,14 @@ def stream( :param bool play: Whether to play audio chunks in real-time to speakers :param bool preprocess: Whether to preprocess text (numbers to words, ๆ) :param int max_phonemes: Maximum phonemes per synthesized chunk + :param custom_dict: Optional custom word dictionary for custom Thai pronunciation in words :param kwargs: Additional parameters passed to the model :return: Generator yielding audio chunks """ + effective_custom_dict = ( + custom_dict if custom_dict is not None else getattr(self, "custom_dict", None) + ) + if self.pretrained in ("fastthaig2p", "FastThaiG2P"): if speaker_idx in ("Linda", None): speaker_idx = "thai_som" @@ -149,6 +239,7 @@ def stream( play=play, preprocess=preprocess, max_phonemes=max_phonemes, + custom_dict=effective_custom_dict, **kwargs, ) else: @@ -156,13 +247,17 @@ def stream( def _gen(): for chunk in stream_text_to_chunks( - text, max_phonemes=max_phonemes, preprocess=preprocess + text, + max_phonemes=max_phonemes, + preprocess=preprocess, + custom_dict=effective_custom_dict, ): yield self.tts( text=chunk, speaker_idx=speaker_idx, return_type=return_type, preprocess=False, + custom_dict=effective_custom_dict, **kwargs, ) @@ -182,6 +277,7 @@ def _gen(): "en_to_ipa", "en_ipa_to_thai_ipa", "preprocess_text", + "load_custom_dict", "num_to_thai", "expand_maiyamok", "chunk_text", diff --git a/pythaitts/preprocess.py b/pythaitts/preprocess.py index 7ac1307..c126dc6 100644 --- a/pythaitts/preprocess.py +++ b/pythaitts/preprocess.py @@ -5,13 +5,99 @@ This module provides text preprocessing functions for Thai Text-to-Speech, including number to Thai text conversion and handling of Thai repetition character (ๆ). """ +import json import re +from pathlib import Path +from typing import Any, Dict, Iterable, Optional, Union # Thai number words THAI_ONES = ["", "หนึ่ง", "สอง", "สาม", "สี่", "ห้า", "หก", "เจ็ด", "แปด", "เก้า"] THAI_TENS = ["", "สิบ", "ยี่สิบ", "สามสิบ", "สี่สิบ", "ห้าสิบ", "หกสิบ", "เจ็ดสิบ", "แปดสิบ", "เก้าสิบ"] +_IPA_TONES = set("˧˨˩˥˩˦˥˩˩˦↓↑→↘↗") +_IPA_CHARS = set("ʔʰ̚ŋɕɔɛɯɤː̯") + + +def _is_ipa(pron: str) -> bool: + """Check if a pronunciation string is already an IPA phoneme representation.""" + pron = pron.strip() + if not pron: + return False + if pron.startswith("/") and pron.endswith("/"): + return True + if any(c in _IPA_TONES for c in pron): + return True + if any(c in _IPA_CHARS for c in pron): + return True + return False + + +def load_custom_dict( + custom_dict: Optional[Union[Dict[str, str], str, Path, Iterable[Any]]] +) -> Dict[str, str]: + """ + Load a custom word pronunciation dictionary. + + Supports: + - Python dict: ``{"word": "pronunciation", ...}`` + - Path or str to a JSON file: ``{"word": "pronunciation", ...}`` + - Path or str to a delimited text file (TSV, CSV, or colon): + ``wordpronunciation``, ``word,pronunciation``, or ``word: pronunciation`` + - Iterable of ``(word, pronunciation)`` pairs or tuples. + + :param custom_dict: Custom dictionary data or file path + :return: Dictionary mapping word to custom pronunciation + :rtype: Dict[str, str] + """ + if custom_dict is None: + return {} + if isinstance(custom_dict, dict): + return {str(k).strip(): str(v).strip() for k, v in custom_dict.items() if str(k).strip()} + + if isinstance(custom_dict, (str, Path)): + p = Path(custom_dict) + if not p.exists(): + raise FileNotFoundError(f"Custom dictionary file not found: {custom_dict}") + if p.suffix.lower() == ".json": + with open(p, "r", encoding="utf-8") as f: + data = json.load(f) + if not isinstance(data, dict): + raise ValueError("JSON custom dictionary must contain a key-value object.") + return {str(k).strip(): str(v).strip() for k, v in data.items() if str(k).strip()} + else: + result = {} + with open(p, "r", encoding="utf-8") as f: + for line in f: + line = line.strip() + if not line or line.startswith("#"): + continue + if "\t" in line: + parts = line.split("\t", 1) + elif "," in line: + parts = line.split(",", 1) + elif ":" in line: + parts = line.split(":", 1) + else: + continue + k = parts[0].strip() + v = parts[1].strip() + if k: + result[k] = v + return result + + if isinstance(custom_dict, Iterable): + result = {} + for item in custom_dict: + if isinstance(item, (list, tuple)) and len(item) >= 2: + k = str(item[0]).strip() + v = str(item[1]).strip() + if k: + result[k] = v + return result + + raise TypeError(f"Unsupported custom_dict type: {type(custom_dict).__name__}") + def _num_to_thai_under_hundred(num: int) -> str: """ @@ -217,13 +303,20 @@ def expand_maiyamok(text: str) -> str: return ''.join(result) -def preprocess_text(text: str, expand_numbers: bool = True, expand_maiyamok_char: bool = True) -> str: +def preprocess_text( + text: str, + expand_numbers: bool = True, + expand_maiyamok_char: bool = True, + custom_dict: Optional[Union[Dict[str, str], str, Path, Iterable[Any]]] = None, +) -> str: """ - Preprocess Thai text for TTS by converting numbers to text and expanding ๆ. + Preprocess Thai text for TTS by converting numbers to text, expanding ๆ, + and optionally applying custom word pronunciation substitutions. :param str text: Input text to preprocess :param bool expand_numbers: Whether to convert numbers to Thai text (default: True) :param bool expand_maiyamok_char: Whether to expand ๆ character (default: True) + :param custom_dict: Optional custom word dictionary for custom Thai pronunciation in words :return: Preprocessed text :rtype: str @@ -234,9 +327,25 @@ def preprocess_text(text: str, expand_numbers: bool = True, expand_maiyamok_char 'ดีดี' >>> preprocess_text("มี 5 คนๆ") 'มี ห้า คนคน' + >>> preprocess_text("ไปเพชรบุรี", custom_dict={"เพชรบุรี": "เพ็ดชะบุรี"}) + 'ไปเพ็ดชะบุรี' """ result = text + # Apply custom word dictionary replacements if provided + if custom_dict is not None: + cd = load_custom_dict(custom_dict) + # Filter for text pronunciations (exclude IPA strings like /.../) + text_replacements = { + k: v for k, v in cd.items() + if not _is_ipa(v) + } + if text_replacements: + pattern = re.compile( + "|".join(re.escape(k) for k in sorted(text_replacements.keys(), key=len, reverse=True)) + ) + result = pattern.sub(lambda m: text_replacements[m.group(0)], result) + # Expand mai yamok (ๆ) first if expand_maiyamok_char: result = expand_maiyamok(result) diff --git a/pythaitts/pretrained/fastthaig2p/__init__.py b/pythaitts/pretrained/fastthaig2p/__init__.py index 818ed26..78632a1 100644 --- a/pythaitts/pretrained/fastthaig2p/__init__.py +++ b/pythaitts/pretrained/fastthaig2p/__init__.py @@ -23,6 +23,7 @@ from .kokoro import ipa_to_kokoro from .normalizer import normalize from .tokenizer import Tokenizer +from pythaitts.preprocess import load_custom_dict __all__ = [ "G2P", @@ -33,6 +34,7 @@ "Tokenizer", "normalize", "ipa_to_kokoro", + "load_custom_dict", "TTS", "FastThaiG2P", "FastThaiG2PEngine", diff --git a/pythaitts/pretrained/fastthaig2p/g2p.py b/pythaitts/pretrained/fastthaig2p/g2p.py index befa4d9..e02d7a7 100644 --- a/pythaitts/pretrained/fastthaig2p/g2p.py +++ b/pythaitts/pretrained/fastthaig2p/g2p.py @@ -20,7 +20,9 @@ import json import re from pathlib import Path +from typing import Any, Dict, Iterable, Optional, Union +from pythaitts.preprocess import _is_ipa, load_custom_dict from ._data import data_path from .en_g2p import en_g2p from .fallback import fallback_g2p @@ -39,18 +41,61 @@ def __init__( dict_path: Path | str = DEFAULT_DICT_PATH, ipa_path: Path | str = DEFAULT_IPA_PATH, use_espeak: bool = False, + custom_dict: Optional[Union[Dict[str, str], str, Path, Iterable[Any]]] = None, ): self._tokenizer = Tokenizer(dict_path) ipa_path = Path(ipa_path) self._ipa = json.loads(ipa_path.read_text(encoding="utf-8")) self.use_espeak = use_espeak + self._custom_dict: Dict[str, str] = {} + self._custom_ipa: Dict[str, str] = {} + if custom_dict is not None: + self.add_custom_dict(custom_dict) - def convert(self, text: str) -> str: - """Convert Thai (and code-switched English) text to IPA phonemes. + def add_custom_dict( + self, custom_dict: Union[Dict[str, str], str, Path, Iterable[Any]] + ) -> None: + """Add custom word pronunciations to the G2P engine.""" + loaded = load_custom_dict(custom_dict) + for word, pron in loaded.items(): + self._custom_dict[word] = pron + self._custom_ipa[word] = self._pronunciation_to_ipa(pron) + self._tokenizer.add_words(loaded.keys()) - Pipeline: normalize → tokenize → dict lookup → English G2P / fallback G2P → join. - """ - normalized = normalize(text) + def set_custom_dict( + self, custom_dict: Union[Dict[str, str], str, Path, Iterable[Any]] + ) -> None: + """Replace all custom word pronunciations with a new dictionary.""" + self.clear_custom_dict() + if custom_dict is not None: + self.add_custom_dict(custom_dict) + + def remove_custom_word(self, word: str) -> None: + """Remove a word from the custom dictionary.""" + word = str(word).strip() + if word in self._custom_dict: + del self._custom_dict[word] + if word in self._custom_ipa: + del self._custom_ipa[word] + self._tokenizer.remove_words(word) + + def clear_custom_dict(self) -> None: + """Clear all custom word pronunciations.""" + keys_to_remove = list(self._custom_dict.keys()) + self._custom_dict.clear() + self._custom_ipa.clear() + self._tokenizer.remove_words(keys_to_remove) + + def _pronunciation_to_ipa(self, pron: str) -> str: + """Convert a custom pronunciation string (Thai text or IPA) to IPA phonemes.""" + pron = pron.strip() + if not pron: + return "" + if _is_ipa(pron): + return pron + + clean_pron = pron.replace("-", " ") + normalized = normalize(clean_pron) tokens = self._tokenizer.tokenize(normalized) phonemes = [] for token in tokens: @@ -64,3 +109,50 @@ def convert(self, text: str) -> str: else: phonemes.append(fallback_g2p(token)) return " ".join(phonemes) + + def convert( + self, + text: str, + custom_dict: Optional[Union[Dict[str, str], str, Path, Iterable[Any]]] = None, + ) -> str: + """Convert Thai (and code-switched English) text to IPA phonemes. + + Pipeline: normalize → tokenize → dict lookup → English G2P / fallback G2P → join. + Supports custom_dict for custom Thai pronunciation in words. + """ + temp_keys = [] + if custom_dict is not None: + temp_loaded = load_custom_dict(custom_dict) + effective_custom_dict = {**self._custom_dict, **temp_loaded} + effective_custom_ipa = { + **self._custom_ipa, + **{k: self._pronunciation_to_ipa(v) for k, v in temp_loaded.items()}, + } + temp_keys = [k for k in temp_loaded.keys() if k not in self._custom_dict] + if temp_keys: + self._tokenizer.add_words(temp_keys) + else: + effective_custom_dict = self._custom_dict + effective_custom_ipa = self._custom_ipa + + try: + normalized = normalize(text, custom_dict=effective_custom_dict) + tokens = self._tokenizer.tokenize(normalized) + phonemes = [] + for token in tokens: + if token.strip() == "": + continue + if token in effective_custom_ipa: + phonemes.append(effective_custom_ipa[token]) + else: + ipa = self._ipa.get(token) + if ipa: + phonemes.append(ipa) + elif _EN_TOKEN_RE.match(token): + phonemes.append(en_g2p(token, use_espeak=self.use_espeak)) + else: + phonemes.append(fallback_g2p(token)) + return " ".join(phonemes) + finally: + if temp_keys: + self._tokenizer.remove_words(temp_keys) diff --git a/pythaitts/pretrained/fastthaig2p/normalizer.py b/pythaitts/pretrained/fastthaig2p/normalizer.py index 7fc10df..37a57c1 100644 --- a/pythaitts/pretrained/fastthaig2p/normalizer.py +++ b/pythaitts/pretrained/fastthaig2p/normalizer.py @@ -442,8 +442,14 @@ def _number_or_id_to_thai(match: re.Match) -> str: return _decimal_to_thai(text) -def _expand_abbreviations(text: str) -> str: - return _ABBREV_PATTERN.sub(lambda m: ABBREVIATIONS[m.group(0)], text) +def _expand_abbreviations(text: str, custom_dict: Optional[dict] = None) -> str: + def _repl(m: re.Match) -> str: + word = m.group(0) + if custom_dict and (word in custom_dict or word.strip() in custom_dict): + return word + return ABBREVIATIONS[word] + + return _ABBREV_PATTERN.sub(_repl, text) # Acronyms / initialisms that are always spelled letter-by-letter in Thai @@ -568,12 +574,14 @@ def _should_spell_latin_token(word: str, full_text: str, start: int, end: int) - return False -def _spell_latin_residue(text: str, spell_all: bool = False) -> str: +def _spell_latin_residue(text: str, spell_all: bool = False, custom_dict: Optional[dict] = None) -> str: """Spell out Latin acronyms/codes Thai-style (A → เอ, B → บี) while preserving pronounceable English words for English-to-Thai-IPA G2P. """ def _repl(m: re.Match) -> str: word = m.group(0) + if custom_dict and (word in custom_dict or word.lower() in custom_dict): + return word if spell_all or _should_spell_latin_token(word, text, m.start(), m.end()): return "".join(LETTER_TO_THAI[c.upper()] for c in word) return word @@ -585,8 +593,15 @@ def _expand_symbols(text: str) -> str: return _SYMBOL_PATTERN.sub(lambda m: SYMBOLS[m.group(0)], text) -def _expand_english_abbrevs(text: str) -> str: - return _ENGLISH_ABBREV_PATTERN.sub(lambda m: ENGLISH_ABBREVS[m.group(0).lower()], text) +def _expand_english_abbrevs(text: str, custom_dict: Optional[dict] = None) -> str: + def _repl(m: re.Match) -> str: + word = m.group(0) + lower_word = word.lower() + if custom_dict and (word in custom_dict or lower_word in custom_dict): + return word + return ENGLISH_ABBREVS[lower_word] + + return _ENGLISH_ABBREV_PATTERN.sub(_repl, text) def _expand_units(text: str) -> str: @@ -605,7 +620,7 @@ def _thai_numerals_to_arabic(text: str) -> str: return text -def normalize(text: str, spell_latin: bool = False) -> str: +def normalize(text: str, spell_latin: bool = False, custom_dict: Optional[dict] = None) -> str: """Normalize text for TTS: convert numbers, abbreviations, and special characters to speakable Thai words. @@ -613,12 +628,14 @@ def normalize(text: str, spell_latin: bool = False) -> str: :param bool spell_latin: If True, spell all remaining Latin words letter-by-letter. If False (default), only acronyms and single letters are spelled out, leaving English words intact for English-to-Thai-IPA G2P. + :param dict custom_dict: Optional custom word dictionary to protect custom words + from being unintentionally expanded by abbreviations or normalizer rules. """ text = _expand_maiyamok(text) text = _thai_numerals_to_arabic(text) text = _EMAIL_PATTERN.sub(_email_to_thai, text) text = _expand_en_contractions(text) - text = _expand_english_abbrevs(text) + text = _expand_english_abbrevs(text, custom_dict=custom_dict) text = _expand_units(text) text = _expand_symbols(text) text = _TIME_PATTERN.sub(_time_to_thai, text) @@ -626,8 +643,8 @@ def normalize(text: str, spell_latin: bool = False) -> str: text = _ALPHANUM_ID_PATTERN.sub(_alphanum_to_thai, text) text = _COMMA_NUMBER_PATTERN.sub(_strip_commas_and_convert, text) text = _NUMBER_PATTERN.sub(_number_or_id_to_thai, text) - text = _expand_abbreviations(text) - text = _spell_latin_residue(text, spell_all=spell_latin) + text = _expand_abbreviations(text, custom_dict=custom_dict) + text = _spell_latin_residue(text, spell_all=spell_latin, custom_dict=custom_dict) # leftover hyphens (ชื่อ-นามสกุล) — connectors, not speakable tokens text = text.replace("-", " ") return text diff --git a/pythaitts/pretrained/fastthaig2p/tokenizer.py b/pythaitts/pretrained/fastthaig2p/tokenizer.py index 202ae5c..0b2ddab 100644 --- a/pythaitts/pretrained/fastthaig2p/tokenizer.py +++ b/pythaitts/pretrained/fastthaig2p/tokenizer.py @@ -17,6 +17,7 @@ """ from __future__ import annotations +from typing import Iterable, Optional, Union from pathlib import Path from pythainlp.tokenize import Tokenizer as _PyThaiTokenizer @@ -27,16 +28,66 @@ class Tokenizer: - def __init__(self, dict_path: Path | str = DEFAULT_DICT_PATH): + def __init__( + self, + dict_path: Path | str = DEFAULT_DICT_PATH, + custom_dict: Optional[Union[Iterable[str], dict, str, Path]] = None, + ): dict_path = Path(dict_path) words = set(dict_path.read_text(encoding="utf-8").splitlines()) words.discard("") - self._dict = words - self._engine = _PyThaiTokenizer(custom_dict=words, engine="newmm") + self._base_dict = words + self._custom_words: set[str] = set() + self._dict = set(self._base_dict) + if custom_dict: + self.add_words(custom_dict) + else: + self._engine = _PyThaiTokenizer(custom_dict=self._dict, engine="newmm") @property def vocab(self) -> set[str]: return self._dict + def add_words(self, words: Union[Iterable[str], dict, str, Path]) -> None: + """Add custom words to the tokenizer vocabulary.""" + if isinstance(words, (str, Path)): + p = Path(words) + if p.exists(): + if p.suffix.lower() == ".json": + import json + + loaded = json.loads(p.read_text(encoding="utf-8")) + words = loaded.keys() if isinstance(loaded, dict) else loaded + else: + words = [ + line.split("\t")[0].split(",")[0].split(":")[0].strip() + for line in p.read_text(encoding="utf-8").splitlines() + if line.strip() and not line.strip().startswith("#") + ] + else: + words = [str(words)] + elif isinstance(words, dict): + words = words.keys() + + for w in words: + w_str = str(w).strip() + if w_str: + self._custom_words.add(w_str) + self._dict = self._base_dict | self._custom_words + self._engine = _PyThaiTokenizer(custom_dict=self._dict, engine="newmm") + + def remove_words(self, words: Union[Iterable[str], dict, str]) -> None: + """Remove custom words from the tokenizer vocabulary.""" + if isinstance(words, str): + words = [words] + elif isinstance(words, dict): + words = words.keys() + + for w in words: + w_str = str(w).strip() + self._custom_words.discard(w_str) + self._dict = self._base_dict | self._custom_words + self._engine = _PyThaiTokenizer(custom_dict=self._dict, engine="newmm") + def tokenize(self, text: str) -> list[str]: return self._engine.word_tokenize(text) diff --git a/pythaitts/pretrained/fastthaig2p/tts.py b/pythaitts/pretrained/fastthaig2p/tts.py index 6c2b8e2..d74d8fc 100644 --- a/pythaitts/pretrained/fastthaig2p/tts.py +++ b/pythaitts/pretrained/fastthaig2p/tts.py @@ -130,6 +130,7 @@ def __init__( device: Optional[str] = None, speed: float = 1.0, intra_op_threads: int = 0, + custom_dict: Optional[Union[dict[str, str], str, Path, Iterable]] = None, ): if model_path is None: model_path = _default_asset("model") @@ -176,15 +177,36 @@ def __init__( ) self._device = device - self._g2p = G2P() + self._g2p = G2P(custom_dict=custom_dict) self.sample_rate = 24000 self.speed = speed - def _text_to_phonemes(self, text: str) -> str: - return ipa_to_kokoro(self._g2p.convert(text)) + def add_custom_dict( + self, custom_dict: Union[dict[str, str], str, Path, Iterable] + ) -> None: + """Add custom word pronunciations to the underlying G2P engine.""" + self._g2p.add_custom_dict(custom_dict) - def generate(self, text: str) -> "np.ndarray": + def set_custom_dict( + self, custom_dict: Union[dict[str, str], str, Path, Iterable] + ) -> None: + """Replace all custom word pronunciations in the underlying G2P engine.""" + self._g2p.set_custom_dict(custom_dict) + + def _text_to_phonemes( + self, + text: str, + custom_dict: Optional[Union[dict[str, str], str, Path, Iterable]] = None, + ) -> str: + return ipa_to_kokoro(self._g2p.convert(text, custom_dict=custom_dict)) + + def generate( + self, + text: str, + custom_dict: Optional[Union[dict[str, str], str, Path, Iterable]] = None, + ) -> "np.ndarray": """Convert Thai text to float32 audio array (24 kHz mono).""" + phonemes = self._text_to_phonemes(text, custom_dict=custom_dict) phonemes = self._text_to_phonemes(text) if not phonemes: return np.array([], dtype=np.float32) @@ -225,10 +247,15 @@ def generate(self, text: str) -> "np.ndarray": audio.cpu().numpy(), durations.cpu().numpy() ) - def synthesize(self, text: str, output_path: str | Path) -> Path: + def synthesize( + self, + text: str, + output_path: str | Path, + custom_dict: Optional[Union[dict[str, str], str, Path, Iterable]] = None, + ) -> Path: """Convert Thai text to WAV file.""" output_path = Path(output_path) - audio = self.generate(text) + audio = self.generate(text, custom_dict=custom_dict) audio_int16 = (np.clip(audio, -1.0, 1.0) * 32767).astype(np.int16) with wave.open(str(output_path), "wb") as wf: wf.setnchannels(1) @@ -243,6 +270,7 @@ def __call__( speaker_idx: str = "thai_som", return_type: str = "file", filename: Optional[str] = None, + custom_dict: Optional[Union[dict[str, str], str, Path, Iterable]] = None, **kwargs, ): """Generate speech from text using FastThaiG2P TTS. @@ -251,6 +279,7 @@ def __call__( :param str speaker_idx: Voice to use (default: "thai_som" or path to custom voicepack) :param str return_type: Return type ("file" or "waveform") :param str filename: Output filename for the generated audio + :param custom_dict: Optional custom word dictionary for custom Thai pronunciation in words :param kwargs: Additional parameters (e.g., speed) :return: File path if return_type is "file", otherwise audio waveform data (numpy.ndarray) """ @@ -266,6 +295,8 @@ def __call__( self.speed = kwargs["speed"] if return_type == "waveform": + if custom_dict is not None: + return self.generate(text, custom_dict=custom_dict) return self.generate(text) else: import tempfile @@ -274,7 +305,10 @@ def __call__( with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as fp: filename = fp.name - self.synthesize(text, filename) + if custom_dict is not None: + self.synthesize(text, filename, custom_dict=custom_dict) + else: + self.synthesize(text, filename) return str(filename) def stream( @@ -285,6 +319,7 @@ def stream( play: bool = False, preprocess: bool = True, max_phonemes: int = 400, + custom_dict: Optional[Union[dict[str, str], str, Path, Iterable]] = None, **kwargs, ): """Stream Thai speech generation chunk by chunk in real-time. @@ -295,6 +330,7 @@ def stream( :param bool play: Whether to play audio chunks in real-time to speakers :param bool preprocess: Whether to preprocess text (numbers to words, ๆ) :param int max_phonemes: Maximum phonemes per synthesized chunk (default: 400) + :param custom_dict: Optional custom word dictionary for custom Thai pronunciation in words :param kwargs: Additional parameters (e.g., speed) :yield: Audio chunk (np.ndarray float32 waveform, int16 bytes, or wav file path) """ @@ -319,10 +355,10 @@ def stream( text, max_phonemes=max_phonemes, preprocess=preprocess, - g2p_converter=self._text_to_phonemes, + g2p_converter=lambda t: self._text_to_phonemes(t, custom_dict=custom_dict), ) for chunk in chunks: - audio = self.generate(chunk) + audio = self.generate(chunk, custom_dict=custom_dict) if len(audio) == 0: continue @@ -349,9 +385,9 @@ def stream( f"Unsupported return_type '{return_type}'. Choose 'waveform', 'bytes', or 'file'." ) finally: + self.speed = orig_speed if player is not None: player.close() - self.speed = orig_speed tts_stream = stream diff --git a/pythaitts/realtime.py b/pythaitts/realtime.py index 7ce5018..9b89dce 100644 --- a/pythaitts/realtime.py +++ b/pythaitts/realtime.py @@ -105,6 +105,7 @@ def chunk_text( max_phonemes: int = 400, preprocess: bool = True, g2p_converter: Optional[Callable[[str], str]] = None, + custom_dict: Optional[Union[dict, str, Path, Iterable]] = None, ) -> List[str]: """Split Thai text into chunks suitable for FastThaiG2P synthesis. @@ -115,6 +116,7 @@ def chunk_text( :param int max_phonemes: Maximum phonemes per chunk (default: 400) :param bool preprocess: Whether to preprocess text (numbers to words, ๆ) :param Callable g2p_converter: Optional custom function converting text to phonemes + :param custom_dict: Optional custom word dictionary for custom Thai pronunciation :return: List of text chunks """ if not text or not str(text).strip(): @@ -155,10 +157,10 @@ def chunk_text( if g2p_converter is None: from pythaitts.pretrained.fastthaig2p import G2P, ipa_to_kokoro - _g2p = G2P() + _g2p = G2P(custom_dict=custom_dict) def _to_phonemes(t: str) -> str: - return ipa_to_kokoro(_g2p.convert(t)) + return ipa_to_kokoro(_g2p.convert(t, custom_dict=custom_dict)) g2p_converter = _to_phonemes @@ -213,6 +215,7 @@ def stream_text_to_chunks( max_buffer_chars: int = 120, preprocess: bool = True, g2p_converter: Optional[Callable[[str], str]] = None, + custom_dict: Optional[Union[dict, str, Path, Iterable]] = None, ) -> Iterator[str]: """Yield text chunks from either a full string or a streaming iterable of text tokens. @@ -224,6 +227,7 @@ def stream_text_to_chunks( :param int max_buffer_chars: Maximum characters to buffer before breaking on word boundary :param bool preprocess: Whether to preprocess text :param Callable g2p_converter: Optional phoneme conversion function + :param custom_dict: Optional custom word dictionary for custom Thai pronunciation :yield: Text chunks ready for synthesis """ if isinstance(text_or_stream, str): @@ -232,6 +236,7 @@ def stream_text_to_chunks( max_phonemes=max_phonemes, preprocess=preprocess, g2p_converter=g2p_converter, + custom_dict=custom_dict, ): yield c return @@ -417,6 +422,7 @@ def __init__( preprocess: bool = True, max_phonemes: int = 400, debug: bool = False, + custom_dict: Optional[Union[dict, str, Path, Iterable]] = None, ): super().__init__() self.engine_name = "fastthaig2p" @@ -433,6 +439,7 @@ def __init__( config_path=config_path, device=device, speed=speed, + custom_dict=custom_dict, ) self.sample_rate = self.model.sample_rate self.set_voice(voice) diff --git a/tests/test_custom_dict.py b/tests/test_custom_dict.py new file mode 100644 index 0000000..b222868 --- /dev/null +++ b/tests/test_custom_dict.py @@ -0,0 +1,214 @@ +# -*- coding: utf-8 -*- +""" +Unit tests for Custom Word Dictionary support for Thai pronunciation +""" +import json +import os +import tempfile +import unittest +from unittest.mock import Mock, patch + +from pythaitts import TTS, preprocess_text, load_custom_dict +from pythaitts.pretrained.fastthaig2p import G2P, FastThaiG2P, ipa_to_kokoro + + +class TestCustomWordDictLoading(unittest.TestCase): + """Test load_custom_dict function with various input formats""" + + def test_load_from_dict(self): + d = {"เพชรบุรี": "เพ็ด-ชะ-บุ-รี", "สถิติ": "สะ-ถิ-ติ"} + loaded = load_custom_dict(d) + self.assertEqual(loaded, d) + + def test_load_from_json_file(self): + data = {"เพชรบุรี": "เพ็ด-ชะ-บุ-รี", "ChatGPT": "แชตจีพีที"} + with tempfile.NamedTemporaryFile("w", suffix=".json", delete=False, encoding="utf-8") as f: + json.dump(data, f, ensure_ascii=False) + filename = f.name + try: + loaded = load_custom_dict(filename) + self.assertEqual(loaded, data) + finally: + os.unlink(filename) + + def test_load_from_tsv_file(self): + content = "# Comment line\nเพชรบุรี\tเพ็ด-ชะ-บุ-รี\nสถิติ\tสะ-ถิ-ติ\n\n" + with tempfile.NamedTemporaryFile("w", suffix=".tsv", delete=False, encoding="utf-8") as f: + f.write(content) + filename = f.name + try: + loaded = load_custom_dict(filename) + self.assertEqual(loaded["เพชรบุรี"], "เพ็ด-ชะ-บุ-รี") + self.assertEqual(loaded["สถิติ"], "สะ-ถิ-ติ") + self.assertEqual(len(loaded), 2) + finally: + os.unlink(filename) + + def test_load_from_csv_file(self): + content = "เพชรบุรี,เพ็ด-ชะ-บุ-รี\nสถิติ,สะ-ถิ-ติ\n" + with tempfile.NamedTemporaryFile("w", suffix=".csv", delete=False, encoding="utf-8") as f: + f.write(content) + filename = f.name + try: + loaded = load_custom_dict(filename) + self.assertEqual(loaded["เพชรบุรี"], "เพ็ด-ชะ-บุ-รี") + self.assertEqual(loaded["สถิติ"], "สะ-ถิ-ติ") + finally: + os.unlink(filename) + + def test_load_from_colon_file(self): + content = "เพชรบุรี: เพ็ด-ชะ-บุ-รี\nสถิติ: สะ-ถิ-ติ\n" + with tempfile.NamedTemporaryFile("w", suffix=".txt", delete=False, encoding="utf-8") as f: + f.write(content) + filename = f.name + try: + loaded = load_custom_dict(filename) + self.assertEqual(loaded["เพชรบุรี"], "เพ็ด-ชะ-บุ-รี") + self.assertEqual(loaded["สถิติ"], "สะ-ถิ-ติ") + finally: + os.unlink(filename) + + def test_load_from_tuples(self): + tuples = [("เพชรบุรี", "เพ็ด-ชะ-บุ-รี"), ("สถิติ", "สะ-ถิ-ติ")] + loaded = load_custom_dict(tuples) + self.assertEqual(loaded["เพชรบุรี"], "เพ็ด-ชะ-บุ-รี") + self.assertEqual(loaded["สถิติ"], "สะ-ถิ-ติ") + + def test_load_none(self): + self.assertEqual(load_custom_dict(None), {}) + + def test_file_not_found(self): + with self.assertRaises(FileNotFoundError): + load_custom_dict("/nonexistent/path/dict.json") + + def test_invalid_type(self): + with self.assertRaises(TypeError): + load_custom_dict(12345) + + +class TestPreprocessWithCustomDict(unittest.TestCase): + """Test text preprocessing with custom word dictionary""" + + def test_preprocess_text_word_substitution(self): + text = "ฉันจะไปเพชรบุรีพรุ่งนี้" + custom_dict = {"เพชรบุรี": "เพ็ดชะบุรี"} + result = preprocess_text(text, custom_dict=custom_dict) + self.assertEqual(result, "ฉันจะไปเพ็ดชะบุรีพรุ่งนี้") + + def test_preprocess_text_with_numbers(self): + text = "ไปเพชรบุรี 2 คน" + custom_dict = {"เพชรบุรี": "เพ็ดชะบุรี"} + result = preprocess_text(text, custom_dict=custom_dict) + self.assertEqual(result, "ไปเพ็ดชะบุรี สอง คน") + + def test_preprocess_text_ignores_ipa(self): + text = "ฉันจะไปเพชรบุรี" + custom_dict = {"เพชรบุรี": "/pʰet̚˦˥.t͡ɕʰaʔ˦˥.bu˨˩.riː˧/"} + # Preprocess text should not replace with IPA string into raw Thai text + result = preprocess_text(text, custom_dict=custom_dict) + self.assertEqual(result, "ฉันจะไปเพชรบุรี") + + +class TestG2PCustomDict(unittest.TestCase): + """Test G2P phoneme conversion with custom word dictionary""" + + def test_g2p_custom_thai_pronunciation(self): + g2p = G2P(custom_dict={"เพชรบุรี": "เพ็ด-ชะ-บุ-รี"}) + result = g2p.convert("ผมไปเพชรบุรี") + # Ensure the syllables for เพ็ด, ชะ, บุ, รี are present in the IPA + self.assertIn("pʰet", result) + self.assertIn("t͡ɕʰa", result) + self.assertIn("riː", result) + + def test_g2p_custom_ipa_pronunciation(self): + custom_ipa = "/pʰet̚˨˩.t͡ɕʰa˦˥.bu˨˩.riː˧/" + g2p = G2P(custom_dict={"เพชรบุรี": custom_ipa}) + result = g2p.convert("ผมไปเพชรบุรี") + self.assertIn(custom_ipa, result) + + def test_g2p_temporary_custom_dict_in_convert(self): + g2p = G2P() + # Default conversion + default_result = g2p.convert("เพชรบุรี") + # With custom_dict in convert call + custom_ipa = "/pʰet̚˨˩.t͡ɕʰa˦˥.bu˨˩.riː˧/" + custom_result = g2p.convert("เพชรบุรี", custom_dict={"เพชรบุรี": custom_ipa}) + self.assertEqual(custom_result, custom_ipa) + # Calling convert again without custom_dict returns default + self.assertEqual(g2p.convert("เพชรบุรี"), default_result) + + def test_g2p_normalizer_protection_for_abbrev(self): + # Without custom dict, "ก.พ." is expanded to "กุมภาพันธ์" + g2p_default = G2P() + default_res = g2p_default.convert("นาย ก.พ. ทำงาน") + self.assertIn("kum", default_res) # from กุมภาพันธ์ + + # With custom dict, "ก.พ." is pronounced as "กอ-พอ" + g2p_custom = G2P(custom_dict={"ก.พ.": "กอ-พอ"}) + custom_res = g2p_custom.convert("นาย ก.พ. ทำงาน") + self.assertNotIn("kum", custom_res) + self.assertIn("kɔː", custom_res) + self.assertIn("pʰɔː", custom_res) + + def test_g2p_add_and_set_and_clear_custom_dict(self): + g2p = G2P() + custom_ipa = "/kʰam˧.pʰiː˧.seːt̚˨˩/" + g2p.add_custom_dict({"คำพิเศษ": custom_ipa}) + self.assertIn(custom_ipa, g2p.convert("นี่คือคำพิเศษ")) + + # remove custom word + g2p.remove_custom_word("คำพิเศษ") + self.assertNotIn(custom_ipa, g2p.convert("นี่คือคำพิเศษ")) + + # set custom dict replaces everything + g2p.set_custom_dict({"อีกคำ": "/ʔiːk̚˨˩.kʰam˧/"}) + self.assertIn("/ʔiːk̚˨˩.kʰam˧/", g2p.convert("อีกคำ")) + + # clear custom dict + g2p.clear_custom_dict() + self.assertEqual(len(g2p._custom_dict), 0) + self.assertEqual(len(g2p._custom_ipa), 0) + + +class TestTTSCustomDictIntegration(unittest.TestCase): + """Test high-level TTS class custom_dict integration""" + + @patch('pythaitts.pretrained.fastthaig2p.FastThaiG2P') + def test_tts_init_with_custom_dict(self, mock_fastthaig2p_cls): + mock_instance = Mock() + mock_fastthaig2p_cls.return_value = mock_instance + + custom_dict = {"เพชรบุรี": "เพ็ด-ชะ-บุ-รี"} + tts = TTS(pretrained="fastthaig2p", custom_dict=custom_dict) + mock_fastthaig2p_cls.assert_called_once() + self.assertEqual(tts.custom_dict, custom_dict) + + @patch('pythaitts.pretrained.fastthaig2p.FastThaiG2P') + def test_tts_call_with_custom_dict(self, mock_fastthaig2p_cls): + mock_instance = Mock() + mock_fastthaig2p_cls.return_value = mock_instance + + tts = TTS(pretrained="fastthaig2p") + custom_dict = {"เพชรบุรี": "เพ็ด-ชะ-บุ-รี"} + tts.tts("ไปเพชรบุรี", custom_dict=custom_dict) + + mock_instance.assert_called_once() + _, kwargs = mock_instance.call_args + self.assertEqual(kwargs.get("custom_dict"), custom_dict) + + @patch('pythaitts.pretrained.fastthaig2p.FastThaiG2P') + def test_tts_add_and_set_custom_dict(self, mock_fastthaig2p_cls): + mock_instance = Mock() + mock_fastthaig2p_cls.return_value = mock_instance + + tts = TTS(pretrained="fastthaig2p") + tts.add_custom_dict({"คำหนึ่ง": "คำ-หนึ่ง"}) + mock_instance.add_custom_dict.assert_called_once() + + tts.set_custom_dict({"คำสอง": "คำ-สอง"}) + mock_instance.set_custom_dict.assert_called_once() + + +if __name__ == "__main__": + unittest.main() +