Rename and cleanup

app.py

@@ -102,7 +102,7 @@ LANGUAGE_CONFIG = {
     },
     "zh": {
         "audio": "https://storage.googleapis.com/chatterbox-demo-samples/mtl_prompts/zh_f2.flac",
-        "text": "上个月，我们达到了一个新的里程碑. 我们的YouTube频道观看次数达到了二十亿次，这绝对令人难以置信。"
+        "text": "上个月，我们达到了一个新的里程碑。 我们的YouTube频道观看次数达到了二十亿次，这绝对令人难以置信。"
     },
 }
 

src/chatterbox/models/tokenizers/tokenizer.py

@@ -1,10 +1,9 @@
 import logging
 import json
-import re
 
 import torch
 from pathlib import Path
-from unicodedata import category
+from unicodedata import category, normalize
 from tokenizers import Tokenizer
 from huggingface_hub import hf_hub_download
 
@@ -33,7 +32,7 @@ class EnTokenizer:
         text_tokens = torch.IntTensor(text_tokens).unsqueeze(0)
         return text_tokens
 
-    def encode( self, txt: str, verbose=False):
+    def encode(self, txt: str):
         """
         clean_text > (append `lang_id`) > replace SPACE > encode text using Tokenizer
         """
@@ -46,8 +45,7 @@ class EnTokenizer:
         if isinstance(seq, torch.Tensor):
             seq = seq.cpu().numpy()
 
-        txt: str = self.tokenizer.decode(seq,
-        skip_special_tokens=False)
+        txt: str = self.tokenizer.decode(seq, skip_special_tokens=False)
         txt = txt.replace(' ', '')
         txt = txt.replace(SPACE, ' ')
         txt = txt.replace(EOT, '')
@@ -61,6 +59,7 @@ REPO_ID = "ResembleAI/chatterbox"
 # Global instances for optional dependencies
 _kakasi = None
 _dicta = None
+_russian_stresser = None
 
 
 def is_kanji(c: str) -> bool:
@@ -207,7 +206,6 @@ class ChineseCangjieConverter:
         index = str(index) if index > 0 else ""
         return code + str(index)
     
-
     
     def __call__(self, text):
         """Convert Chinese characters in text to Cangjie tokens."""
@@ -235,53 +233,30 @@ class ChineseCangjieConverter:
         return "".join(output)
 
 
-class RussianStressLabeler:
-    """Adds stress marks to Russian text when the optional dependency is available."""
-
-    def __init__(self):
-        self._stresser = None
-        self._available = False
-        self._error_logged = False
-        self._initialize()
-
-    def _initialize(self):
-        try:
+def add_russian_stress(text: str) -> str:
+    """Russian text normalization: adds stress marks to Russian text."""
+    global _russian_stresser
+    
+    try:
+        if _russian_stresser is None:
             from russian_text_stresser.text_stresser import RussianTextStresser
-        except ImportError:
-            logger.warning("russian_text_stresser not available - Russian stress labeling skipped")
-            self._error_logged = True
-            return
-        except Exception as exc:
-            logger.warning(f"Failed to import RussianTextStresser: {exc}")
-            self._error_logged = True
-            return
-
-        try:
-            self._stresser = RussianTextStresser()
-            self._available = True
-        except Exception as exc:
-            logger.warning(f"Failed to initialize RussianTextStresser: {exc}")
-            self._error_logged = True
-
-    def __call__(self, text: str) -> str:
-        if not text or not self._available:
-            return text
-
-        try:
-            return self._stresser.stress_text(text)
-        except Exception as exc:
-            if not self._error_logged:
-                logger.warning(f"Russian stress labeling failed: {exc}")
-                self._error_logged = True
-            return text
+            _russian_stresser = RussianTextStresser()
         
+        return _russian_stresser.stress_text(text)
         
+    except ImportError:
+        logger.warning("russian_text_stresser not available - Russian stress labeling skipped")
+        return text
+    except Exception as e:
+        logger.warning(f"Russian stress labeling failed: {e}")
+        return text
+
+
 class MTLTokenizer:
     def __init__(self, vocab_file_path):
         self.tokenizer: Tokenizer = Tokenizer.from_file(vocab_file_path)
         model_dir = Path(vocab_file_path).parent
         self.cangjie_converter = ChineseCangjieConverter(model_dir)
-        self.russian_stress_labeler = RussianStressLabeler()
         self.check_vocabset_sot_eot()
 
     def check_vocabset_sot_eot(self):
@@ -289,12 +264,26 @@ class MTLTokenizer:
         assert SOT in voc
         assert EOT in voc
 
-    def text_to_tokens(self, text: str, language_id: str = None):
-        text_tokens = self.encode(text, language_id=language_id)
+    def preprocess_text(self, raw_text: str, language_id: str = None, lowercase: bool = True, nfkd_normalize: bool = True):
+        """
+        Text preprocessor that handles lowercase conversion and NFKD normalization.
+        """
+        preprocessed_text = raw_text
+        if lowercase:
+            preprocessed_text = preprocessed_text.lower()
+        if nfkd_normalize:
+            preprocessed_text = normalize("NFKD", preprocessed_text)
+        
+        return preprocessed_text
+
+    def text_to_tokens(self, text: str, language_id: str = None, lowercase: bool = True, nfkd_normalize: bool = True):
+        text_tokens = self.encode(text, language_id=language_id, lowercase=lowercase, nfkd_normalize=nfkd_normalize)
         text_tokens = torch.IntTensor(text_tokens).unsqueeze(0)
         return text_tokens
 
-    def encode(self, txt: str, language_id: str = None):
+    def encode(self, txt: str, language_id: str = None, lowercase: bool = True, nfkd_normalize: bool = True):
+        txt = self.preprocess_text(txt, language_id=language_id, lowercase=lowercase, nfkd_normalize=nfkd_normalize)
+        
         # Language-specific text processing
         if language_id == 'zh':
             txt = self.cangjie_converter(txt)
@@ -305,11 +294,7 @@ class MTLTokenizer:
         elif language_id == 'ko':
             txt = korean_normalize(txt)
         elif language_id == 'ru':
-            txt = self.russian_stress_labeler(txt)
-        elif language_id == 'pl':
-            # Polish text normalization: ensure diacritic characters are preserved
-            import unicodedata
-            txt = unicodedata.normalize('NFC', txt)
+            txt = add_russian_stress(txt)
         
         # Prepend language token
         if language_id: