Update inference script for NeMo 2.7.2 with Georgian-aware chunking

README.md

@@ -21,6 +21,7 @@ pipeline_tag: text-to-speech
 A fine-tuned [MagPIE TTS](https://huggingface.co/nvidia/magpie_tts_multilingual_357m) model for Georgian (ქართული) text-to-speech synthesis.
 
 This is the **open-source TTS model fine-tuned specifically for Georgian**, produced as part of the [Georgian TTS Benchmark](https://github.com/NikaGaworworw/TTS_pipelines).
+
 ## Evaluation Results
 
 Evaluated on the full [FLEURS Georgian](https://huggingface.co/datasets/google/fleurs) test set (979 samples) using round-trip intelligibility:
@@ -37,9 +38,8 @@ Evaluated on the full [FLEURS Georgian](https://huggingface.co/datasets/google/f
 ### Installation
 
 ```bash
-# MagPIE TTS requires NeMo 2.8+ (not yet on PyPI — install from source)
-git clone https://github.com/NVIDIA/NeMo.git
-cd NeMo && pip install -e ".[tts]"
+# Requires NeMo 2.7.2 (install from source at the tested commit)
+pip install nemo_toolkit[tts]@git+https://github.com/NVIDIA-NeMo/NeMo.git@3d73c48aca1ae3be44657267b81f25dc3201161a
 pip install huggingface_hub
 ```
 
@@ -48,29 +48,89 @@ pip install huggingface_hub
 ### Inference
 
 ```python
+import re
 import torch
 import torchaudio
 from huggingface_hub import hf_hub_download
-from nemo.collections.tts.models import MagpieTTSModel
-from nemo.collections.tts.parts.utils.tts_dataset_utils import chunk_text_for_inference
+from nemo.collections.tts.models.magpietts import MagpieTTSModel
 
 # Download and load model
 nemo_path = hf_hub_download(repo_id="NMikka/Magpie-TTS-Geo-357m", filename="magpie_tts_georgian.nemo")
 model = MagpieTTSModel.restore_from(nemo_path, map_location="cpu")
 model = model.eval().cuda()
 
+TOKENIZER_NAME = "text_ce_tokenizer"
+MAX_TOKENS_PER_CHUNK = 400  # ~133 Georgian chars, keeps well under 500 decoder steps
+
+
+def split_georgian_text(text: str) -> list[str]:
+    """Split Georgian text into chunks suitable for TTS inference.
+
+    Splitting priority:
+    1. Sentence-ending punctuation (. ! ?)
+    2. Clause-level punctuation (, ; : —)
+    3. Word boundaries (whitespace) as last resort for very long spans
+    """
+    sentences = re.split(r'(?<=[.!?])\s+', text)
+
+    chunks = []
+    for sentence in sentences:
+        est_tokens = len(sentence.encode('utf-8'))
+        if est_tokens <= MAX_TOKENS_PER_CHUNK:
+            chunks.append(sentence)
+            continue
+
+        clauses = re.split(r'(?<=[,;:—])\s+', sentence)
+        current = ""
+        for clause in clauses:
+            combined = f"{current} {clause}".strip() if current else clause
+            if len(combined.encode('utf-8')) <= MAX_TOKENS_PER_CHUNK:
+                current = combined
+            else:
+                if current:
+                    chunks.append(current)
+                if len(clause.encode('utf-8')) > MAX_TOKENS_PER_CHUNK:
+                    words = clause.split()
+                    current = ""
+                    for word in words:
+                        combined = f"{current} {word}".strip() if current else word
+                        if len(combined.encode('utf-8')) <= MAX_TOKENS_PER_CHUNK:
+                            current = combined
+                        else:
+                            if current:
+                                chunks.append(current)
+                            current = word
+                else:
+                    current = clause
+        if current:
+            chunks.append(current)
+
+    return [c for c in chunks if c.strip()]
+
+
+def tokenize_chunks(chunks: list[str], tokenizer, eos_id: int):
+    """Tokenize pre-split text chunks, appending EOS to each."""
+    chunked_tokens = []
+    chunked_tokens_len = []
+    for chunk in chunks:
+        tokens = tokenizer.encode(text=chunk, tokenizer_name=TOKENIZER_NAME)
+        tokens = tokens + [eos_id]
+        tokens = torch.tensor(tokens, dtype=torch.int32)
+        chunked_tokens.append(tokens)
+        chunked_tokens_len.append(tokens.shape[0])
+    return chunked_tokens, chunked_tokens_len
+
+
 # Synthesize
 text = "გამარჯობა, მე მქვია მაგპაი და ქართულად ვლაპარაკობ."
 
-chunked_tokens, chunked_tokens_len, _ = chunk_text_for_inference(
-    text=text,
-    language="ka",
-    tokenizer_name="text_ce_tokenizer",
-    text_tokenizer=model.tokenizer,
-    eos_token_id=model.eos_id,
-)
+if text[-1] not in ".!?,:;":
+    text += "."
 
-chunk_state = model.create_chunk_state(batch_size=1)
+chunks = split_georgian_text(text)
+chunked_tokens, chunked_tokens_len = tokenize_chunks(chunks, model.tokenizer, model.eos_id)
+
+chunk_state = model.create_longform_chunk_state(batch_size=1)
 all_codes = []
 
 for i, (toks, toks_len) in enumerate(zip(chunked_tokens, chunked_tokens_len)):
@@ -80,7 +140,7 @@ for i, (toks, toks_len) in enumerate(zip(chunked_tokens, chunked_tokens_len)):
         "speaker_indices": 1,  # speaker index (0-4)
     }
     with torch.no_grad():
-        output = model.generate_speech(
+        output = model.generate_long_form_speech(
             batch,
             chunk_state=chunk_state,
             end_of_text=[i == len(chunked_tokens) - 1],
@@ -117,15 +177,14 @@ def synthesize(model, text, speaker=1, use_cfg=True):
     Returns:
         waveform (torch.Tensor): Audio tensor, shape (1, num_samples), 22050 Hz
     """
-    chunked_tokens, chunked_tokens_len, _ = chunk_text_for_inference(
-        text=text,
-        language="ka",
-        tokenizer_name="text_ce_tokenizer",
-        text_tokenizer=model.tokenizer,
-        eos_token_id=model.eos_id,
-    )
-
-    chunk_state = model.create_chunk_state(batch_size=1)
+    text = text.strip()
+    if text[-1] not in ".!?,:;":
+        text += "."
+
+    chunks = split_georgian_text(text)
+    chunked_tokens, chunked_tokens_len = tokenize_chunks(chunks, model.tokenizer, model.eos_id)
+
+    chunk_state = model.create_longform_chunk_state(batch_size=1)
     all_codes = []
 
     for i, (toks, toks_len) in enumerate(zip(chunked_tokens, chunked_tokens_len)):
@@ -135,7 +194,7 @@ def synthesize(model, text, speaker=1, use_cfg=True):
             "speaker_indices": speaker,
         }
         with torch.no_grad():
-            output = model.generate_speech(
+            output = model.generate_long_form_speech(
                 batch,
                 chunk_state=chunk_state,
                 end_of_text=[i == len(chunked_tokens) - 1],
@@ -172,6 +231,16 @@ MagPIE TTS is an **encoder-decoder transformer** (not a diffusion or flow model)
 
 **Classifier-Free Guidance (CFG)** runs two forward passes (with/without text conditioning) and interpolates. Set `use_cfg=False` for ~2x faster inference with slightly lower quality.
 
+## Text Chunking
+
+Georgian text requires custom chunking because NeMo's built-in `split_by_sentence` doesn't handle Georgian properly (incorrect capitalization, no splitting of long sentences). The chunker included above splits text with this priority:
+
+1. **Sentence-ending punctuation** (`.` `!` `?`)
+2. **Clause-level punctuation** (`,` `;` `:` `—`)
+3. **Word boundaries** as a last resort
+
+Each chunk is limited to 400 bytes (~133 Georgian characters), keeping well under the model's 500 decoder step limit.
+
 ## Speakers
 
 The model has 5 baked speaker embeddings from pretraining. Set via `speaker_indices` in the batch dict.