sesame

第7章 · 模型后训练 · 配套项目 chapter7/sesame

项目说明

Sesame CSM (1B) TTS - Text-to-Speech Fine-tuning

This directory contains scripts for fine-tuning and running inference with the Sesame CSM text-to-speech model using Unsloth.

Files

• sesame_csm_sft_unsloth.py - Training script for fine-tuning the model with LoRA
• inference.py - Single inference script for generating speech from text
• batch_inference.py - Batch inference script for processing multiple texts
• example_inputs.json - Example input file for batch inference
• requirements.txt - Python dependencies

Installation

# Install dependencies
pip install -r requirements.txt

# For Conda users, install ffmpeg
conda install -c conda-forge "ffmpeg>=6.0" -y
conda install -c conda-forge libiconv -y

Training

To train a model with your own dataset:

python sesame_csm_sft_unsloth.py

This will: 1. Load the base model unsloth/csm-1b 2. Add LoRA adapters 3. Fine-tune on your dataset 4. Save the LoRA adapters to lora_model/

Inference

Single Text Inference

Generate speech from a single text:

# Without context (simple generation)
python inference.py \
    --lora-path lora_model \
    --text "We just finished fine tuning a text to speech model... and it's pretty good!" \
    --output example_without_context_1.wav

# With voice context (for voice consistency) using dataset index 3
python inference.py \
    --lora-path lora_model \
    --text "Sesame is a super cool TTS model which can be fine tuned with Unsloth." \
    --dataset-context-idx 3 \
    --output example_with_context_1.wav

# Using base model only (without LoRA)
python inference.py \
    --text "Hello world, this is a test of the text to speech system."

# With custom speaker ID and longer generation
python inference.py \
    --lora-path lora_model \
    --text "This is a longer sentence that needs more tokens." \
    --speaker-id 0 \
    --max-tokens 250 \
    --output long_speech.wav

# With 4-bit quantization (lower memory)
python inference.py \
    --lora-path lora_model \
    --text "Memory efficient inference." \
    --load-in-4bit

Batch Inference

Process multiple texts at once:

# From JSON file
python batch_inference.py \
    --lora-path lora_model \
    --input-file example_inputs.json \
    --output-dir batch_outputs

# From plain text file (one text per line)
python batch_inference.py \
    --lora-path lora_model \
    --input-file texts.txt \
    --output-dir batch_outputs

Input File Formats

JSON format (example_inputs.json):

Without context:

[
    {
        "text": "We just finished fine tuning a text to speech model... and it's pretty good!",
        "speaker_id": 0,
        "output": "example_without_context_1.wav"
    },
    {
        "text": "Sesame is a super cool TTS model which can be fine tuned with Unsloth.",
        "speaker_id": 0,
        "output": "example_without_context_2.wav"
    }
]

With context (for voice consistency using dataset indices):

[
    {
        "text": "Sesame is a super cool TTS model which can be fine tuned with Unsloth.",
        "speaker_id": 0,
        "dataset_context_idx": 3,
        "output": "example_with_context_1.wav"
    },
    {
        "text": "We just finished fine tuning a text to speech model... and it's pretty good!",
        "speaker_id": 0,
        "dataset_context_idx": 4,
        "output": "example_with_context_2.wav"
    }
]

Note: dataset_context_idx refers to the index in the training dataset (e.g., MrDragonFox/Elise). Indices 3 and 4 are used in the training script examples.

Plain text format (one sentence per line, no context support):

Hello world, this is the first sentence.
This is the second sentence.
This is the third sentence.

Parameters

Common Parameters

• --base-model: Base model name or path (default: unsloth/csm-1b)
• --lora-path: Path to saved LoRA adapters (optional)
• --load-in-4bit: Load model in 4-bit quantization to reduce memory usage

Inference Parameters

• --text: Text to convert to speech
• --speaker-id: Speaker ID for multi-speaker models (default: 0)
• --output: Output audio file path (default: output.wav)
• --max-tokens: Maximum tokens to generate (125 ≈ 10 seconds of audio)
• --dataset-context-idx: Dataset index to use for voice consistency (e.g., 3 or 4 from training examples)
• --dataset-name: Dataset name to load context from (default: MrDragonFox/Elise)

Batch Inference Parameters

• --input-file: Input file (JSON or plain text)
• --output-dir: Output directory for audio files (default: outputs)

Model Information

• Base Model: Sesame CSM (1B) - A compact text-to-speech model
• Output Format: 24kHz WAV audio
• Token-to-Time Ratio: Approximately 125 tokens = 10 seconds of audio
• Multi-speaker: Supports multiple speakers via speaker IDs

Advanced Usage

Voice Consistency with Context

For better voice consistency, you can provide audio context from the dataset (see sesame_csm_sft_unsloth.py lines 320-390 for examples):

from inference import load_model, generate_speech

model, processor = load_model("unsloth/csm-1b", "lora_model")

# Use dataset index 3 for voice consistency (same as training script)
generate_speech(
    model=model,
    processor=processor,
    text="Sesame is a super cool TTS model which can be fine tuned with Unsloth.",
    dataset_context_idx=3,
    dataset_name="MrDragonFox/Elise",
    output_path="output_with_context.wav"
)

