Feat/add chunk size control arg (#56)

This pull request addresses issue #38 

**Changes included:**

- Added a `max_tokens` parameter to the typer cli for the `generate`
command, allowing user to set a default number of tokens per chunk.
- Introduced a default variable `MAX_TOKEN_PER_CHUNK` to centralize the
default value.
- Fixed a tokenization issue by inserting a comma, allowing the model to
correctly read each chunk as intended rather than treating one chunk as
a single sentence.

**Future improvement suggestion :**

- Consider enhancing the `split_into_best_sentences` function using
regex to split text more accurately without modifying it, which would
improve clarity and maintainability.

---------

Co-authored-by: Václav Volhejn <8401624+vvolhejn@users.noreply.github.com>

pocket_tts/data/audio_utils.py

@@ -1,4 +1,4 @@
-"""Various utilities for audio convertion (pcm format, sample rate and channels),
+"""Various utilities for audio conversion (pcm format, sample rate and channels),
 and volume normalization."""
 
 import torch

pocket_tts/default_parameters.py

@@ -5,3 +5,4 @@ DEFAULT_LSD_DECODE_STEPS = 1
 DEFAULT_NOISE_CLAMP = None
 DEFAULT_EOS_THRESHOLD = -4.0
 DEFAULT_FRAMES_AFTER_EOS = None
+MAX_TOKEN_PER_CHUNK = 50

pocket_tts/main.py

@@ -22,6 +22,7 @@ from pocket_tts.default_parameters import (
     DEFAULT_NOISE_CLAMP,
     DEFAULT_TEMPERATURE,
     DEFAULT_VARIANT,
+    MAX_TOKEN_PER_CHUNK,
 )
 from pocket_tts.models.tts_model import TTSModel
 from pocket_tts.utils.logging_utils import enable_logging
@@ -223,6 +224,9 @@ def generate(
         str, typer.Option(help="Output path for generated audio")
     ] = "./tts_output.wav",
     device: Annotated[str, typer.Option(help="Device to use")] = "cpu",
+    max_tokens: Annotated[
+        int, typer.Option(help="Maximum number of tokens per chunk.")
+    ] = MAX_TOKEN_PER_CHUNK,
 ):
     """Generate speech using Kyutai Pocket TTS."""
     if "cuda" in device:
@@ -242,6 +246,7 @@ def generate(
             model_state=model_state_for_voice,
             text_to_generate=text,
             frames_after_eos=frames_after_eos,
+            max_tokens=max_tokens,
         )
 
         stream_audio_chunks(output_path, audio_chunks, tts_model.config.mimi.sample_rate)

pocket_tts/models/tts_model.py

@@ -23,6 +23,7 @@ from pocket_tts.default_parameters import (
     DEFAULT_NOISE_CLAMP,
     DEFAULT_TEMPERATURE,
     DEFAULT_VARIANT,
+    MAX_TOKEN_PER_CHUNK,
 )
 from pocket_tts.models.flow_lm import FlowLMModel
 from pocket_tts.models.mimi import MimiModel
@@ -363,6 +364,7 @@ class TTSModel(nn.Module):
         self,
         model_state: dict,
         text_to_generate: str,
+        max_tokens: int = MAX_TOKEN_PER_CHUNK,
         frames_after_eos: int | None = None,
         copy_state: bool = True,
     ) -> torch.Tensor:
@@ -405,6 +407,7 @@ class TTSModel(nn.Module):
             text_to_generate=text_to_generate,
             frames_after_eos=frames_after_eos,
             copy_state=copy_state,
+            max_tokens=max_tokens,
         ):
             audio_chunks.append(chunk)
         return torch.cat(audio_chunks, dim=0)
@@ -414,6 +417,7 @@ class TTSModel(nn.Module):
         self,
         model_state: dict,
         text_to_generate: str,
+        max_tokens: int = MAX_TOKEN_PER_CHUNK,
         frames_after_eos: int | None = None,
         copy_state: bool = True,
     ):
@@ -458,7 +462,9 @@ class TTSModel(nn.Module):
         # by using teacher forcing, but it would be a bit slower.
         # TODO: add the teacher forcing method for long texts where we use the audio of one chunk
         # as conditioning for the next chunk.
-        chunks = split_into_best_sentences(self.flow_lm.conditioner.tokenizer, text_to_generate)
+        chunks = split_into_best_sentences(
+            self.flow_lm.conditioner.tokenizer, text_to_generate, max_tokens
+        )
 
         for chunk in chunks:
             text_to_generate, frames_after_eos_guess = prepare_text_prompt(chunk)
@@ -733,7 +739,7 @@ def prepare_text_prompt(text: str) -> tuple[str, int]:
     return text, frames_after_eos_guess
 
 
-def split_into_best_sentences(tokenizer, text_to_generate: str) -> list[str]:
+def split_into_best_sentences(tokenizer, text_to_generate: str, max_tokens: int) -> list[str]:
     text_to_generate, _ = prepare_text_prompt(text_to_generate)
     text_to_generate = text_to_generate.strip()
     tokens = tokenizer(text_to_generate)
@@ -761,7 +767,7 @@ def split_into_best_sentences(tokenizer, text_to_generate: str) -> list[str]:
         text = tokenizer.sp.decode(list_of_tokens[start:end])
         nb_tokens_and_sentences.append((end - start, text))
 
-    max_nb_tokens_in_a_chunk = 50
+    max_nb_tokens_in_a_chunk = max_tokens
     chunks = []
     current_chunk = ""
     current_nb_of_tokens_in_chunk = 0