Multilingual v2 update (resemble-ai#295)

* multilingual v2 vocab and russian stresser update

* multilinugal tokenizer fix

Author: ZihanJin

pyproject.toml

@@ -28,6 +28,7 @@ dependencies = [
     "spacy-pkuseg",
     "pykakasi==2.3.0",
     # "gradio==5.44.1",
+    "russian-text-stresser @ git+https://github.com/Vuizur/add-stress-to-epub",
 ]
 
 [project.urls]

src/chatterbox/models/t3/modules/t3_config.py

@@ -28,7 +28,7 @@ def n_channels(self):
 
     @property
     def is_multilingual(self):
-        return self.text_tokens_dict_size == 2352
+        return self.text_tokens_dict_size == 2454
 
     @classmethod
     def english_only(cls):
@@ -38,4 +38,4 @@ def english_only(cls):
     @classmethod 
     def multilingual(cls):
         """Create configuration for multilingual TTS model."""
-        return cls(text_tokens_dict_size=2352)
+        return cls(text_tokens_dict_size=2454)

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 @@ def text_to_tokens(self, text: str):
         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 @@ def decode(self, seq):
         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 @@ def decode(self, seq):
 # Global instances for optional dependencies
 _kakasi = None
 _dicta = None
+_russian_stresser = None
 
 
 def is_kanji(c: str) -> bool:
@@ -281,6 +280,25 @@ def __call__(self, text):
         return "".join(output)
 
 
+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
+            _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)
@@ -293,12 +311,26 @@ def check_vocabset_sot_eot(self):
         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)
@@ -310,6 +342,8 @@ def encode(self, txt: str, language_id: str = None):
             txt = korean_normalize(txt)
         elif language_id == 'fr': # Author: Rouxin  
             txt = decompose_french_text(txt)
+        elif language_id == 'ru':
+            txt = add_russian_stress(txt)
 
         # Prepend language token
         if language_id:

src/chatterbox/mtl_tts.py

@@ -168,7 +168,7 @@ def from_local(cls, ckpt_dir, device) -> 'ChatterboxMultilingualTTS':
         ve.to(device).eval()
 
         t3 = T3(T3Config.multilingual())
-        t3_state = load_safetensors(ckpt_dir / "t3_23lang.safetensors")
+        t3_state = load_safetensors(ckpt_dir / "t3_mtl23ls_v2.safetensors")
         if "model" in t3_state.keys():
             t3_state = t3_state["model"][0]
         t3.load_state_dict(t3_state)
@@ -181,7 +181,7 @@ def from_local(cls, ckpt_dir, device) -> 'ChatterboxMultilingualTTS':
         s3gen.to(device).eval()
 
         tokenizer = MTLTokenizer(
-            str(ckpt_dir / "mtl_tokenizer.json")
+            str(ckpt_dir / "grapheme_mtl_merged_expanded_v1.json")
         )
 
         conds = None
@@ -197,7 +197,7 @@ def from_pretrained(cls, device: torch.device) -> 'ChatterboxMultilingualTTS':
                 repo_id=REPO_ID,
                 repo_type="model",
                 revision="main", 
-                allow_patterns=["ve.pt", "t3_23lang.safetensors", "s3gen.pt", "mtl_tokenizer.json", "conds.pt", "Cangjie5_TC.json"],
+                allow_patterns=["ve.pt", "t3_mtl23ls_v2.safetensors", "s3gen.pt", "grapheme_mtl_merged_expanded_v1.json", "conds.pt", "Cangjie5_TC.json"],
                 token=os.getenv("HF_TOKEN"),
             )
         )