The inference scripts automatically load the audio and text from the specified dataset index, ensuring consistency with the training approach.

Memory Requirements

• Base model: ~4-6GB VRAM
• Base model + LoRA training: ~8-12GB VRAM
• With 4-bit quantization: ~2-3GB VRAM
• Inference only: ~2-4GB VRAM

Use --load-in-4bit flag if you have limited GPU memory.

Troubleshooting

AssertionError during generation

If you encounter AssertionError during model.generate(), ensure you're passing tensors correctly:

# ❌ Wrong
audio_values = model.generate(**inputs, max_new_tokens=125)

# ✅ Correct
audio_values = model.generate(
    input_ids=inputs["input_ids"],
    attention_mask=inputs.get("attention_mask"),
    max_new_tokens=125,
    output_audio=True
)

Out of memory

• Use --load-in-4bit flag
• Reduce max_new_tokens
• Process texts one at a time instead of batching
• Use a smaller batch size during training

Resources

• Unsloth Documentation
• Unsloth TTS Guide
• Unsloth Discord
• Unsloth GitHub

License

This project uses the Unsloth library and Sesame CSM model. Please refer to their respective licenses.

源代码

batch_inference.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Sesame CSM (1B) TTS - Batch Inference Script

This script loads a trained LoRA model and generates speech from multiple texts.
"""

import argparse
import json
import torch
import soundfile as sf
from pathlib import Path
from typing import List, Dict
from tqdm import tqdm
from datasets import load_dataset, Audio
from unsloth import FastModel
from transformers import CsmForConditionalGeneration
from peft import PeftModel


def load_model(base_model_name: str, lora_path: str = None, load_in_4bit: bool = False):
    """Load the base model and optionally apply LoRA adapters."""
    print(f"Loading base model: {base_model_name}")
    model, processor = FastModel.from_pretrained(
        model_name=base_model_name,
        max_seq_length=2048,
        dtype=None,
        auto_model=CsmForConditionalGeneration,
        load_in_4bit=load_in_4bit,
    )

    if lora_path:
        print(f"Loading LoRA adapters from: {lora_path}")
        model = PeftModel.from_pretrained(model, lora_path)

    return model, processor


def load_texts_from_file(input_file: str) -> List[Dict]:
    """
    Load texts from a JSON file.

    Expected format:
    [
        {"text": "Hello world", "speaker_id": 0, "output": "hello.wav"},
        {"text": "Another sentence", "speaker_id": 0, "output": "another.wav"}
    ]

    Or simple text file (one text per line):
    Hello world
    Another sentence
    """
    input_path = Path(input_file)

    if input_path.suffix == '.json':
        with open(input_path, 'r', encoding='utf-8') as f:
            return json.load(f)
    else:
        # Plain text file
        with open(input_path, 'r', encoding='utf-8') as f:
            lines = [line.strip() for line in f if line.strip()]
        return [
            {
                "text": line,
                "speaker_id": 0,
                "output": f"output_{i:04d}.wav"
            }
            for i, line in enumerate(lines)
        ]


def load_dataset_for_context(dataset_name: str = "MrDragonFox/Elise", split: str = "train"):
    """Load the dataset for voice context examples."""
    raw_ds = load_dataset(dataset_name, split=split)
    target_sampling_rate = 24000
    raw_ds = raw_ds.cast_column("audio", Audio(sampling_rate=target_sampling_rate))
    return raw_ds


def generate_speech_batch(
    model,
    processor,
    texts: List[Dict],
    output_dir: str,
    max_new_tokens: int = 125,
    dataset_name: str = "MrDragonFox/Elise",
):
    """Generate speech for multiple texts."""
    output_path = Path(output_dir)
    output_path.mkdir(parents=True, exist_ok=True)

    device = "cuda" if torch.cuda.is_available() else "cpu"

    # Load dataset once if any item needs context
    raw_ds = None
    needs_context = any(item.get("dataset_context_idx") is not None for item in texts)
    if needs_context:
        print(f"Loading dataset: {dataset_name}")
        raw_ds = load_dataset_for_context(dataset_name)
        print(f"Loaded {len(raw_ds)} examples from dataset")

    for item in tqdm(texts, desc="Generating speech"):
        text = item.get("text")
        if not text:
            raise ValueError(f"Each item must have a non-empty 'text' field, got: {item}")
        speaker_id = item.get("speaker_id", 0)
        output_name = item.get("output") or f"output_{hash(text)}.wav"
        output_file = output_path / output_name

        # Check if dataset context is provided
        dataset_context_idx = item.get("dataset_context_idx")

        if dataset_context_idx is not None:
            # Generate with voice context from dataset
            context_example = raw_ds[dataset_context_idx]
            context_audio = context_example["audio"]["array"]
            context_text = context_example["text"]

            conversation = [
                {
                    "role": str(speaker_id),
                    "content": [
                        {"type": "text", "text": context_text},
                        {"type": "audio", "path": context_audio}
                    ]
                },
                {
                    "role": str(speaker_id),
                    "content": [{"type": "text", "text": text}]
                },
            ]
            inputs = processor.apply_chat_template(
                conversation,
                tokenize=True,
                return_dict=True,
            ).to(device)
        else:
            # Generate without context
            inputs = processor(
                f"[{speaker_id}]{text}",
                add_special_tokens=True,
                return_tensors="pt"
            ).to(device)

        # Generate audio
        with torch.no_grad():
            audio_values = model.generate(
                input_ids=inputs["input_ids"],
                attention_mask=inputs.get("attention_mask"),
                max_new_tokens=max_new_tokens,
                output_audio=True,
            )

        # Save audio
        audio = audio_values[0].to(torch.float32).cpu().numpy()
        sf.write(output_file, audio, 24000)

    print(f"\nGenerated {len(texts)} audio files in: {output_dir}")


def main():
    parser = argparse.ArgumentParser(
        description="Batch generate speech using Sesame CSM TTS model"
    )
    parser.add_argument(
        "--base-model",
        type=str,
        default="unsloth/csm-1b",
        help="Base model name or path (default: unsloth/csm-1b)"
    )
    parser.add_argument(
        "--lora-path",
        type=str,
        default=None,
        help="Path to saved LoRA adapters (optional)"
    )
    parser.add_argument(
        "--input-file",
        type=str,
        required=True,
        help="Input file (JSON or plain text, one text per line)"
    )
    parser.add_argument(
        "--output-dir",
        type=str,
        default="outputs",
        help="Output directory for audio files (default: outputs)"
    )
    parser.add_argument(
        "--max-tokens",
        type=int,
        default=125,
        help="Maximum tokens to generate (125 ≈ 10 seconds) (default: 125)"
    )
    parser.add_argument(
        "--load-in-4bit",
        action="store_true",
        help="Load model in 4-bit quantization to reduce memory usage"
    )
    parser.add_argument(
        "--dataset-name",
        type=str,
        default="MrDragonFox/Elise",
        help="Dataset name to load context from (default: MrDragonFox/Elise)"
    )

    args = parser.parse_args()

    # Load texts
    print(f"Loading texts from: {args.input_file}")
    texts = load_texts_from_file(args.input_file)
    print(f"Loaded {len(texts)} texts")

    # Load model
    model, processor = load_model(
        base_model_name=args.base_model,
        lora_path=args.lora_path,
        load_in_4bit=args.load_in_4bit
    )

    # Generate speech
    generate_speech_batch(
        model=model,
        processor=processor,
        texts=texts,
        output_dir=args.output_dir,
        max_new_tokens=args.max_tokens,
        dataset_name=args.dataset_name,
    )


if __name__ == "__main__":
    main()

inference.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Sesame CSM (1B) TTS - Inference Script

This script loads a trained LoRA model and generates speech from text.
"""

import argparse
import torch
import soundfile as sf
from pathlib import Path
from datasets import load_dataset, Audio
from unsloth import FastModel
from transformers import CsmForConditionalGeneration
from peft import PeftModel


def load_model(base_model_name: str, lora_path: str = None, load_in_4bit: bool = False):
    """
    Load the base model and optionally apply LoRA adapters.

    Args:
        base_model_name: Name or path of the base model
        lora_path: Path to saved LoRA adapters (optional)
        load_in_4bit: Whether to load model in 4-bit quantization

    Returns:
        model, processor
    """
    print(f"Loading base model: {base_model_name}")
    model, processor = FastModel.from_pretrained(
        model_name=base_model_name,
        max_seq_length=2048,
        dtype=None,  # Auto-detection
        auto_model=CsmForConditionalGeneration,
        load_in_4bit=load_in_4bit,
    )

    if lora_path:
        print(f"Loading LoRA adapters from: {lora_path}")
        model = PeftModel.from_pretrained(model, lora_path)

    return model, processor


def load_dataset_for_context(dataset_name: str = "MrDragonFox/Elise", split: str = "train"):
    """
    Load the dataset for voice context examples.

    Args:
        dataset_name: Name of the dataset
        split: Dataset split to use

    Returns:
        Dataset
    """
    print(f"Loading dataset: {dataset_name}")
    raw_ds = load_dataset(dataset_name, split=split)
    target_sampling_rate = 24000
    raw_ds = raw_ds.cast_column("audio", Audio(sampling_rate=target_sampling_rate))
    print(f"Loaded {len(raw_ds)} examples from dataset")
    return raw_ds


def generate_speech(
    model,
    processor,
    text: str,
    speaker_id: int = 0,
    max_new_tokens: int = 125,
    output_path: str = "output.wav",
    dataset_context_idx: int = None,
    dataset_name: str = "MrDragonFox/Elise",
):
    """
    Generate speech from text.

    Args:
        model: The loaded model
        processor: The processor
        text: Text to convert to speech
        speaker_id: Speaker ID (for multi-speaker models)
        max_new_tokens: Maximum number of tokens to generate (125 tokens ≈ 10 seconds)
        output_path: Path to save the output audio file
        dataset_context_idx: Optional dataset index to use for voice consistency
        dataset_name: Name of the dataset to load context from
    """
    device = "cuda" if torch.cuda.is_available() else "cpu"

    # Prepare inputs based on whether context is provided
    if dataset_context_idx is not None:
        print(f"Generating speech with voice context from dataset index: {dataset_context_idx}")
        print(f"Target text: '{text}'")

        # Load dataset and get context example
        raw_ds = load_dataset_for_context(dataset_name)
        context_example = raw_ds[dataset_context_idx]
        context_audio = context_example["audio"]["array"]
        context_text = context_example["text"]

        print(f"Context text: '{context_text}'")

        # Use conversation format with audio context for voice consistency
        conversation = [
            {
                "role": str(speaker_id),
                "content": [
                    {"type": "text", "text": context_text},
                    {"type": "audio", "path": context_audio}
                ]
            },
            {
                "role": str(speaker_id),
                "content": [{"type": "text", "text": text}]
            },
        ]
        inputs = processor.apply_chat_template(
            conversation,
            tokenize=True,
            return_dict=True,
        ).to(device)
    else:
        print(f"Generating speech without context for: '{text}'")

        # Simple text-only input
        inputs = processor(
            f"[{speaker_id}]{text}",
            add_special_tokens=True,
            return_tensors="pt"
        ).to(device)

    # Generate audio
    print("Generating audio...")
    with torch.no_grad():
        audio_values = model.generate(
            input_ids=inputs["input_ids"],
            attention_mask=inputs.get("attention_mask"),
            max_new_tokens=max_new_tokens,
            output_audio=True,
        )

    # Save audio
    audio = audio_values[0].to(torch.float32).cpu().numpy()
    sf.write(output_path, audio, 24000)
    print(f"Audio saved to: {output_path}")


def main():
    parser = argparse.ArgumentParser(
        description="Generate speech using Sesame CSM TTS model"
    )
    parser.add_argument(
        "--base-model",
        type=str,
        default="unsloth/csm-1b",
        help="Base model name or path (default: unsloth/csm-1b)"
    )
    parser.add_argument(
        "--lora-path",
        type=str,
        default=None,
        help="Path to saved LoRA adapters (optional)"
    )
    parser.add_argument(
        "--text",
        type=str,
        required=True,
        help="Text to convert to speech"
    )
    parser.add_argument(
        "--speaker-id",
        type=int,
        default=0,
        help="Speaker ID (default: 0)"
    )
    parser.add_argument(
        "--output",
        type=str,
        default="output.wav",
        help="Output audio file path (default: output.wav)"
    )
    parser.add_argument(
        "--max-tokens",
        type=int,
        default=125,
        help="Maximum tokens to generate (125 ≈ 10 seconds) (default: 125)"
    )
    parser.add_argument(
        "--load-in-4bit",
        action="store_true",
        help="Load model in 4-bit quantization to reduce memory usage"
    )
    parser.add_argument(
        "--dataset-context-idx",
        type=int,
        default=None,
        help="Dataset index to use for voice consistency (e.g., 3 or 4 from training examples)"
    )
    parser.add_argument(
        "--dataset-name",
        type=str,
        default="MrDragonFox/Elise",
        help="Dataset name to load context from (default: MrDragonFox/Elise)"
    )

    args = parser.parse_args()

    # Load model
    model, processor = load_model(
        base_model_name=args.base_model,
        lora_path=args.lora_path,
        load_in_4bit=args.load_in_4bit
    )

    # Generate speech
    generate_speech(
        model=model,
        processor=processor,
        text=args.text,
        speaker_id=args.speaker_id,
        max_new_tokens=args.max_tokens,
        output_path=args.output,
        dataset_context_idx=args.dataset_context_idx,
        dataset_name=args.dataset_name,
    )


if __name__ == "__main__":
    main()

sesame_csm_sft_unsloth.py

# -*- coding: utf-8 -*-
"""Sesame_CSM_(1B)-TTS.ipynb

Automatically generated by Colab.

Original file is located at
    https://colab.research.google.com/github/unslothai/notebooks/blob/main/nb/Sesame_CSM_(1B)-TTS.ipynb

To run this, press "*Runtime*" and press "*Run all*" on a **free** Tesla T4 Google Colab instance!
<div class="align-center">
<a href="https://unsloth.ai/"><img src="https://github.com/unslothai/unsloth/raw/main/images/unsloth%20new%20logo.png" width="115"></a>
<a href="https://discord.gg/unsloth"><img src="https://github.com/unslothai/unsloth/raw/main/images/Discord button.png" width="145"></a>
<a href="https://docs.unsloth.ai/"><img src="https://github.com/unslothai/unsloth/blob/main/images/documentation%20green%20button.png?raw=true" width="125"></a></a> Join Discord if you need help + ⭐ <i>Star us on <a href="https://github.com/unslothai/unsloth">Github</a> </i> ⭐
</div>

To install Unsloth on your own computer, follow the installation instructions on our Github page [here](https://docs.unsloth.ai/get-started/installing-+-updating).

You will learn how to do [data prep](#Data), how to [train](#Train), how to [run the model](#Inference), & [how to save it](#Save)

### News

Unsloth now supports [gpt-oss RL](https://docs.unsloth.ai/new/gpt-oss-reinforcement-learning) with the fastest inference & lowest VRAM. Try our [new notebook](https://colab.research.google.com/github/unslothai/notebooks/blob/main/nb/gpt-oss-(20B)-GRPO.ipynb) which automatically creates kernels!

[Vision RL](https://docs.unsloth.ai/new/vision-reinforcement-learning-vlm-rl) is now supported! Train Qwen2.5-VL, Gemma 3 etc. with GSPO or GRPO.

Introducing Unsloth [Standby for RL](https://docs.unsloth.ai/basics/memory-efficient-rl): GRPO is now faster, uses 30% less memory with 2x longer context.

Unsloth now supports Text-to-Speech (TTS) models. Read our [guide here](https://docs.unsloth.ai/basics/text-to-speech-tts-fine-tuning).

Visit our docs for all our [model uploads](https://docs.unsloth.ai/get-started/all-our-models) and [notebooks](https://docs.unsloth.ai/get-started/unsloth-notebooks).

### Installation
"""

# Commented out IPython magic to ensure Python compatibility.
# %%capture
# import os, re
# if "COLAB_" not in "".join(os.environ.keys()):
#     !pip install unsloth
# else:
#     # Do this only in Colab notebooks! Otherwise use pip install unsloth
#     import torch; v = re.match(r"[0-9\.]{3,}", str(torch.__version__)).group(0)
#     xformers = "xformers==" + ("0.0.32.post2" if v == "2.8.0" else "0.0.29.post3")
#     !pip install --no-deps bitsandbytes accelerate {xformers} peft trl triton cut_cross_entropy unsloth_zoo
#     !pip install sentencepiece protobuf "datasets>=3.4.1,<4.0.0" "huggingface_hub>=0.34.0" hf_transfer
#     !pip install --no-deps unsloth
# !pip install transformers==4.52.3
# !pip install --no-deps trl==0.22.2
# !pip install torchcodec
# !pip install soundfile

# Install ffmpeg in conda environment
# !conda install -c conda-forge "ffmpeg>=6.0" -y
# !conda install -c conda-forge libiconv -y

"""### Unsloth

`FastModel` supports loading nearly any model now! This includes Vision and Text models!
"""

from unsloth import FastModel
from transformers import CsmForConditionalGeneration
import torch

model, processor = FastModel.from_pretrained(
    model_name = "unsloth/csm-1b",
    max_seq_length= 2048, # Choose any for long context!
    dtype = None, # Leave as None for auto-detection
    auto_model = CsmForConditionalGeneration,
    load_in_4bit = False, # Select True for 4bit - reduces memory usage
)

"""We now add LoRA adapters so we only need to update 1 to 10% of all parameters!"""

model = FastModel.get_peft_model(
    model,
    r = 32, # Choose any number > 0 ! Suggested 8, 16, 32, 64, 128
    target_modules = ["q_proj", "k_proj", "v_proj", "o_proj",
                      "gate_proj", "up_proj", "down_proj",],
    lora_alpha = 32,
    lora_dropout = 0, # Supports any, but = 0 is optimized
    bias = "none",    # Supports any, but = "none" is optimized
    # [NEW] "unsloth" uses 30% less VRAM, fits 2x larger batch sizes!
    use_gradient_checkpointing = "unsloth", # True or "unsloth" for very long context
    random_state = 3407,
    use_rslora = False,  # We support rank stabilized LoRA
    loftq_config = None, # And LoftQ
)

"""<a name="Data"></a>
### Data Prep  

We will use the `MrDragonFox/Elise`, which is designed for training TTS models. Ensure that your dataset follows the required format: **text, audio** for single-speaker models or **source, text, audio** for multi-speaker models. You can modify this section to accommodate your own dataset, but maintaining the correct structure is essential for optimal training.
"""

#@title Dataset Prep functions
from datasets import load_dataset, Audio, Dataset
import os
from transformers import AutoProcessor
processor = AutoProcessor.from_pretrained("unsloth/csm-1b")

raw_ds = load_dataset("MrDragonFox/Elise", split="train")

# Getting the speaker id is important for multi-speaker models and speaker consistency
speaker_key = "source"
if "source" not in raw_ds.column_names and "speaker_id" not in raw_ds.column_names:
    print("Unsloth: No speaker found, adding default \"source\" of 0 for all examples")
    new_column = ["0"] * len(raw_ds)
    raw_ds = raw_ds.add_column("source", new_column)
elif "source" not in raw_ds.column_names and "speaker_id" in raw_ds.column_names:
    speaker_key = "speaker_id"

target_sampling_rate = 24000
raw_ds = raw_ds.cast_column("audio", Audio(sampling_rate=target_sampling_rate))

def preprocess_example(example):
    conversation = [
        {
            "role": str(example[speaker_key]),
            "content": [
                {"type": "text", "text": example["text"]},
                {"type": "audio", "path": example["audio"]["array"]},
            ],
        }
    ]

    try:
        model_inputs = processor.apply_chat_template(
            conversation,
            tokenize=True,
            return_dict=True,
            output_labels=True,
            text_kwargs = {
                "padding": "max_length", # pad to the max_length
                "max_length": 256, # this should be the max length of audio
                "pad_to_multiple_of": 8,
                "padding_side": "right",
            },
            audio_kwargs = {
                "sampling_rate": 24_000,
                "max_length": 240001, # max input_values length of the whole dataset
                "padding": "max_length",
            },
            common_kwargs = {"return_tensors": "pt"},
        )
    except Exception as e:
        print(f"Error processing example with text '{example['text'][:50]}...': {e}")
        return None

    required_keys = ["input_ids", "attention_mask", "labels", "input_values", "input_values_cutoffs"]
    processed_example = {}
    # print(model_inputs.keys())
    for key in required_keys:
        if key not in model_inputs:
            print(f"Warning: Required key '{key}' not found in processor output for example.")
            return None

        value = model_inputs[key][0]
        processed_example[key] = value


    # Final check (optional but good)
    if not all(isinstance(processed_example[key], torch.Tensor) for key in processed_example):
         print(f"Error: Not all required keys are tensors in final processed example. Keys: {list(processed_example.keys())}")
         return None

    return processed_example

processed_ds = raw_ds.map(
    preprocess_example,
    remove_columns=raw_ds.column_names,
    desc="Preprocessing dataset",
)

"""<a name="Train"></a>
### Train the model
Now let's use Huggingface  `Trainer`! More docs here: [Transformers docs](https://huggingface.co/docs/transformers/main_classes/trainer).
"""

from transformers import TrainingArguments, Trainer
from unsloth import is_bfloat16_supported

trainer = Trainer(
    model = model,
    train_dataset = processed_ds,
    args = TrainingArguments(
        per_device_train_batch_size = 2,
        gradient_accumulation_steps = 4,
        warmup_steps = 5,
        max_steps = 60,
        learning_rate = 2e-4,
        fp16 = not is_bfloat16_supported(),
        bf16 = is_bfloat16_supported(),
        logging_steps = 1,
        optim = "adamw_8bit",
        weight_decay = 0.01, # Turn this on if overfitting
        lr_scheduler_type = "linear",
        seed = 42,
        output_dir = "outputs",
        report_to = "none", # Use this for WandB etc
    ),
)

# @title Show current memory stats
gpu_stats = torch.cuda.get_device_properties(0)
start_gpu_memory = round(torch.cuda.max_memory_reserved() / 1024 / 1024 / 1024, 3)
max_memory = round(gpu_stats.total_memory / 1024 / 1024 / 1024, 3)
print(f"GPU = {gpu_stats.name}. Max memory = {max_memory} GB.")
print(f"{start_gpu_memory} GB of memory reserved.")

trainer_stats = trainer.train()

# @title Show final memory and time stats
used_memory = round(torch.cuda.max_memory_reserved() / 1024 / 1024 / 1024, 3)
used_memory_for_lora = round(used_memory - start_gpu_memory, 3)
used_percentage = round(used_memory / max_memory * 100, 3)
lora_percentage = round(used_memory_for_lora / max_memory * 100, 3)
print(f"{trainer_stats.metrics['train_runtime']} seconds used for training.")
print(
    f"{round(trainer_stats.metrics['train_runtime']/60, 2)} minutes used for training."
)
print(f"Peak reserved memory = {used_memory} GB.")
print(f"Peak reserved memory for training = {used_memory_for_lora} GB.")
print(f"Peak reserved memory % of max memory = {used_percentage} %.")
print(f"Peak reserved memory for training % of max memory = {lora_percentage} %.")

"""<a name="Save"></a>
### Saving, loading finetuned models
To save the final model as LoRA adapters, either use Huggingface's `push_to_hub` for an online save or `save_pretrained` for a local save.

**[NOTE]** This ONLY saves the LoRA adapters, and not the full model. To save to 16bit or GGUF, scroll down!
"""

model.save_pretrained("lora_model")  # Local saving
processor.save_pretrained("lora_model")
# model.push_to_hub("your_name/lora_model", token = "...") # Online saving
# processor.push_to_hub("your_name/lora_model", token = "...") # Online saving

"""### Saving to float16

We also support saving to `float16` directly. Select `merged_16bit` for float16 or `merged_4bit` for int4. We also allow `lora` adapters as a fallback. Use `push_to_hub_merged` to upload to your Hugging Face account! You can go to https://huggingface.co/settings/tokens for your personal tokens.
"""

# Merge to 16bit
if False: model.save_pretrained_merged("model", processor, save_method = "merged_16bit",)
if False: model.push_to_hub_merged("hf/model", processor, save_method = "merged_16bit", token = "")

# Merge to 4bit
if False: model.save_pretrained_merged("model", processor, save_method = "merged_4bit",)
if False: model.push_to_hub_merged("hf/model", processor, save_method = "merged_4bit", token = "")

# Just LoRA adapters
if False:
    model.save_pretrained("model")
    processor.save_pretrained("model")
if False:
    model.push_to_hub("hf/model", token = "")
    processor.push_to_hub("hf/model", token = "")

"""<a name="Inference"></a>
### Inference
Let's run the model! You can change the prompts
"""

import soundfile as sf

text = "We just finished fine tuning a text to speech model... and it's pretty good!"
speaker_id = 0
inputs = processor(f"[{speaker_id}]{text}", add_special_tokens=True, return_tensors="pt").to("cuda")
audio_values = model.generate(
    input_ids=inputs["input_ids"],
    attention_mask=inputs.get("attention_mask"),
    max_new_tokens=125, # 125 tokens is 10 seconds of audio, for longer speech increase this
    # play with these parameters to tweak results
    # depth_decoder_top_k=0,
    # depth_decoder_top_p=0.9,
    # depth_decoder_do_sample=True,
    # depth_decoder_temperature=0.9,
    # top_k=0,
    # top_p=1.0,
    # temperature=0.9,
    # do_sample=True,
    #########################################################
    output_audio=True
)
audio = audio_values[0].to(torch.float32).cpu().numpy()
sf.write("example_without_context_1.wav", audio, 24000)

text = "Sesame is a super cool TTS model which can be fine tuned with Unsloth."

speaker_id = 0
# Another equivalent way to prepare the inputs
conversation = [
    {"role": str(speaker_id), "content": [{"type": "text", "text": text}]},
]
inputs = processor.apply_chat_template(
    conversation,
    tokenize=True,
    return_dict=True,
).to("cuda")
audio_values = model.generate(
    input_ids=inputs["input_ids"],
    attention_mask=inputs.get("attention_mask"),
    max_new_tokens=125, # 125 tokens is 10 seconds of audio, for longer speech increase this
    # play with these parameters to tweak results
    # depth_decoder_top_k=0,
    # depth_decoder_top_p=0.9,
    # depth_decoder_do_sample=True,
    # depth_decoder_temperature=0.9,
    # top_k=0,
    # top_p=1.0,
    # temperature=0.9,
    # do_sample=True,
    #########################################################
    output_audio=True
)
audio = audio_values[0].to(torch.float32).cpu().numpy()
sf.write("example_without_context_2.wav", audio, 24000)

"""#### Voice and style consistency

Sesame CSM's power comes from providing audio context for each speaker. Let's pass a sample utterance from our dataset to ground speaker identity and style.
"""

speaker_id = 0

utterance = raw_ds[5]["audio"]["array"]
utterance_text = raw_ds[5]["text"]
text = "Sesame is a super cool TTS model which can be fine tuned with Unsloth."

# CSM will fill in the audio for the last text.
# You can even provide a conversation history back in as you generate new audio

inputs = processor(f"[{speaker_id}]{text}", add_special_tokens=True, return_tensors="pt").to("cuda")
audio_values = model.generate(
    input_ids=inputs["input_ids"],
    attention_mask=inputs.get("attention_mask"),
    max_new_tokens=125, # 125 tokens is 10 seconds of audio, for longer speech increase this
    # play with these parameters to tweak results
    # depth_decoder_top_k=0,
    # depth_decoder_top_p=0.9,
    # depth_decoder_do_sample=True,
    # depth_decoder_temperature=0.9,
    # top_k=0,
    # top_p=1.0,
    # temperature=0.9,
    # do_sample=True,
    #########################################################
    output_audio=True
)
audio = audio_values[0].to(torch.float32).cpu().numpy()
sf.write("example_with_context_1.wav", audio, 24000)

# Example 2
utterance = raw_ds[4]["audio"]["array"]
utterance_text = raw_ds[4]["text"]

conversation = [
    {"role": str(speaker_id), "content": [{"type": "text", "text": utterance_text},{"type": "audio", "path": utterance}]},
    {"role": str(speaker_id), "content": [{"type": "text", "text": text}]},
]

text = "We just finished fine tuning a text to speech model... and it's pretty good!"

conversation = [
    {"role": str(speaker_id), "content": [{"type": "text", "text": utterance_text},{"type": "audio", "path": utterance}]},
    {"role": str(speaker_id), "content": [{"type": "text", "text": text}]},
]

inputs = processor.apply_chat_template(
        conversation,
        tokenize=True,
        return_dict=True,
    ).to("cuda")
audio_values = model.generate(
    input_ids=inputs["input_ids"],
    attention_mask=inputs.get("attention_mask"),
    max_new_tokens=125, # 125 tokens is 10 seconds of audio, for longer text increase this
    # play with these parameters to tweak results
    # depth_decoder_top_k=0,
    # depth_decoder_top_p=0.9,
    # depth_decoder_do_sample=True,
    # depth_decoder_temperature=0.9,
    # top_k=0,
    # top_p=1.0,
    # temperature=0.9,
    # do_sample=True,
    #########################################################
    output_audio=True
)
audio = audio_values[0].to(torch.float32).cpu().numpy()
sf.write("example_with_context_2.wav", audio, 24000)

"""And we're done! If you have any questions on Unsloth, we have a [Discord](https://discord.gg/unsloth) channel! If you find any bugs or want to keep updated with the latest LLM stuff, or need help, join projects etc, feel free to join our Discord!

Some other links:
1. Train your own reasoning model - Llama GRPO notebook [Free Colab](https://colab.research.google.com/github/unslothai/notebooks/blob/main/nb/Llama3.1_(8B)-GRPO.ipynb)
2. Saving finetunes to Ollama. [Free notebook](https://colab.research.google.com/github/unslothai/notebooks/blob/main/nb/Llama3_(8B)-Ollama.ipynb)
3. Llama 3.2 Vision finetuning - Radiography use case. [Free Colab](https://colab.research.google.com/github/unslothai/notebooks/blob/main/nb/Llama3.2_(11B)-Vision.ipynb)
6. See notebooks for DPO, ORPO, Continued pretraining, conversational finetuning and more on our [documentation](https://docs.unsloth.ai/get-started/unsloth-notebooks)!

<div class="align-center">
  <a href="https://unsloth.ai"><img src="https://github.com/unslothai/unsloth/raw/main/images/unsloth%20new%20logo.png" width="115"></a>
  <a href="https://discord.gg/unsloth"><img src="https://github.com/unslothai/unsloth/raw/main/images/Discord.png" width="145"></a>
  <a href="https://docs.unsloth.ai/"><img src="https://github.com/unslothai/unsloth/blob/main/images/documentation%20green%20button.png?raw=true" width="125"></a>

  Join Discord if you need help + ⭐️ <i>Star us on <a href="https://github.com/unslothai/unsloth">Github</a> </i> ⭐️
</div>

"""

test_batch_inference.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Regression tests for batch_inference.py JSON input handling:

- Items missing the required "text" field must raise a clear ValueError
  (previously a bare KeyError: 'text' aborted the whole batch).
- An explicit JSON null "output" must fall back to the default filename
  (previously `output_path / None` raised TypeError mid-batch).

Heavy dependencies (torch, unsloth, transformers, peft, datasets, soundfile,
tqdm) are stubbed via sys.modules so the real module can be imported and
generate_speech_batch driven directly with a mock model/processor -- the bugs
lived in the per-item parsing, before any real model work.
"""

import contextlib
import os
import sys
import types
from unittest.mock import MagicMock

import pytest


def _stub(name, **attrs):
    module = types.ModuleType(name)
    for key, value in attrs.items():
        setattr(module, key, value)
    sys.modules[name] = module


_stub("torch",
      cuda=types.SimpleNamespace(is_available=lambda: False),
      no_grad=lambda: contextlib.nullcontext(),
      float32="float32")
_stub("soundfile", write=lambda *a, **k: None)
_stub("datasets", load_dataset=lambda *a, **k: None, Audio=lambda *a, **k: None)
_stub("unsloth", FastModel=object)
_stub("transformers", CsmForConditionalGeneration=object)
_stub("peft", PeftModel=object)
_stub("tqdm", tqdm=lambda it, desc=None: it)

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import batch_inference  # noqa: E402


def _run_batch(texts, tmp_path):
    """Run generate_speech_batch over texts, returning the sf.write call paths."""
    written = []
    sys.modules["soundfile"].write = lambda path, *a, **k: written.append(str(path))
    batch_inference.generate_speech_batch(
        model=MagicMock(),
        processor=MagicMock(),
        texts=texts,
        output_dir=str(tmp_path),
    )
    return written


def test_missing_text_raises_clear_error(tmp_path):
    """Item without a 'text' key must fail loudly with a clear message."""
    with pytest.raises(ValueError, match="text"):
        _run_batch([{"speaker_id": 0, "output": "hello.wav"}], tmp_path)


def test_empty_text_raises_clear_error(tmp_path):
    """Item with empty 'text' must also fail loudly."""
    with pytest.raises(ValueError, match="text"):
        _run_batch([{"text": ""}], tmp_path)


def test_null_output_falls_back_to_default_filename(tmp_path):
    """Explicit JSON null 'output' must use the generated default filename."""
    written = _run_batch([{"text": "Hello world", "output": None}], tmp_path)
    assert len(written) == 1
    assert os.path.basename(written[0]).startswith("output_")
    assert written[0].endswith(".wav")


def test_valid_item_still_works(tmp_path):
    """A well-formed item still generates to its explicit output path."""
    written = _run_batch([{"text": "Hi", "speaker_id": 0, "output": "hi.wav"}], tmp_path)
    assert len(written) == 1
    assert written[0].endswith("hi.wav")

example_inputs.json

[
    {
        "text": "We just finished fine tuning a text to speech model... and it's pretty good!",
        "speaker_id": 0,
        "output": "example_without_context_1.wav",
        "comment": "Example without context - generates speech without voice reference"
    },
    {
        "text": "Sesame is a super cool TTS model which can be fine tuned with Unsloth.",
        "speaker_id": 0,
        "output": "example_without_context_2.wav",
        "comment": "Example without context - generates speech without voice reference"
    },
    {
        "text": "Sesame is a super cool TTS model which can be fine tuned with Unsloth.",
        "speaker_id": 0,
        "dataset_context_idx": 3,
        "output": "example_with_context_1.wav",
        "comment": "Example with context - uses dataset index 3 for voice consistency (same as training script)"
    },
    {
        "text": "We just finished fine tuning a text to speech model... and it's pretty good!",
        "speaker_id": 0,
        "dataset_context_idx": 4,
        "output": "example_with_context_2.wav",
        "comment": "Example with context - uses dataset index 4 for voice consistency (same as training script)"
    }
]