Thoroughly explore /home/ubuntu/training/maya-tts-llama32-1b-pretraining to understand the full training setup. I need to understand:

1. What is the model architecture? (base model, any modifications, custom layers, audio codecs used)
2. What is the task? (TTS = text-to-speech, but how exactly - what are inputs/outputs, tokenization strategy for audio)
3. Training schedule and strategy (stages, learning rate, epochs, batch size, distributed config)
4. Data pipeline - how is data loaded, what format, what preprocessing
5. How checkpoints are saved and what format they're in
6. Any existing inference scripts or evaluation code
7. The model's config files and any custom tokenizer setup
8. Any README or documentation

Read ALL key files - configs, training scripts, model definitions, data loaders, tokenizer code, evaluation scripts. Be very thorough - read the full content of important files, not just headers. Pay special attention to:
- How audio is encoded/decoded (codec, codebook structure)
- The model's forward pass and generation logic
- Any special tokens or vocabulary extensions
- Checkpoint saving format and what's included

Return a comprehensive summary with file paths and line numbers for key components.


I'll thoroughly explore the training setup. Let me start by examining the directory structure and key files.


Bash

find /home/ubuntu/training/maya-tts-llama32-1b-pretraining -type f -name "*.py" -o -name "*.json" -o -name "*.md" -o -name "*.txt" -o -name "*.yaml" -o -name "*.yml" | head -50


Let me continue listing files and get a fuller picture:


find /home/ubuntu/training/maya-tts-llama32-1b-pretraining -type f \( -name "*.json" -o -name "*.md" -o -name "*.yaml" -o -name "*.yml" \) | sort


Now let me read the key documentation and configuration files:


     1→# Inworld Text-To-Speech Trainer
     2→
     3→[![Lint Code](https://github.com/inworld-ai/tts/actions/workflows/linter.yaml/badge.svg)](https://github.com/inworld-ai/tts/actions/workflows/linter.yaml)
     4→[![Python 3.10](https://img.shields.io/badge/python-3.10+-blue.svg)](https://www.python.org/downloads/)
     5→[![CUDA Support](https://img.shields.io/badge/CUDA-12.4%20%7C%2012.8-green.svg)](https://developer.nvidia.com/cuda-toolkit)
     6→[![PyTorch](https://img.shields.io/badge/PyTorch-2.6%20%7C%202.7-orange.svg)](https://pytorch.org/)
     7→[![License: MIT](https://img.shields.io/badge/License-MIT-yellow.svg)](https://opensource.org/licenses/MIT)
     8→[![ArXiv](https://img.shields.io/badge/arXiv-Paper-<COLOR>.svg)](https://arxiv.org/abs/2507.21138)
     9→[![Playground](https://img.shields.io/badge/Playground-Try%20Now-blue.svg)](https://inworld.ai/tts)
    10→
    11→This repository contains the training and modeling code used to create [Inworld TTS-1 and TTS-1-Max models](https://inworld.ai/blog/introducing-inworld-tts).
    12→
    13→You can use this code to pre-train, fine-tune, or align with RL your arbitrary SpeechLM-based TTS model, no matter if you're using a single GPU machine or a multi-GPU cluster.
    14→
    15→## Inworld TTS
    16→
    17→* [Playground](https://inworld.ai/tts)
    18→* [Examples](https://inworld-ai.github.io/tts/)
    19→* [Technical Report](https://arxiv.org/abs/2507.21138)
    20→
    21→## TTS Trainer Features
    22→
    23→* **Modeling**: SpeechLM and 1D audio-codecs
    24→* **Distributed Training**: DDP, DeepSpeed and FSDP for training arbitrary SpeechLM
    25→* **Data Pipeline**: Ready-to-use scripts to vectorize and prepare your audio data for training
    26→
    27→## Requirements
    28→
    29→The code is only tested on Ubuntu 22.04.
    30→
    31→| Component | Version | Notes |
    32→|-----------|---------|-------|
    33→| **Python** | 3.10 | Required for all features |
    34→| **CUDA** | 12.4 or 12.8 | Auto-detected on Linux when possible, or set with `CUDA_VERSION` |
    35→| **PyTorch** | 2.6 (CUDA 12.4) or 2.7 (CUDA 12.8) | Installed from `uv.lock` |
    36→
    37→## Quick Start
    38→
    39→### Prerequisites
    40→
    41→This project depends on **Python 3.10** and **uv** for package management.
    42→
    43→#### Install uv
    44→
    45→Install [uv](https://docs.astral.sh/uv/) for fast Python package management:
    46→
    47→```bash
    48→curl -LsSf https://astral.sh/uv/install.sh | sh
    49→```
    50→
    51→### One-Command Setup
    52→
    53→`make install` creates `.venv`, activates it inside the setup script, installs dependencies from `uv.lock` when present, and then installs the matching flash-attn prebuilt wheel on Linux.
    54→
    55→**Auto-detect CUDA on Linux:**
    56→
    57→```bash
    58→make install
    59→```
    60→
    61→**Specify CUDA version explicitly if detection is unavailable or wrong:**
    62→
    63→```bash
    64→# For CUDA 12.4 + PyTorch 2.6
    65→make install CUDA_VERSION=12.4
    66→
    67→# For CUDA 12.8 + PyTorch 2.7
    68→make install CUDA_VERSION=12.8
    69→```
    70→
    71→This automatically:
    72→* Creates Python 3.10 virtual environment
    73→* Syncs project and development dependencies from `uv.lock`
    74→* Installs CUDA-optimized PyTorch with the proper [flash attention](https://github.com/Dao-AILab/flash-attention) implementation
    75→* Sets up the local development toolchain (`pytest`, `ruff`, `pre-commit`)
    76→
    77→### Optional: Install vLLM for Faster Inference
    78→
    79→For faster inference and RLHF training, install vLLM:
    80→
    81→```bash
    82→# Activate virtual environment first
    83→source .venv/bin/activate
    84→
    85→# Install vLLM
    86→uv pip install vllm==0.10.0
    87→```
    88→
    89→**Note:** vLLM is required for:
    90→- Fast inference with `--use_vllm=true` flag
    91→- RLHF training (serves the base model during training)
    92→
    93→## Training Example
    94→
    95→In order to train a SpeechLM you need to first vectorize your audio-data into the audio-codes.
    96→Combined together with the transcript, the model learns how being conditioned on it generates the audio.
    97→Below example shows how to get started with a simple SFT training.
    98→
    99→### 1. Data Preparation
   100→
   101→Process your raw audio dataset into a JSONL file where each line contains a sample with the following format:
   102→
   103→```json
   104→{
   105→  "transcript": "Then they would swiftly dart at their prey and bear it to the ground.",
   106→  "language": "en",
   107→  "wav_path": "/path/to/audio.wav",
   108→  "duration": 3.42,
   109→  "sample_rate": 24000
   110→}
   111→```
   112→
   113→**Required fields:**
   114→* `transcript`: Text transcription of the audio
   115→* `language`: Language code (e.g., "en" for English)
   116→* `wav_path`: Absolute path to the audio file
   117→* `duration`: Audio duration in seconds
   118→* `sample_rate`: Audio sample rate in Hz
   119→
   120→**Example dataset:** You can reference the [LibriTTS dataset](https://huggingface.co/datasets/mythicinfinity/libritts) which contains ~585 hours of English speech from 2, 456 speakers at 24kHz sampling rate.
   121→
   122→**Sample files:** We provide real example data from LibriTTS in this repository:
   123→* [`./example/configs/samples.jsonl`](./example/configs/samples.jsonl) - 100 real LibriTTS samples with proper JSONL format
   124→* `./example/wavs/` - Corresponding audio files (audio_1.wav through audio_100.wav)
   125→
   126→This gives you a working example to test the data vectorization and training pipeline with actual audio data.
   127→
   128→### 2. Data Vectorization
   129→
   130→Vectorize audio data using codec's encoder (We also made the codec compatible with [xcodec2](https://huggingface.co/HKUSTAudio/xcodec2/tree/main/ckpt), so you can use the publicly available checkpoint if you prefer not to train the codec from scratch):
   131→
   132→**Test with provided samples:**
   133→```bash
   134→WANDB_PROJECT="your_project" \
   135→torchrun --nproc_per_node 8 ./tools/data/data_vectorizer.py \
   136→    --codec_model_path=/path/to/codec/model.pt \
   137→    --batch_size=16 \
   138→    --dataset_path=./example/configs/samples.jsonl \
   139→    --output_dir=/path/to/output_directory \
   140→    --use_wandb \
   141→    --run_name=test_vectorization
   142→```
   143→
   144→**With your own dataset:**
   145→```bash
   146→WANDB_PROJECT="your_project" \
   147→torchrun --nproc_per_node 8 ./tools/data/data_vectorizer.py \
   148→    --codec_model_path=/path/to/codec/model.pt \
   149→    --batch_size=16 \
   150→    --dataset_path=/path/to/your_data.jsonl \
   151→    --output_dir=/path/to/output_directory \
   152→    --use_wandb \
   153→    --run_name=vectorization_run
   154→```
   155→
   156→After vectorization completes, you'll have multiple shard files in your output directory like below:
   157→
   158→```
   159→train_codes_{0..n}.npy         # Vectorized audio codes for training
   160→train_codes_index_{0..n}.npy   # Index mappings for training codes
   161→train_samples_{0..n}.jsonl     # Training sample metadata
   162→val_codes_{0..n}.npy           # Vectorized audio codes for validation
   163→val_codes_index_{0..n}.npy     # Index mappings for validation codes
   164→val_samples_{0..n}.jsonl       # Validation sample metadata
   165→```
   166→
   167→Feel free to customize [filtering logic](./tts/data/data_utils.py) to implement your custom filtering logic.
   168→
   169→### 3. Merge Shards
   170→
   171→Combine vectorized shards into unified dataset to save space:
   172→
   173→```bash
   174→python tools/data/data_merger.py \
   175→    --dataset_path /path/to/your/vectorized_dataset \
   176→    --remove_shards
   177→```
   178→
   179→After merging, your dataset folder will contain:
   180→
   181→```
   182→train_codes.npy        # Merged training codes
   183→train_codes_index.npy  # Merged training code indices
   184→train_samples.jsonl    # Merged training samples
   185→val_codes.npy          # Merged validation codes
   186→val_codes_index.npy    # Merged validation code indices
   187→val_samples.jsonl      # Merged validation samples
   188→```
   189→
   190→### 3-1. Text Pretraining Corpus
   191→
   192→If you want to mix text-only PT data into training, use the registry-driven text corpus pipeline. It keeps a canonical document store for auditability, exports compiler-ready JSONL shards, and compiles them into the flat memmaps expected by the current trainer.
   193→
   194→**Registry example:** `./example/configs/text_corpus_sources.json`
   195→
   196→```bash
   197→python tools/data/build_text_corpus.py \
   198→    --registry_path=./example/configs/text_corpus_sources.json \
   199→    --output_dir=/path/to/text_corpus \
   200→    --tokenizer_name=meta-llama/Llama-3.2-1B-Instruct \
   201→    --compile
   202→```
   203→
   204→By default this collects the enabled Wave 1 backbone sources from the registry. You can target specific sources or waves, and add local mirrors for later-wave corpora:
   205→
   206→```bash
   207→python tools/data/build_text_corpus.py \
   208→    --registry_path=./example/configs/text_corpus_sources.json \
   209→    --output_dir=/path/to/text_corpus \
   210→    --waves=1,2 \
   211→    --source_ids=fineweb_edu,sangraha_verified,indiccorp_v2 \
   212→    --local_input_root=/path/to/local_text_mirrors \
   213→    --include_review_required=false \
   214→    --compile
   215→```
   216→
   217→The pipeline writes:
   218→
   219→* `canonical/` document-level records with source, language, license, and split metadata
   220→* `exports/` JSONL shards with a `text` field for each wave view
   221→* `compiled/` `train_pretraining_tokens.npy` and `val_pretraining_tokens.npy` views built by `tools/data/build_pretraining_data.py`
   222→* `reports/` per-source and per-wave corpus statistics
   223→
   224→For large-scale tokenization (25B+ tokens), use the parallel tokenizer which saturates all available CPUs:
   225→
   226→```bash
   227→OPENBLAS_NUM_THREADS=1 OMP_NUM_THREADS=1 TOKENIZERS_PARALLELISM=false \
   228→python tools/data/parallel_tokenize.py \
   229→    --input_dir=text_corpus/compiled_25B/train \
   230→    --output_dir=text_corpus/compiled_25B/pt_output \
   231→    --tokenizer_name=meta-llama/Llama-3.2-1B-Instruct \
   232→    --max_seq_len=2048 \
   233→    --workers=200
   234→```
   235→
   236→Use the compiled text PT directory in training configs with the `[text]` suffix:
   237→
   238→```json
   239→{
   240→    "train_weighted_datasets": {
   241→        "/path/to/text_corpus/compiled/all_approved[text]": 0.1
   242→    }
   243→}
   244→```
   245→
   246→To bootstrap text-to-speech alignment during PT, you can also compile a dedicated paired PT subset from a transcript-bearing merged triplet. The source must preserve `samples.jsonl` rows compatible with `Sample.from_json()` rather than the minimal `segment_id`/`token_count` metadata export:
   247→
   248→```bash
   249→python tools/data/build_pretraining_data.py \
   250→    --mode=paired \
   251→    --input_dir=/path/to/transcript_merged_dataset \
   252→    --output_dir=/path/to/paired_pt_subset \
   253→    --allowed_languages=en,hi \
   254→    --min_dnsmos_score=3.0 \
   255→    --min_sample_rate=16000 \
   256→    --min_audio_duration=0.5 \
   257→    --shuffle=true \
   258→    --split=train
   259→```
   260→
   261→Use the paired PT artifact directory in training configs with the `[paired]` suffix:
   262→
   263→```json
   264→{
   265→    "train_weighted_datasets": {
   266→        "/path/to/paired_pt_subset[paired]": 0.1,
   267→        "/path/to/text_corpus/compiled/all_approved[text]": 0.1
   268→    }
   269→}
   270→```
   271→
   272→### 4. Configuration
   273→
   274→Create training config ( `./example/configs/sft.json` ). Below shows key configuration sections - see the full configuration file at `./example/configs/sft.json` for all available options:
   275→
   276→```json
   277→{
   278→    "training": {
   279→        "seed": 777,
   280→        "logging_steps": 150,
   281→        "eval_steps": 300,
   282→        "learning_rate": 1e-04,
   283→        "batch_size": 4,
   284→        "precision": "bf16",
   285→        "strategy": "ddp"
   286→    },
   287→    "modeling": {
   288→        "parameters": {
   289→            "codebook_size": 65536,
   290→            "max_seq_len": 2048,
   291→            "model_name": "meta-llama/Llama-3.2-1B-Instruct",
   292→            "enable_text_normalization": true
   293→        }
   294→    },
   295→    "checkpointing": {
   296→        "save_steps": 10000,
   297→        "keep_only_last_n_checkpoints": 10
   298→    },
   299→    "train_weighted_datasets": {
   300→        "/path/to/your/vectorized_dataset": 1.0
   301→    },
   302→    "val_weighted_datasets": {
   303→        "/path/to/your/vectorized_dataset": 1.0
   304→    },
   305→    "dataset": {
   306→        "enable_rlhf_training": false
   307→    }
   308→}
   309→```
   310→
   311→**Important**:
   312→* Update dataset paths to point to your vectorized data directory
   313→* This shows only key parameters - refer to `./example/configs/sft.json` for the complete configuration with all available options
   314→* To resume from a checkpoint: Add `"checkpoint_file_to_resume_from": "/path/to/your/checkpoint.pt"` to the `checkpointing` section
   315→* If `save_intermediate_generations` is enabled, set real `codec_encoder_path`, `codec_decoder_path`, and `canary_prompts`; the built-in defaults are placeholders and the validator now fails fast when they do not exist
   316→
   317→### 5. SFT Training
   318→
   319→
   320→```bash
   321→fabric run --devices=$NUM_GPU tts/training/main.py \
   322→    --config_path=./example/configs/sft.json \
   323→    --use_wandb \
   324→    --run_name=my_tts_training
   325→```
   326→
   327→After training completes, you'll find the trained model at `./experiments/my_tts_training/final_model.pt` along with model and tokenizer configuration files.
   328→
   329→**Additional options:**
   330→* `--dry_run`: Test pipeline without training
   331→* `--compile_model`: Enable torch.compile optimization (works well only if all your batch' samples have the same length)
   332→
   333→### 6. RLHF Training
   334→
   335→After completing SFT training, you can further improve your model using Reinforcement Learning from Human Feedback (RLHF). This process uses reward functions to align the model with human preferences.
   336→
   337→#### 6-1: Convert Model for Serving
   338→
   339→First, convert your trained SFT model to a serving-friendly format:
   340→
   341→```bash
   342→python tools/serving/convert_checkpoint.py \
   343→  --checkpoint_path=/path/to/your/sft_model.pt \
   344→  --output_path=/path/to/serving/model
   345→```
   346→
   347→#### 6-2: RLHF Configuration
   348→
   349→Create an RLHF training config (`./example/configs/rlhf.json`). Key sections include:
   350→
   351→```json
   352→{
   353→    "training": {
   354→        "seed": 777,
   355→        "logging_steps": 50,
   356→        "eval_steps": 100,
   357→        "learning_rate": 5e-06,
   358→        "batch_size": 2,
   359→        "precision": "bf16",
   360→        "strategy": "ddp"
   361→    },
   362→    "rlhf_training": {
   363→        "base_model_dir": "/path/to/your/serving/model",
   364→        "top_p": 0.9,
   365→        "top_k": 75,
   366→        "repetition_penalty": 1.1,
   367→        "temperature": 1.1,
   368→        "num_generations": 8,
   369→        "max_prompt_length": 1024,
   370→        "max_completion_length": 1024,
   371→        "min_completion_length": 25,
   372→        "use_vllm": true,
   373→        "reward_funcs": ["WERRewardFunc"],
   374→        "reward_weights": [1.0],
   375→        "save_completions_steps": 50,
   376→        "per_device_train_batch_size": 4,
   377→        "num_iterations": 1,
   378→        "scale_rewards": false,
   379→        "kl_beta": 0.04
   380→    },
   381→    "train_weighted_datasets": {
   382→        "/path/to/your/vectorized_dataset": 1.0
   383→    },
   384→    "val_weighted_datasets": {
   385→        "/path/to/your/vectorized_dataset": 1.0
   386→    },
   387→    "dataset": {
   388→        "enable_rlhf_training": true,
   389→        "allowed_languages": ["en", "es", "fr", "de", "it", "pt", "ru", "zh", "ko", "ja", "nl", "pl"]
   390→    }
   391→}
   392→```
   393→
   394→**Key RLHF parameters:**
   395→- `base_model_dir`: Path to your converted serving model
   396→- `reward_funcs`: List of reward functions (e.g., "WERRewardFunc" for Word Error Rate)
   397→- `num_generations`: Number of completions to generate per prompt
   398→- `kl_beta`: KL divergence penalty weight
   399→- `use_vllm`: Enable vLLM for faster inference during training
   400→
   401→#### 6-3: Multi-Node RLHF Training
   402→
   403→RLHF training requires two components running simultaneously:
   404→1. **Training node**: Runs the RLHF training loop
   405→2. **vLLM node**: Serves the base model for generation
   406→
   407→**Update the script configuration:**
   408→
   409→Edit `./tts/training/rlhf/run_rlhf_combine.sh` and set:
   410→```bash
   411→# Update this path to your converted serving model
   412→VLLM_MODEL_PATH="/path/to/your/serving/model"
   413→```
   414→
   415→**Launch RLHF training:**
   416→
   417→```bash
   418→# Submit SLURM job with 2 nodes
   419→sbatch --nodes=2 --gpus-per-node=8 --cpus-per-task=114 \
   420→  --nodelist=gpu-node1,gpu-node2 --partition=compute \
   421→  ./tts/training/rlhf/run_rlhf_combine.sh
   422→```
   423→
   424→
   425→
   426→### 7. Monitoring
   427→
   428→Track progress via:
   429→* **Weights & Biases**: Loss curves and training metrics
   430→* **Checkpoints**: Saved every `save_steps` iterations
   431→* **Console logs**: Real-time training information
   432→
   433→## Inference
   434→
   435→Once you have a trained model, you can use it for inference to generate speech from text and an audio prompt.
   436→
   437→Use the provided inference script for easy speech generation:
   438→
   439→**Standard PyTorch inference:**
   440→```bash
   441→python tools/serving/inference.py \
   442→    --model_checkpoint_path /path/to/your/serving/model \
   443→    --audio_encoder_path /path/to/encoder.pt \
   444→    --audio_decoder_path /path/to/decoder.pt \
   445→    --prompt_wav_path /path/to/your_prompt.wav \
   446→    --prompt_transcription "This is what the speaker says in the prompt." \
   447→    --text "Hello, this is a test of the text-to-speech system." \
   448→    --output_path ./audios/output.wav
   449→```
   450→
   451→**vLLM inference (faster):**
   452→```bash
   453→python tools/serving/inference.py \
   454→    --model_checkpoint_path /path/to/your/serving/model \
   455→    --audio_encoder_path /path/to/encoder.pt \
   456→    --audio_decoder_path /path/to/decoder.pt \
   457→    --prompt_wav_path /path/to/your_prompt.wav \
   458→    --prompt_transcription "This is what the speaker says in the prompt." \
   459→    --text "Hello, this is a test of the text-to-speech system." \
   460→    --output_path ./audios/output.wav \
   461→    --use_vllm=true \
   462→    --seed=42
   463→```
   464→
   465→**Required components:**
   466→- Trained model in serving format (converted using `convert_checkpoint.py`)
   467→- Audio encoder checkpoint (codec `.pt`/`.ckpt` file)
   468→- Audio decoder checkpoint (codec `.pt`/`.ckpt` file + `model_config.json` in same directory)
   469→- Audio prompt file (`.wav` format)
   470→- Prompt transcription (text of what's spoken in the audio prompt)
   471→
   472→**Note:** If you don't want to retrain the decoder, you can use the same checkpoint from [xcodec2](https://huggingface.co/HKUSTAudio/xcodec2/tree/main/ckpt) for both encoder and decoder paths. We provided a xcodec2 compatible `model_config.json` file in `./example/codec`.
   473→
   474→## Development
   475→
   476→### Available Commands
   477→
   478→```bash
   479→make help           # Show all available commands
   480→make install        # Install development environment
   481→make test           # Run test suite
   482→make test-coverage  # Run tests with coverage report
   483→make lint           # Run code linting across tts, tools, and tests
   484→make lint-fix       # Auto-fix linting issues across tts, tools, and tests
   485→make version        # Show current version or bump version
   486→make version patch  # Bump patch version (1.0.0 → 1.0.1), also support `minor`, `major`
   487→```
   488→
   489→### Development Workflow
   490→
   491→```bash
   492→git clone https://github.com/inworld-ai/tts.git
   493→cd tts
   494→
   495→# Set up development environment
   496→make install
   497→
   498→# Activate the local environment for direct CLI tools like pre-commit
   499→source .venv/bin/activate
   500→
   501→# Install pre-commit hooks
   502→pre-commit install
   503→
   504→# Make your changes and test
   505→make lint-fix
   506→make test-coverage
   507→```
   508→
   509→## Contributing
   510→
   511→We welcome contributions! Please see our contributing guidelines:
   512→
   513→1. **Fork the repository**
   514→2. **Create a feature branch**: `git checkout -b feature/amazing-feature`
   515→3. **Run tests**: `make test`
   516→4. **Run linting**: `make lint-fix`
   517→5. **Commit changes**: `git commit -m 'Add amazing feature'`
   518→6. **Push to branch**: `git push origin feature/amazing-feature`
   519→7. **Open a Pull Request**
   520→
   521→## Acknowledgments
   522→
   523→* Meta AI team for open-sourcing [LLaMA LLMs](https://www.llama.com/)
   524→* The PyTorch and Hugging Face communities
   525→* Codec architecture inspired by [Llasa](https://arxiv.org/abs/2502.04128)
   526→
   527→## Support
   528→
   529→* **Bug Reports**: [GitHub Issues](https://github.com/inworld-ai/tts/issues)
   530→* **General Questions**: For general inquiries and support, please [email us](mailto:opensource@inworld.ai)
   531→


Error: File content (11287 tokens) exceeds maximum allowed tokens (10000). Use offset and limit parameters to read specific portions of the file, or search for specific content instead of reading the whole file.


     1→# Maya TTS Training Guide (LlamaFactory)
     2→
     3→## Overview
     4→
     5→This repo trains a multimodal TTS LLM based on `Llama-3.2-1B-Instruct` with an
     6→expanded vocabulary (`193,856`) that includes:
     7→
     8→- codec tokens: `<|s_0|>` ... `<|s_65535|>`
     9→- speech boundaries: `<|speech_start|>`, `<|speech_end|>`
    10→- control tokens for prompt formatting
    11→
    12→Training stack is now **LlamaFactory** with a token-id-native mixed PT stream.
    13→
    14→## Data Mixture
    15→
    16→Mixed pretraining uses three sources:
    17→
    18→- Audio PT: bounded codec stream windows (`{split}_pretraining_codes.npy`)
    19→- Text PT: token windows (`{split}_pretraining_tokens.npy`)
    20→- Paired PT: transcript-to-speech triplets (`{split}_codes.npy` + index + jsonl)
    21→
    22→Default weight ratio is `84:10:6` (audio:text:paired).
    23→
    24→## Quick Start
    25→
    26→### 1) Install
    27→
    28→```bash
    29→pip install -e .
    30→git clone --depth 1 https://github.com/hiyouga/LLaMA-Factory.git llamafactory_ref
    31→cd llamafactory_ref
    32→pip install -e .
    33→pip install -r requirements/deepspeed.txt
    34→cd ..
    35→```
    36→
    37→### 2) Build seed checkpoint
    38→
    39→```bash
    40→python3 tools/model/create_maya_seed_checkpoint.py \
    41→  --model_name meta-llama/Llama-3.2-1B-Instruct \
    42→  --output_dir checkpoints/hf_seed/maya-llama3.2-1b-seed
    43→```
    44→
    45→### 3) Smoke run
    46→
    47→```bash
    48→bash llamafactory_experiment/scripts/run_pt_smoke.sh
    49→```
    50→
    51→### 4) Full pretraining run
    52→
    53→```bash
    54→bash llamafactory_experiment/scripts/run_pt_full.sh
    55→```
    56→
    57→## Key Paths and Env Overrides
    58→
    59→The PT scripts use these defaults (override via env vars):
    60→
    61→- `MODEL_PATH=checkpoints/hf_seed/maya-llama3.2-1b-seed`
    62→- `AUDIO_PT_DIR=pretrain-encoded/pt_output`
    63→- `TEXT_PT_DIR=text_corpus/compiled_25B/pt_output`
    64→- `PAIRED_PT_DIR=finalsftdata/sft4b`
    65→
    66→Important knobs:
    67→
    68→- `MAYA_PT_MIX_WEIGHTS=84,10,6`
    69→- `MAYA_TRAIN_SAMPLES=200000000`
    70→- `NPROC_PER_NODE=8`
    71→- `MASTER_PORT=29610`
    72→
    73→## Notes
    74→
    75→- PT uses `stage: sft` in LlamaFactory intentionally so paired samples can keep
    76→  prompt-masked labels (`-100`) exactly.
    77→- Data is consumed as token ids directly. No decode/re-encode roundtrip.
    78→- If you want strict dependency guards, set `DISABLE_VERSION_CHECK=0`.
    79→
    80→## Pre-SFT Loss Check
    81→
    82→Use the HF-only evaluator to inspect source losses before starting SFT:
    83→
    84→```bash
    85→python3 tools/analysis/eval_maya_pretrain_losses.py \
    86→  --model_path checkpoints/hf_final/step-545243-hf-fixed \
    87→  --audio_pt_dir pretrain-encoded/pt_output \
    88→  --text_pt_dir text_corpus/compiled_25B/pt_output \
    89→  --paired_pt_dir finalsftdata/sft4b \
    90→  --split train \
    91→  --audio_samples 512 \
    92→  --text_samples 512 \
    93→  --paired_samples 512 \
    94→  --batch_size 4 \
    95→  --device cuda
    96→```
    97→


Now let me read the key source files for architecture and modeling:


     1→import contextlib
     2→import json
     3→import os
     4→import time
     5→
     6→import lightning.fabric as lightning_fabric
     7→import torch
     8→import transformers
     9→from absl import logging
    10→
    11→from tts.core import constants, lora
    12→from tts.data import caching
    13→from tts.utils import configuration
    14→
    15→# Use FlashAttention-3 (Hopper-optimized) when available, fall back to FA2.
    16→try:
    17→    from transformers.utils import is_flash_attn_3_available
    18→    _ATTN_IMPLEMENTATION = "flash_attention_3" if is_flash_attn_3_available() else "flash_attention_2"
    19→except ImportError:
    20→    _ATTN_IMPLEMENTATION = "flash_attention_2"
    21→
    22→
    23→def _str_to_torch_dtype(dtype: str) -> torch.dtype:
    24→    """Converts a string to a torch dtype."""
    25→    str_to_dtype = {
    26→        "fp32": torch.float32,
    27→        "fp16": torch.float16,
    28→        "bf16": torch.bfloat16,
    29→        "int8": torch.int8,
    30→    }
    31→    if dtype not in str_to_dtype:
    32→        raise ValueError(f"Unknown dtype [{dtype}].")
    33→    return str_to_dtype[dtype]
    34→
    35→
    36→def _construct_model(
    37→    model_name: str, torch_dtype: torch.dtype | str, vocab_size: int, cache_dir: str
    38→) -> torch.nn.Module:
    39→    model = transformers.AutoModelForCausalLM.from_pretrained(
    40→        model_name,
    41→        torch_dtype=torch_dtype,
    42→        attn_implementation=_ATTN_IMPLEMENTATION,
    43→        cache_dir=cache_dir,
    44→    )
    45→
    46→    # LLAMA_ASSUMPTION: `.model.embed_tokens` is the LLaMA-family attribute path.
    47→    # Other architectures (e.g. Qwen) may use `.transformer.wte` or similar.
    48→    current_model_vocab_size = model.model.embed_tokens.weight.size(0)
    49→    logging.info("[%s] vocab size: [%d].", model_name, current_model_vocab_size)
    50→
    51→    if current_model_vocab_size != vocab_size:
    52→        logging.info(
    53→            "Resizing model from [%d] to [%d]...", current_model_vocab_size, vocab_size
    54→        )
    55→        model.resize_token_embeddings(vocab_size)
    56→        model.vocab_size = vocab_size
    57→        if vocab_size != model.model.embed_tokens.weight.size(0):
    58→            raise ValueError("Embedding size mismatch!")
    59→        if vocab_size != model.lm_head.weight.size(0):
    60→            raise ValueError("LM head size mismatch!")
    61→
    62→    return model
    63→
    64→
    65→def load_tokenizer_config_and_model(
    66→    checkpoint_path: str,
    67→) -> tuple[
    68→    transformers.PreTrainedTokenizerBase,
    69→    dict,
    70→    torch.nn.Module,
    71→    configuration.LoraConfig | None,
    72→]:
    73→    """Loads tokenizer, config, and model from a checkpoint.
    74→
    75→    This is a common pattern used across multiple scripts that load models
    76→    from checkpoints. It handles loading the tokenizer, training config,
    77→    extracting LoRA config if present, and loading the model.
    78→
    79→    Args:
    80→        checkpoint_path: Path to the checkpoint file or directory.
    81→        set_pad_token: Whether to set the tokenizer's pad_token to eos_token.
    82→
    83→    Returns:
    84→        Tuple of (tokenizer, config, model, lora_config).
    85→    """
    86→    checkpoint_dir = (
    87→        checkpoint_path
    88→        if os.path.isdir(checkpoint_path)
    89→        else os.path.dirname(checkpoint_path)
    90→    )
    91→
    92→    # Load tokenizer
    93→    start_time = time.time()
    94→    tokenizer = transformers.AutoTokenizer.from_pretrained(checkpoint_dir)
    95→    logging.info(
    96→        "Tokenizer with size: %d loaded in %.2f seconds.",
    97→        len(tokenizer),
    98→        time.time() - start_time,
    99→    )
   100→
   101→    # Load training config
   102→    training_config_file = os.path.join(checkpoint_dir, constants.CONFIG_FILE_NAME)
   103→    with open(training_config_file) as f:
   104→        config = json.load(f)
   105→
   106→    # Extract LoRA config if present
   107→    lora_config = None
   108→    if "lora" in config and config["lora"] is not None:
   109→        lora_config = configuration.LoraConfig.from_dict(config["lora"])
   110→
   111→    # Load model
   112→    start_time = time.time()
   113→    model = load_model_from_checkpoint(
   114→        model_name=config["modeling"]["parameters"]["model_name"],
   115→        vocab_size=len(tokenizer),
   116→        checkpoint_path=checkpoint_path,
   117→        precision=config["training"]["precision"],
   118→        lora_config=lora_config,
   119→    )
   120→    logging.info("Model loaded in %.2f seconds.", time.time() - start_time)
   121→
   122→    return tokenizer, config, model, lora_config
   123→
   124→
   125→def build_model(
   126→    fabric: lightning_fabric.Fabric,
   127→    model_name: str,
   128→    precision: str,
   129→    vocab_size: int,
   130→    deepspeed: bool = False,
   131→    gradient_checkpointing: bool = False,
   132→    use_liger_kernel: bool = False,
   133→) -> torch.nn.Module:
   134→    """Builds a TTS model on top of a pre-trained LLM."""
   135→    cache_dir = caching.get_hf_cache_dir()
   136→    logging.info(
   137→        "Loading pretrained model [%s] using the following HF cache directory: [%s]...",
   138→        model_name,
   139→        cache_dir,
   140→    )
   141→
   142→    # Liger fused kernels must be applied BEFORE model construction so the
   143→    # monkey-patched modules are used when HF instantiates the model.
   144→    if use_liger_kernel:
   145→        try:
   146→            from liger_kernel.transformers import apply_liger_kernel_to_llama
   147→            apply_liger_kernel_to_llama(
   148→                fused_linear_cross_entropy=True,
   149→                rope=True,
   150→                swiglu=True,
   151→                rms_norm=True,
   152→            )
   153→            logging.info("Applied Liger fused kernels (CE, RoPE, SwiGLU, RMSNorm).")
   154→        except ImportError:
   155→            logging.warning("Liger kernel requested but not installed; skipping.")
   156→
   157→    # TODO: debug why using DDP multi-GPU training leaks memory.
   158→    model_init_context = contextlib.nullcontext()
   159→    if fabric.world_size == 1 or deepspeed:
   160→        model_init_context = fabric.init_module()
   161→    with model_init_context:
   162→        model = _construct_model(
   163→            model_name, _str_to_torch_dtype(precision), vocab_size, cache_dir
   164→        )
   165→    if gradient_checkpointing:
   166→        logging.info("Enabling gradient checkpointing...")
   167→        model.gradient_checkpointing_enable()
   168→
   169→    return model
   170→
   171→
   172→def load_model_from_checkpoint(
   173→    model_name: str,
   174→    vocab_size: int,
   175→    checkpoint_path: str,
   176→    precision: str,
   177→    lora_config: dict | None = None,
   178→) -> torch.nn.Module:
   179→    """Loads a model from a checkpoint file to CPU."""
   180→    cache_dir = caching.get_hf_cache_dir()
   181→    torch_dtype = _str_to_torch_dtype(precision)
   182→    logging.info(
   183→        "Loading model from [%s] using the following HF cache directory: [%s]...",
   184→        checkpoint_path,
   185→        cache_dir,
   186→    )
   187→    model = _construct_model(model_name, torch_dtype, vocab_size, cache_dir)
   188→    if lora_config is not None:
   189→        model = lora.apply_lora(model, lora_config)
   190→        logging.info("Applied LoRA adapter [%s] to the base model.", lora_config)
   191→    checkpoint = torch.load(checkpoint_path, weights_only=False)["model"]
   192→    model.load_state_dict(checkpoint)
   193→    return model
   194→


     1→import transformers
     2→from absl import logging
     3→
     4→from tts.core import constants
     5→from tts.data import caching
     6→
     7→# LLAMA_ASSUMPTION: This vocab size is specific to Llama 3.1/3.2 base tokenizers
     8→# (128256 base tokens) plus speech tokens and padding to a multiple of 64.
     9→# Swapping to a different backbone (e.g. Qwen) requires recalculating this.
    10→_EXPECTED_VOCAB_SIZE = 193856
    11→
    12→
    13→def build_tokenizer(
    14→    model_name: str, max_seq_len: int, codebook_size: int, is_lora: bool
    15→) -> transformers.AutoTokenizer:
    16→    """Creates a tokenizer for the model."""
    17→    cache_dir = caching.get_hf_cache_dir()
    18→    logging.info(
    19→        "Loading tokenizer using the following HF cache directory: [%s]...", cache_dir
    20→    )
    21→    tokenizer = transformers.AutoTokenizer.from_pretrained(
    22→        model_name,
    23→        model_max_length=max_seq_len,
    24→        padding_side="right",
    25→        cache_dir=cache_dir,
    26→    )
    27→    tokenizer.pad_token = tokenizer.eos_token
    28→    original_vocab_size = len(tokenizer)
    29→    logging.info(
    30→        "Loaded [%s] tokenizer with [%d] tokens.", model_name, original_vocab_size
    31→    )
    32→
    33→    # LLAMA_ASSUMPTION: If the tokenizer already matches the expected size,
    34→    # we assume it was previously extended (e.g. loaded from a checkpoint).
    35→    expected_vocab_size = _EXPECTED_VOCAB_SIZE
    36→    if original_vocab_size == expected_vocab_size:
    37→        logging.info("Tokenizer already has the correct size.")
    38→        return tokenizer
    39→
    40→    new_tokens = [
    41→        constants.SPEECH_START_TOKEN,
    42→        constants.SPEECH_END_TOKEN,
    43→        constants.TEXT_PROMPT_START_TOKEN,
    44→        constants.TEXT_PROMPT_END_TOKEN,
    45→        constants.VOICE_DESCRIPTION_START_TOKEN,
    46→        constants.VOICE_DESCRIPTION_END_TOKEN,
    47→        constants.SOUND_EFFECT_START_TOKEN,
    48→        constants.SOUND_EFFECT_END_TOKEN,
    49→    ]
    50→    new_tokens.extend(
    51→        [constants.SPEECH_TOKEN_PATTERN.format(i) for i in range(codebook_size)]
    52→    )
    53→    num_added_tokens = tokenizer.add_tokens(sorted(new_tokens))
    54→    new_vocab_size = len(tokenizer)
    55→
    56→    if new_vocab_size < expected_vocab_size:
    57→        num_extra_tokens = expected_vocab_size - new_vocab_size
    58→        extra_tokens = [f"<|extra_token_{i}|>" for i in range(num_extra_tokens)]
    59→        num_added_tokens += tokenizer.add_tokens(extra_tokens)
    60→        new_vocab_size = len(tokenizer)
    61→        logging.info(
    62→            "Added another [%d] token(s) to the tokenizer so the vocab size is [%d].",
    63→            num_extra_tokens,
    64→            new_vocab_size,
    65→        )
    66→
    67→    if new_vocab_size != expected_vocab_size:
    68→        raise ValueError(
    69→            f"Expected tokenizer size to be {expected_vocab_size}, "
    70→            f"but got {new_vocab_size}!"
    71→        )
    72→    logging.info(
    73→        "Added [%d] speech tokens to the tokenizer total. Final vocab size: [%d].",
    74→        num_added_tokens,
    75→        new_vocab_size,
    76→    )
    77→
    78→    return tokenizer
    79→


     1→# Training-time constants.
     2→LOSS_IGNORE_TOKEN_ID = -100
     3→
     4→# Tokenization.
     5→SPEECH_TOKEN_PATTERN = "<|s_{}|>"
     6→SPEECH_START_TOKEN = "<|speech_start|>"
     7→SPEECH_END_TOKEN = "<|speech_end|>"
     8→TEXT_PROMPT_START_TOKEN = "<|text_prompt_start|>"
     9→TEXT_PROMPT_END_TOKEN = "<|text_prompt_end|>"
    10→VOICE_DESCRIPTION_START_TOKEN = "<|voice_description_start|>"
    11→VOICE_DESCRIPTION_END_TOKEN = "<|voice_description_end|>"
    12→SOUND_EFFECT_START_TOKEN = "<|sound_effect_start|>"
    13→SOUND_EFFECT_END_TOKEN = "<|sound_effect_end|>"
    14→END_HEADER_ID = "<|end_header_id|>"
    15→
    16→# File names.
    17→CONFIG_FILE_NAME = "training_config.json"
    18→
    19→# Audio.
    20→CODEC_TOKENS_RATE = 50
    21→CODEC_SAMPLE_RATE = 16000
    22→
    23→# Splits.
    24→TRAIN_SPLIT = "train"
    25→VAL_SPLIT = "val"
    26→
    27→# vLLM related constants.
    28→DEFAULT_MODEL_INSTRUCTION = "Convert the text to speech:"
    29→
    30→# Metrics.
    31→TOTAL_SOURCE = "total"
    32→
    33→# Reward functions.
    34→WER_REWARD_FUNC = "WERRewardFunc"
    35→DNSMOS_REWARD_FUNC = "DNSMOSRewardFunc"
    36→SIMILARITY_REWARD_FUNC = "SimilarityRewardFunc"
    37→
    38→# 21 nonverbal tokens.
    39→NONVERBAL_TOKENS = [
    40→    "<breathe>",
    41→    "<burp>",
    42→    "<chew>",
    43→    "<clear_throat>",
    44→    "<cough>",
    45→    "<cry>",
    46→    "<gasp>",
    47→    "<grunt>",
    48→    "<hiccup>",
    49→    "<laugh>",
    50→    "<moan>",
    51→    "<pant>",
    52→    "<scream>",
    53→    "<sigh>",
    54→    "<sing>",
    55→    "<slurp>",
    56→    "<sneeze>",
    57→    "<sniff>",
    58→    "<snort>",
    59→    "<whistle>",
    60→    "<yawn>",
    61→]
    62→


Now let me read training-related files:


     1→import contextlib
     2→import os
     3→import time
     4→from collections.abc import Iterator
     5→
     6→import lightning.fabric as lightning_fabric
     7→import torch
     8→import wandb
     9→from absl import logging
    10→
    11→from tts.core import constants
    12→from tts.inference import quality_validation
    13→from tts.training import checkpointing, evaluation, r2_upload
    14→from tts.utils import configuration, custom_logging
    15→
    16→_PERFORMANCE_METRICS_STEP_OFFSET = 5
    17→
    18→# Watchdog thresholds (configurable via env vars for production tuning).
    19→_WATCHDOG_PAD_RATIO_WARN = float(os.environ.get("MAYA_WATCHDOG_PAD_RATIO", "0.50"))
    20→_WATCHDOG_THROUGHPUT_DROP_FRAC = float(os.environ.get("MAYA_WATCHDOG_THROUGHPUT_DROP", "0.30"))
    21→_WATCHDOG_LOSS_EXPLOSION = float(os.environ.get("MAYA_WATCHDOG_LOSS_EXPLOSION", "100.0"))
    22→
    23→
    24→def _running_average(old: float, new: float) -> float:
    25→    """Computes a running average of the old and new values."""
    26→    if old == -1.0:
    27→        return new
    28→    else:
    29→        return 0.9 * old + 0.1 * new
    30→
    31→
    32→def _distributed_any(
    33→    fabric: lightning_fabric.Fabric,
    34→    value: bool | torch.Tensor,
    35→) -> bool:
    36→    """Returns True if any rank reports a truthy scalar flag."""
    37→    if isinstance(value, torch.Tensor):
    38→        tensor = value.detach().reshape(()).to(device=fabric.device, dtype=torch.int32)
    39→    else:
    40→        tensor = torch.tensor(int(bool(value)), device=fabric.device, dtype=torch.int32)
    41→
    42→    if (
    43→        fabric.world_size <= 1
    44→        or not torch.distributed.is_available()
    45→        or not torch.distributed.is_initialized()
    46→    ):
    47→        return bool(tensor.item())
    48→
    49→    return bool(fabric.all_reduce(tensor, reduce_op="max").item())
    50→
    51→
    52→def _get_first_rank_error(
    53→    fabric: lightning_fabric.Fabric, local_error: str | None
    54→) -> str | None:
    55→    """Collects and returns the first non-empty error string across ranks."""
    56→    if (
    57→        fabric.world_size <= 1
    58→        or not torch.distributed.is_available()
    59→        or not torch.distributed.is_initialized()
    60→    ):
    61→        return local_error
    62→
    63→    gathered: list[str | None] = [None] * fabric.world_size
    64→    torch.distributed.all_gather_object(gathered, local_error)
    65→    for error in gathered:
    66→        if error:
    67→            return error
    68→    return None
    69→
    70→
    71→def _get_batch_sampler(train_data_loader):
    72→    """Extracts the batch sampler from a fabric-wrapped dataloader, if present."""
    73→    dl = getattr(train_data_loader, "_dataloader", train_data_loader)
    74→    bs = getattr(dl, "batch_sampler", None)
    75→    return bs
    76→
    77→
    78→def _get_sampler_state(train_data_loader) -> dict | None:
    79→    """Returns batch sampler state_dict if the sampler supports it."""
    80→    bs = _get_batch_sampler(train_data_loader)
    81→    if bs is not None and hasattr(bs, "state_dict"):
    82→        return bs.state_dict()
    83→    return None
    84→
    85→
    86→def _compute_sampler_resume_state(
    87→    step: int,
    88→    gradient_accumulation_steps: int,
    89→    batches_per_rank: int,
    90→) -> dict:
    91→    """Computes the exact sampler resume state from the training step count.
    92→
    93→    This is prefetch-safe: with ``num_workers > 0`` the sampler may have
    94→    *yielded* more batches than the training loop actually consumed (due to
    95→    DataLoader prefetching).  Rather than reading the sampler's internal
    96→    ``_batches_yielded`` counter (which can be ahead of reality), we
    97→    recompute the correct epoch and position from the number of micro-batches
    98→    the training loop has definitely consumed:
    99→
   100→        total_consumed = step * gradient_accumulation_steps
   101→        epoch          = total_consumed // batches_per_rank
   102→        position       = total_consumed %  batches_per_rank
   103→    """
   104→    total_consumed = step * gradient_accumulation_steps
   105→    epoch = total_consumed // batches_per_rank
   106→    position = total_consumed % batches_per_rank
   107→    return {"epoch": epoch, "batches_yielded": position}
   108→
   109→
   110→def _validate_resume_schedule_compatibility(
   111→    config: configuration.ExperimentConfig,
   112→    saved_config: dict | None,
   113→    checkpoint_file_to_resume_from: str,
   114→) -> None:
   115→    """Fails fast when a full-checkpoint resume would alter the LR schedule."""
   116→    if not saved_config:
   117→        return
   118→
   119→    saved_training = saved_config.get("training") or {}
   120→    current_training = config.training
   121→    scheduler_fields = (
   122→        "max_training_steps",
   123→        "warmup_ratio",
   124→        "learning_rate",
   125→        "min_lr_ratio",
   126→    )
   127→    mismatches = []
   128→    for field_name in scheduler_fields:
   129→        saved_value = saved_training.get(field_name)
   130→        current_value = getattr(current_training, field_name)
   131→        if saved_value != current_value:
   132→            mismatches.append(
   133→                f"{field_name}: checkpoint={saved_value}, current={current_value}"
   134→            )
   135→
   136→    if mismatches:
   137→        raise ValueError(
   138→            "Checkpoint resume would change the learning-rate schedule. "
   139→            f"Checkpoint: [{checkpoint_file_to_resume_from}]. "
   140→            "Mismatched fields: "
   141→            + "; ".join(mismatches)
   142→            + ". Resume with the original scheduler settings, or use "
   143→            "`only_load_model_weights=true` for a fresh schedule."
   144→        )
   145→
   146→    saved_dataset = saved_config.get("dataset") or {}
   147→    saved_total_training_steps = saved_dataset.get("total_training_steps")
   148→    current_total_training_steps = config.dataset.total_training_steps
   149→    if (
   150→        saved_total_training_steps is not None


     1→"""Main entry point for the Inworld Text-to-Speech training pipeline."""
     2→
     3→import dataclasses
     4→import os
     5→import time
     6→from collections.abc import Sequence
     7→
     8→import lightning.fabric as lightning_fabric
     9→import torch
    10→import wandb
    11→from absl import app, flags, logging
    12→
    13→from tts.core import constants, lora, modeling, optimization, tokenization
    14→from tts.data import text_normalization, tts_datasets
    15→from tts.inference import quality_validation
    16→from tts.training import data_loading
    17→from tts.training import checkpointing, environment, r2_upload, training_loop
    18→from tts.utils import configuration, custom_logging
    19→
    20→FLAGS = flags.FLAGS
    21→
    22→
    23→def _define_flag_or_get_existing(definer, name: str, *args, **kwargs):
    24→    if name in FLAGS:
    25→        return FLAGS[name]
    26→    return definer(name, *args, **kwargs)
    27→
    28→
    29→_CONFIG_PATH = _define_flag_or_get_existing(
    30→    flags.DEFINE_string,
    31→    "config_path", None, "Path to the training config json file."
    32→)
    33→_SLURM_DISTRIBUTED = _define_flag_or_get_existing(
    34→    flags.DEFINE_bool,
    35→    "slurm_distributed", False, "Whether the script is running via SLURM."
    36→)
    37→_RUN_NAME = _define_flag_or_get_existing(
    38→    flags.DEFINE_string, "run_name", None, "Run name for wandb logging."
    39→)
    40→_DRY_RUN = _define_flag_or_get_existing(
    41→    flags.DEFINE_bool, "dry_run", False, "Whether to run in dry run mode."
    42→)
    43→_EXPERIMENT_DIR = _define_flag_or_get_existing(
    44→    flags.DEFINE_string,
    45→    "experiment_dir", "./experiments", "Experiment directory."
    46→)
    47→_PRETRAINING_MODE = _define_flag_or_get_existing(
    48→    flags.DEFINE_bool,
    49→    "pretraining_mode", False, "Whether to run in pretraining mode."
    50→)
    51→_COMPILE_MODEL = _define_flag_or_get_existing(
    52→    flags.DEFINE_bool, "compile_model", False, "Whether to compile model."
    53→)
    54→_USE_LIGER_KERNEL = _define_flag_or_get_existing(
    55→    flags.DEFINE_bool,
    56→    "use_liger_kernel", False, "Whether to use Liger fused kernels (RoPE, SwiGLU, RMSNorm, CE)."
    57→)
    58→_SEED = _define_flag_or_get_existing(
    59→    flags.DEFINE_integer, "seed", 777, "Random seed."
    60→)
    61→_OVERFIT_SAMPLES = _define_flag_or_get_existing(
    62→    flags.DEFINE_integer,
    63→    "overfit_samples",
    64→    0,
    65→    "If >0, run a tiny overfit test on this many samples instead of full training.",
    66→)
    67→
    68→# Wandb logging flags.
    69→_PROJECT_NAME = _define_flag_or_get_existing(
    70→    flags.DEFINE_string,
    71→    "project_name", None, "Project name for wandb logging."
    72→)
    73→_USE_WANDB = _define_flag_or_get_existing(
    74→    flags.DEFINE_bool, "use_wandb", False, "Whether to use wandb for logging."
    75→)
    76→
    77→
    78→def _broadcast_validation_error(
    79→    fabric: lightning_fabric.Fabric, validation_error: str | None
    80→) -> str | None:
    81→    """Broadcasts a validation error so every rank can fail consistently."""
    82→    return _broadcast_object_from_global_zero(fabric, validation_error)
    83→
    84→
    85→def _broadcast_object_from_global_zero(
    86→    fabric: lightning_fabric.Fabric, value: object | None
    87→) -> object | None:
    88→    """Broadcasts a small Python object from rank 0 to all ranks."""
    89→    if (
    90→        fabric.world_size <= 1
    91→        or not torch.distributed.is_available()
    92→        or not torch.distributed.is_initialized()
    93→    ):
    94→        return value
    95→
    96→    payload = [value if fabric.is_global_zero else None]
    97→    torch.distributed.broadcast_object_list(payload, src=0)
    98→    return payload[0]
    99→
   100→
   101→def _setup_directories(
   102→    experiment_dir: str, checkpoint_dir: str, is_global_zero: bool
   103→) -> tuple[str, str]:
   104→    """Makes path absolute and creates directory if it doesn't exist."""
   105→    experiment_dir = os.path.abspath(experiment_dir)
   106→    checkpoint_dir = os.path.join(experiment_dir, checkpoint_dir)
   107→    if not os.path.exists(checkpoint_dir) and is_global_zero:
   108→        os.makedirs(checkpoint_dir)
   109→        logging.info("Checkpoint directory was created: %s.", checkpoint_dir)
   110→
   111→    return experiment_dir, checkpoint_dir
   112→
   113→
   114→def run_training(
   115→    fabric: lightning_fabric.Fabric,
   116→    config: configuration.ExperimentConfig,
   117→    use_wandb: bool,
   118→    pretraining_mode: bool,
   119→    dry_run: bool,
   120→) -> None:
   121→    """Launches the training pipeline."""
   122→    modeling_params = config.modeling.parameters
   123→    max_seq_len = modeling_params["max_seq_len"]
   124→    is_lora = config.lora is not None
   125→    eval_enabled = config.training.eval_steps > 0
   126→
   127→    # ------------------ Tokenizer ------------------ #
   128→    tokenizer = tokenization.build_tokenizer(
   129→        model_name=modeling_params["model_name"],
   130→        max_seq_len=max_seq_len,
   131→        codebook_size=modeling_params["codebook_size"],
   132→        is_lora=is_lora,
   133→    )
   134→
   135→    if fabric.is_global_zero:
   136→        with custom_logging.Timer() as t:
   137→            tokenizer.save_pretrained(config.checkpointing.directory)
   138→        logging.info(
   139→            "Tokenizer saved to %s in %.2f seconds.",
   140→            config.checkpointing.directory,
   141→            t.get_duration(),
   142→        )
   143→
   144→    # ------------------ Model ------------------ #
   145→    vocab_size = len(tokenizer)
   146→    with custom_logging.Timer() as t:
   147→        deepspeed = config.training.strategy == configuration.TrainingStrategy.DEEPSPEED
   148→        model = modeling.build_model(
   149→            fabric,
   150→            modeling_params["model_name"],


     1→import dataclasses
     2→import os
     3→import random
     4→
     5→import lightning.fabric as lightning_fabric
     6→import numpy as np
     7→import torch
     8→import wandb
     9→from absl import logging
    10→
    11→from tts.core import constants
    12→from tts.training import r2_upload
    13→from tts.utils import configuration, custom_logging
    14→
    15→_BEST_CHECKPOINT_NAME = "best_model.pt"
    16→
    17→
    18→@dataclasses.dataclass
    19→class LoadedCheckpointState:
    20→    """Checkpoint metadata needed to resume training safely."""
    21→
    22→    statistics: custom_logging.Statistics | None
    23→    rng_state: dict | None
    24→    rng_state_per_rank: list[dict] | None
    25→    sampler_state: dict | None
    26→    saved_config: dict | None
    27→    best_eval_loss: float | None
    28→
    29→
    30→def _capture_rng_state() -> dict:
    31→    """Captures the full RNG state so exact resume is possible without fast-forward."""
    32→    state = {
    33→        "torch_rng": torch.random.get_rng_state(),
    34→        "python_rng": random.getstate(),
    35→        "numpy_rng": np.random.get_state(),
    36→    }
    37→    if torch.cuda.is_available():
    38→        state["cuda_rng"] = torch.cuda.get_rng_state_all()
    39→    return state
    40→
    41→
    42→def _restore_rng_state(rng_state: dict) -> None:
    43→    """Restores RNG state captured by _capture_rng_state."""
    44→    torch.random.set_rng_state(rng_state["torch_rng"])
    45→    random.setstate(rng_state["python_rng"])
    46→    np.random.set_state(rng_state["numpy_rng"])
    47→    if "cuda_rng" in rng_state and torch.cuda.is_available():
    48→        torch.cuda.set_rng_state_all(rng_state["cuda_rng"])
    49→
    50→
    51→def _gather_rng_state_per_rank(
    52→    fabric: lightning_fabric.Fabric, local_rng_state: dict
    53→) -> list[dict] | None:
    54→    """Collects rank-local RNG states so distributed resume can be exact."""
    55→    if (
    56→        fabric.world_size <= 1
    57→        or not torch.distributed.is_available()
    58→        or not torch.distributed.is_initialized()
    59→    ):
    60→        return None
    61→
    62→    gathered: list[dict | None] = [None] * fabric.world_size
    63→    torch.distributed.all_gather_object(gathered, local_rng_state)
    64→    return [state for state in gathered if state is not None]
    65→
    66→
    67→def restore_checkpoint_rng_state(
    68→    *,
    69→    global_rank: int,
    70→    rng_state: dict | None,
    71→    rng_state_per_rank: list[dict] | None,
    72→) -> None:
    73→    """Restores the correct RNG state payload for the current rank."""
    74→    if rng_state_per_rank:
    75→        if global_rank >= len(rng_state_per_rank):
    76→            raise ValueError(
    77→                "Checkpoint RNG state does not contain an entry for "
    78→                f"global rank {global_rank}."
    79→            )
    80→        _restore_rng_state(rng_state_per_rank[global_rank])
    81→        return
    82→
    83→    if rng_state is not None:
    84→        _restore_rng_state(rng_state)
    85→
    86→
    87→def load_from_checkpoint(
    88→    fabric: lightning_fabric.Fabric,
    89→    model: torch.nn.Module,
    90→    optimizer: torch.optim.Optimizer,
    91→    checkpoint_file_to_resume_from: str,
    92→    load_full_checkpoint: bool = True,
    93→) -> tuple[torch.nn.Module, torch.optim.Optimizer, LoadedCheckpointState]:
    94→    """Loads checkpoint and updates model state.
    95→
    96→    Returns the model, optimizer, and resume metadata.
    97→    """
    98→    checkpoint = {"model": model}
    99→    if load_full_checkpoint:
   100→        checkpoint.update({
   101→            "optimizer": optimizer,
   102→            "loss_statistics": {},
   103→            "rng_state": {},
   104→            "rng_state_per_rank": [],
   105→            "sampler_state": {},
   106→            "config": {},
   107→            "best_eval_loss": None,
   108→        })
   109→
   110→    fabric.load(checkpoint_file_to_resume_from, checkpoint, strict=True)
   111→    loaded_state = LoadedCheckpointState(
   112→        statistics=None,
   113→        rng_state=None,
   114→        rng_state_per_rank=None,
   115→        sampler_state=None,
   116→        saved_config=None,
   117→        best_eval_loss=None,
   118→    )
   119→    if load_full_checkpoint:
   120→        loaded_state.statistics = custom_logging.Statistics.from_dict(
   121→            checkpoint["loss_statistics"]
   122→        )
   123→        rng_state = checkpoint.get("rng_state")
   124→        if rng_state and isinstance(rng_state, dict) and "torch_rng" in rng_state:
   125→            logging.info("RNG state found in checkpoint; will restore for exact resume.")
   126→            loaded_state.rng_state = rng_state
   127→
   128→        rng_state_per_rank = checkpoint.get("rng_state_per_rank")
   129→        if rng_state_per_rank:
   130→            loaded_state.rng_state_per_rank = rng_state_per_rank
   131→            logging.info(
   132→                "Per-rank RNG state found in checkpoint for %d rank(s).",
   133→                len(rng_state_per_rank),
   134→            )
   135→
   136→        sampler_state = checkpoint.get("sampler_state")
   137→        if sampler_state and isinstance(sampler_state, dict) and "epoch" in sampler_state:
   138→            logging.info(
   139→                "Sampler state found in checkpoint (epoch=%d, batches_yielded=%d).",
   140→                sampler_state["epoch"],
   141→                sampler_state.get("batches_yielded", 0),
   142→            )
   143→            loaded_state.sampler_state = sampler_state
   144→
   145→        saved_config = checkpoint.get("config")
   146→        if saved_config:
   147→            loaded_state.saved_config = saved_config
   148→
   149→        best_eval_loss = checkpoint.get("best_eval_loss")
   150→        if isinstance(best_eval_loss, (int, float)):
   151→            loaded_state.best_eval_loss = float(best_eval_loss)
   152→
   153→    return model, optimizer, loaded_state
   154→
   155→
   156→def save_to_checkpoint(
   157→    fabric: lightning_fabric.Fabric,
   158→    model: torch.nn.Module,
   159→    config: configuration.ExperimentConfig,
   160→    optimizer: torch.optim.Optimizer,
   161→    statistics: custom_logging.Statistics,
   162→    checkpoint_name: str | None = None,
   163→    sampler_state: dict | None = None,
   164→    best_eval_loss: float | None = None,
   165→) -> str:
   166→    """Saves the model, optimizer, RNG state, and training stats to a checkpoint."""
   167→    checkpoint_name = checkpoint_name or f"checkpoint_{statistics.step}.pt"
   168→    checkpoint_file = os.path.join(config.checkpointing.directory, checkpoint_name)
   169→    local_rng_state = _capture_rng_state()
   170→
   171→    checkpoint = {
   172→        "model": model,
   173→        "loss_statistics": statistics.as_dict(),
   174→        "optimizer": optimizer,
   175→        "config": config.to_dict(),
   176→        "rng_state": local_rng_state,
   177→    }
   178→    rng_state_per_rank = _gather_rng_state_per_rank(fabric, local_rng_state)
   179→    if rng_state_per_rank is not None:
   180→        checkpoint["rng_state_per_rank"] = rng_state_per_rank
   181→    if sampler_state is not None:
   182→        checkpoint["sampler_state"] = sampler_state
   183→    if best_eval_loss is not None:
   184→        checkpoint["best_eval_loss"] = best_eval_loss
   185→    fabric.save(path=checkpoint_file, state=checkpoint)
   186→
   187→    if fabric.is_global_zero:
   188→        keep_only_last_n_checkpoints = config.checkpointing.keep_only_last_n_checkpoints
   189→        if keep_only_last_n_checkpoints is not None:
   190→            _cleanup_old_checkpoints(
   191→                checkpoint_dir=config.checkpointing.directory,
   192→                keep_local=keep_only_last_n_checkpoints,
   193→            )
   194→
   195→    return checkpoint_file
   196→
   197→
   198→def _cleanup_old_checkpoints(
   199→    checkpoint_dir: str,
   200→    keep_local: int,
   201→) -> None:
   202→    """Deletes old checkpoints that are confirmed uploaded to R2.
   203→
   204→    Safety rules:
   205→      1. Always keep the newest *keep_local* checkpoints on disk regardless.
   206→      2. Among older checkpoints, only delete those confirmed uploaded to R2.
   207→      3. Never delete best_model.pt.
   208→    """
   209→    checkpoint_files = [
   210→        f
   211→        for f in os.listdir(checkpoint_dir)
   212→        if f.startswith("checkpoint_") and f.endswith(".pt")
   213→        and f != _BEST_CHECKPOINT_NAME
   214→    ]
   215→    if len(checkpoint_files) <= keep_local:
   216→        return
   217→
   218→    checkpoint_files.sort(key=lambda x: int(x.split("_")[-1].split(".")[0]))
   219→    # The newest keep_local checkpoints are always protected.
   220→    candidates_for_deletion = checkpoint_files[:-keep_local]
   221→
   222→    uploaded = r2_upload.get_uploaded_checkpoints()
   223→    for fname in candidates_for_deletion:
   224→        full_path = os.path.join(checkpoint_dir, fname)
   225→        if full_path in uploaded:
   226→            logging.info(
   227→                "Removing checkpoint %s (confirmed uploaded to R2).", fname
   228→            )
   229→            os.remove(full_path)
   230→        else:
   231→            logging.info(
   232→                "Keeping checkpoint %s (not yet confirmed in R2).", fname
   233→            )
   234→
   235→
   236→def maybe_save_best_checkpoint(
   237→    fabric: lightning_fabric.Fabric,
   238→    model: torch.nn.Module,
   239→    config: configuration.ExperimentConfig,
   240→    optimizer: torch.optim.Optimizer,
   241→    statistics: custom_logging.Statistics,
   242→    eval_loss: float,
   243→    best_eval_loss: float,
   244→    sampler_state: dict | None = None,
   245→) -> float:
   246→    """Saves a 'best_model.pt' checkpoint when eval_loss improves."""
   247→    if eval_loss < best_eval_loss:
   248→        logging.info(
   249→            "New best eval loss %.5f < %.5f, saving best checkpoint.",
   250→            eval_loss,
   251→            best_eval_loss,
   252→        )
   253→        save_to_checkpoint(
   254→            fabric, model, config, optimizer, statistics,
   255→            checkpoint_name=_BEST_CHECKPOINT_NAME,
   256→            sampler_state=sampler_state,
   257→            best_eval_loss=eval_loss,
   258→        )
   259→        return eval_loss
   260→    return best_eval_loss
   261→
   262→
   263→def save_config(
   264→    experiment_config: configuration.ExperimentConfig,
   265→    checkpoint_dir: str,
   266→    use_wandb: bool,
   267→):
   268→    """Saves model config to a file."""
   269→    config_file = os.path.join(checkpoint_dir, constants.CONFIG_FILE_NAME)
   270→    with open(config_file, "w") as f:
   271→        f.write(str(experiment_config))
   272→
   273→    # Config might be with new values after first initialized, to ensure
   274→    # consistency, the config here should be updated with wandb.
   275→    #
   276→    # TODO: sweep run set value can be overriden by python's training code
   277→    #       leading to config here and one shown in the W&B UI being different.
   278→    if use_wandb:
   279→        wandb.config.update(experiment_config.to_dict(), allow_val_change=True)
   280→


Now let me look at codec files and data loading:


     1→"""Codec encoder implementation."""
     2→
     3→import collections
     4→
     5→import torch
     6→import transformers
     7→import vector_quantize_pytorch as vq
     8→from absl import logging
     9→
    10→from tts.core import constants
    11→from tts.core.codec import encoder_modules
    12→
    13→_HOP_LENGTH = 320
    14→_HALF_HOP_LENGTH = _HOP_LENGTH // 2
    15→
    16→
    17→class Encoder(torch.nn.Module):
    18→    """The audio encoder model."""
    19→
    20→    def __init__(self, model_path: str):
    21→        super().__init__()
    22→
    23→        # model properties
    24→        self.sample_rate = constants.CODEC_SAMPLE_RATE
    25→        self.token_rate = constants.CODEC_TOKENS_RATE
    26→
    27→        # encoders
    28→        self.semantic_encoder = encoder_modules.SemanticEncoder(
    29→            input_channels=1024,
    30→            output_channels=1024,
    31→            encode_channels=1024,
    32→            kernel_size=3,
    33→        )
    34→        self.acoustic_encoder = encoder_modules.AcousticEncoder(
    35→            num_generator_features=48,
    36→            initial_conv_kernel_size=7,
    37→            final_conv_kernel_size=3,
    38→            up_ratios=[2, 2, 4, 4, 5],
    39→            dilations=(1, 3, 9),
    40→            output_dim=1024,
    41→        )
    42→        self.fusion_layer = torch.nn.Linear(2048, 2048)
    43→
    44→        # vector quantization.
    45→        self.quantizer = vq.ResidualFSQ(
    46→            dim=2048, levels=[4, 4, 4, 4, 4, 4, 4, 4], num_quantizers=1
    47→        )
    48→        self.load_from_checkpoint(model_path)
    49→
    50→        # wav2vec model for semantic features extraction.
    51→        self.wav2vec_feature_extractor = (
    52→            transformers.AutoFeatureExtractor.from_pretrained("facebook/w2v-bert-2.0")
    53→        )
    54→        self.wav2vec_model = transformers.Wav2Vec2BertModel.from_pretrained(
    55→            "facebook/w2v-bert-2.0", output_hidden_states=True
    56→        )
    57→
    58→    def forward(self, wavs: torch.Tensor, feats: torch.Tensor) -> torch.Tensor:
    59→        """Computes VQ codes for a batch of audio files."""
    60→        acoustic_hidden_states = self.acoustic_encoder(wavs)
    61→        acoustic_hidden_states = acoustic_hidden_states.transpose(1, 2)
    62→
    63→        wav2vec_features = self.wav2vec_model(feats[:, 0, :, :]).hidden_states[16]
    64→        semantic_hidden_states = self.semantic_encoder(wav2vec_features.transpose(1, 2))
    65→
    66→        hidden_states = torch.cat(
    67→            [semantic_hidden_states, acoustic_hidden_states], dim=1
    68→        )
    69→        hidden_states = self.fusion_layer(hidden_states.transpose(1, 2)).transpose(1, 2)
    70→
    71→        return self.quantize(hidden_states)
    72→
    73→    def quantize(self, hidden_states: torch.Tensor) -> torch.Tensor:
    74→        """Quantize the hidden states to the VQ codes."""
    75→        hidden_states = hidden_states.permute(0, 2, 1)
    76→        _, vq_code = self.quantizer(hidden_states)
    77→        vq_code = vq_code.permute(0, 2, 1)
    78→        return vq_code
    79→
    80→    def load_from_checkpoint(self, checkpoint_path: str):
    81→        """Loads the encoder from a checkpoint file."""
    82→        logging.info("Loading encoder checkpoint from %s", checkpoint_path)
    83→        ckpt = torch.load(checkpoint_path, map_location="cpu", weights_only=False)
    84→
    85→        if "state_dict" in ckpt.keys():
    86→            # to provide compatibility for the checkpoint in
    87→            # https://huggingface.co/HKUSTAudio/xcodec2/tree/main/ckpt
    88→            ckpt = ckpt["state_dict"]
    89→
    90→            filtered_state_dict_codec = collections.OrderedDict()
    91→            filtered_state_dict_semantic_encoder = collections.OrderedDict()
    92→            filtered_state_dict_gen = collections.OrderedDict()
    93→            filtered_state_dict_fc_prior = collections.OrderedDict()
    94→            for key, value in ckpt.items():
    95→                if key.startswith("CodecEnc."):
    96→                    new_key = key[len("CodecEnc.") :]
    97→                    filtered_state_dict_codec[new_key] = value
    98→                elif key.startswith("generator.quantizer."):
    99→                    new_key = key[len("generator.quantizer.") :]
   100→                    filtered_state_dict_gen[new_key] = value


     1→"""Decoder model for audio waveform generation."""
     2→
     3→import collections
     4→import itertools
     5→import math
     6→from typing import Any
     7→
     8→import torch
     9→from absl import logging
    10→
    11→from tts.core.codec import criterion, decoder_modules, discriminator, upsampler
    12→
    13→
    14→class Decoder(torch.nn.Module):
    15→    """The decoder model for audio waveform generation."""
    16→
    17→    def __init__(
    18→        self,
    19→        sample_rate: int,
    20→        hop_length: int,
    21→        upsample_factors: list[int],
    22→        kernel_sizes: list[int],
    23→        checkpoint_path: str | None = None,
    24→    ):
    25→        super().__init__()
    26→        self.sample_rate = sample_rate
    27→        self.hop_length = hop_length
    28→        self.upsample_factors = upsample_factors
    29→        self.kernel_sizes = kernel_sizes
    30→
    31→        total_ups = math.prod(self.upsample_factors) if self.upsample_factors else 1
    32→        if self.sample_rate // self.hop_length // total_ups != 50:
    33→            raise ValueError(
    34→                f"Current hop length {self.hop_length} and upsample "
    35→                f"factors {self.upsample_factors} do not match the target "
    36→                f"sample rate {self.sample_rate}."
    37→            )
    38→
    39→        logging.info(
    40→            "Creating audio decoder with output sample_rate"
    41→            f" {self.sample_rate}, upsample factors {self.upsample_factors}, "
    42→            f"kernel sizes {self.kernel_sizes}, hop_length {self.hop_length}."
    43→        )
    44→
    45→        # decoder includes the quantizer and the vocoder.
    46→        self.decoder = decoder_modules.Generator(hop_length=hop_length)
    47→
    48→        if self.upsample_factors:
    49→            self.upsampler = upsampler.UpSamplerBlock(
    50→                in_channels=1024,
    51→                upsample_factors=self.upsample_factors,
    52→                kernel_sizes=self.kernel_sizes,
    53→            )
    54→            logging.info(
    55→                "Upsampler created with upsample factors %s and kernel sizes %s",
    56→                self.upsample_factors,
    57→                self.kernel_sizes,
    58→            )
    59→        else:
    60→            logging.info("No upsampler created.")
    61→            self.upsampler = None
    62→
    63→        self.fc_post_a = torch.nn.Linear(2048, 1024)
    64→
    65→        if checkpoint_path is not None:
    66→            logging.info("Loading codec checkpoint from %s", checkpoint_path)
    67→            self.load_from_checkpoint(checkpoint_path)
    68→
    69→    def forward(self, vq_codes: torch.Tensor) -> torch.Tensor:
    70→        """Computes VQ codes for a batch of audio files."""
    71→
    72→        # vq_codes: (batch_size, codes_length) or (batch_size, 1, code_length)
    73→        if len(vq_codes.shape) == 2:
    74→            vq_codes = vq_codes.unsqueeze(1)
    75→
    76→        vq_codes = vq_codes.transpose(1, 2).to(self.fc_post_a.weight.device)
    77→        vq_post_emb = self.decoder.quantizer.get_output_from_indices(vq_codes)
    78→
    79→        vq_post_emb_fc = self.fc_post_a(vq_post_emb)
    80→
    81→        # vocos_backbone -> upsampler -> istft_head
    82→        decoder_hidden = self.decoder.backbone(vq_post_emb_fc)
    83→        if self.upsampler:
    84→            upsampled_hidden = self.upsampler(decoder_hidden.transpose(1, 2))
    85→        else:
    86→            upsampled_hidden = decoder_hidden
    87→
    88→        recon_audio = self.decoder.head(upsampled_hidden)
    89→        return recon_audio
    90→
    91→    def load_from_checkpoint(self, checkpoint_path: str) -> None:
    92→        ckpt = torch.load(checkpoint_path, map_location="cpu", weights_only=False)
    93→
    94→        if "state_dict" in ckpt.keys():
    95→            # to provide compatibility for the checkpoint in
    96→            # https://huggingface.co/HKUSTAudio/xcodec2/tree/main/ckpt
    97→            ckpt = ckpt["state_dict"]
    98→
    99→            filtered_state_dict_generator = collections.OrderedDict()
   100→            filtered_state_dict_fc_post = collections.OrderedDict()


     1→import torch
     2→
     3→
     4→def can_inject_distributed_sampler(
     5→    data_loader: torch.utils.data.DataLoader,
     6→) -> bool:
     7→    """Returns whether Lightning can safely inject a distributed sampler."""
     8→    batch_sampler = getattr(data_loader, "batch_sampler", None)
     9→    return batch_sampler is None or isinstance(
    10→        batch_sampler, torch.utils.data.BatchSampler
    11→    )
    12→


Let me check the data structures and more training details:


     1→import dataclasses
     2→import uuid
     3→from typing import Any
     4→
     5→_DEFAULT_VALUES = {
     6→    "speaker_id": "",
     7→    "emotion": "",
     8→    "language": "unknown",
     9→    "dnsmos_mos_ovr": 0.0,
    10→    "style": "",
    11→}
    12→_REQUIRED_FIELDS = frozenset(("wav_path",))
    13→
    14→
    15→@dataclasses.dataclass
    16→class Sample:
    17→    """A TTS sample. Can and should be adjusted for the specific use-cases.
    18→
    19→    Args:
    20→        id: The id of the sample.
    21→        wav_path: The path to the wav file.
    22→        speaker_id: The speaker id.
    23→
    24→        language: The language of the sample.
    25→        emotion: The emotion of the sample.
    26→        transcript: The transcript of the sample.
    27→        voice_description: The voice description of the sample.
    28→        sound_effect: The sound effect of the sample.
    29→
    30→        duration: The duration of the sample in seconds.
    31→        sample_rate: Original sample rate of the wav file.
    32→        dataset_name: The name of the dataset the sample is from.
    33→        dnsmos_mos_ovr: The dnsmos mos ovr of the sample.
    34→        style: The style of the sample (e.g., whispering, shouting, laughing, etc.).
    35→        original_data: The original sample data used to support flexible formatting.
    36→    """
    37→
    38→    # Metadata.
    39→    id: str
    40→    wav_path: str
    41→    speaker_id: str
    42→
    43→    # Annotations.
    44→    language: str
    45→    emotion: str
    46→    transcript: str
    47→    voice_description: str
    48→    sound_effect: str
    49→
    50→    # Properties.
    51→    duration: float
    52→    sample_rate: int
    53→    dataset_name: str
    54→    dnsmos_mos_ovr: float
    55→    style: str
    56→    original_data: dict[str, Any] = dataclasses.field(default_factory=dict)
    57→
    58→    def __post_init__(self):
    59→        """Validates sample fields."""
    60→        if not self.transcript and not self.voice_description and not self.sound_effect:
    61→            raise ValueError(
    62→                "At least one of transcript, voice_description, "
    63→                "or sound_effect must be set."
    64→            )
    65→
    66→    def to_json(self) -> dict[str, str]:
    67→        return {k: v for k, v in dataclasses.asdict(self).items() if v}
    68→
    69→    @classmethod
    70→    def from_json(cls, json_data: dict[str, Any], dataset_name: str) -> "Sample":
    71→        if not dataset_name:
    72→            raise ValueError("dataset_name is required")
    73→        for field in _REQUIRED_FIELDS:
    74→            if json_data.get(field) is None:
    75→                raise ValueError(f"{field} is required for sample: {json_data}")
    76→
    77→        return Sample(
    78→            id=json_data.get("id", str(uuid.uuid4())),
    79→            wav_path=json_data["wav_path"],
    80→            speaker_id=json_data.get("speaker_id", _DEFAULT_VALUES["speaker_id"]),
    81→            emotion=json_data.get("emotion", _DEFAULT_VALUES["emotion"]).lower(),
    82→            transcript=json_data.get("transcript", ""),
    83→            voice_description=json_data.get("voice_description", ""),
    84→            sound_effect=json_data.get("sound_effect", ""),
    85→            language=json_data.get("language", _DEFAULT_VALUES["language"]),
    86→            duration=json_data.get("duration", -1.0),
    87→            sample_rate=json_data.get("sample_rate", -1),
    88→            dataset_name=dataset_name,
    89→            # Accept tx_quality_score as an alias for dnsmos_mos_ovr so that
    90→            # SFT samples exported with tx_quality_score still get filtered
    91→            # correctly when min_dnsmos_score > 0.
    92→            dnsmos_mos_ovr=json_data.get(
    93→                "dnsmos_mos_ovr",
    94→                json_data.get("tx_quality_score", _DEFAULT_VALUES["dnsmos_mos_ovr"]),
    95→            ),
    96→            style=json_data.get("style", _DEFAULT_VALUES["style"]).lower(),
    97→            original_data=json_data.get("original_data", {}),
    98→        )
    99→


     1→import dataclasses
     2→import hashlib
     3→import math
     4→import os
     5→from bisect import bisect_right
     6→from collections.abc import Callable
     7→from typing import Any
     8→
     9→import datasets
    10→import numpy as np
    11→import torch
    12→import transformers
    13→from absl import logging
    14→
    15→from tts.core import constants, prompting
    16→from tts.data import data_utils, text_normalization
    17→from tts.data.datasets import finetuning, pretraining, rlhf
    18→from tts.utils import configuration
    19→
    20→
    21→def _get_sft_dataset(dataset_path: str, split: str) -> datasets.Dataset:
    22→    """Handles creation of a text fine-tuning dataset from a dataset path."""
    23→    dataset = datasets.load_from_disk(dataset_path.replace("[text]", ""))
    24→    subset = split if split == constants.TRAIN_SPLIT else "test"
    25→    return dataset.train_test_split(test_size=0.001)[subset]
    26→
    27→
    28→def _build_dataset(
    29→    *,
    30→    tokenizer: transformers.PreTrainedTokenizer,
    31→    dataset_path: str,
    32→    max_seq_len: int,
    33→    split: str,
    34→    pretraining_mode: bool,
    35→    text_normalizer: text_normalization.TextNormalizer,
    36→    dataset_config: configuration.DatasetConfig,
    37→    codebook_size: int | None = None,
    38→) -> tuple[torch.utils.data.Dataset, str]:
    39→    """Builds a dataset from a dataset path."""
    40→    dataset_name = os.path.basename(dataset_path)
    41→    text_dataset = dataset_name.endswith("[text]")
    42→    paired_dataset = dataset_name.endswith("[paired]")
    43→    if pretraining_mode:
    44→        if text_dataset:
    45→            dataset_dir = dataset_path.replace("[text]", "")
    46→            return pretraining.TextPretrainingDataset(
    47→                dataset_dir=dataset_dir,
    48→                split=split,
    49→                max_seq_len=max_seq_len,
    50→                tokenizer_name=tokenizer.name_or_path,
    51→            ), dataset_name
    52→        if paired_dataset:
    53→            dataset_dir = dataset_path.replace("[paired]", "")
    54→            if codebook_size is None:
    55→                raise ValueError("`codebook_size` is required for paired PT datasets.")
    56→            return pretraining.PairedPretrainingDataset(
    57→                dataset_dir=dataset_dir,
    58→                split=split,
    59→                max_seq_len=max_seq_len,
    60→                tokenizer=tokenizer,
    61→                text_normalizer=text_normalizer,
    62→                dataset_config=dataset_config,
    63→                codebook_size=codebook_size,
    64→            ), dataset_name
    65→        if codebook_size is None:
    66→            raise ValueError("`codebook_size` is required for audio pretraining datasets.")
    67→        return pretraining.TtsPretrainingDataset(
    68→            dataset_dir=dataset_path,
    69→            split=split,
    70→            max_seq_len=max_seq_len,
    71→            tokenizer=tokenizer,
    72→            codebook_size=codebook_size,
    73→            dataset_config=dataset_config,
    74→        ), dataset_name
    75→    if paired_dataset:
    76→        raise ValueError("`[paired]` datasets are only supported in pretraining mode.")
    77→    if text_dataset:
    78→        dataset = _get_sft_dataset(dataset_path, split)
    79→        return finetuning.TextFineTuningDataset(
    80→            dataset, tokenizer, max_seq_len
    81→        ), dataset_name
    82→
    83→    codes, samples, indexes, _ = data_utils.load_and_filter_audio_codes_and_samples(
    84→        dataset_dir=dataset_path, split=split, dataset_config=dataset_config
    85→    )
    86→
    87→    if dataset_config.enable_rlhf_training:
    88→        return rlhf.TtsRLHFDataset(
    89→            dataset_name=dataset_name,
    90→            samples=samples,
    91→            codes=codes,
    92→            indexes=indexes,
    93→            tokenizer=tokenizer,
    94→            prompt_compiler=prompting.InferencePromptCompiler(),
    95→            text_normalizer=text_normalizer,
    96→        ), dataset_name
    97→    return finetuning.TtsFineTuningDataset(
    98→        dataset_name=dataset_name,
    99→        samples=samples,
   100→        codes=codes,
   101→        indexes=indexes,
   102→        tokenizer=tokenizer,
   103→        max_seq_len=max_seq_len,
   104→        prompt_compiler=prompting.TrainingPromptCompiler(),
   105→        text_normalizer=text_normalizer,
   106→    ), dataset_name
   107→
   108→
   109→@dataclasses.dataclass
   110→class WeightedDataset:
   111→    """A torch dataset with a weight.
   112→
   113→    Attrs:
   114→        name: Name of the dataset.
   115→        dataset: The dataset instance.
   116→        epochs: The number of epochs to run on the dataset.
   117→    """
   118→
   119→    name: str
   120→    dataset: torch.utils.data.Dataset
   121→    epochs: float
   122→
   123→
   124→_UINT64_MASK = (1 << 64) - 1
   125→
   126→
   127→def _stable_hash_to_uint64(value: str) -> int:
   128→    """Returns a stable 64-bit hash for deterministic sampling."""
   129→    digest = hashlib.sha256(value.encode("utf-8")).digest()
   130→    return int.from_bytes(digest[:8], byteorder="little", signed=False)
   131→
   132→
   133→def _mix_uint64(value: int) -> int:
   134→    """SplitMix64 mixer for deterministic pseudo-randomization."""
   135→    value = (value + 0x9E3779B97F4A7C15) & _UINT64_MASK
   136→    value = (value ^ (value >> 30)) * 0xBF58476D1CE4E5B9 & _UINT64_MASK
   137→    value = (value ^ (value >> 27)) * 0x94D049BB133111EB & _UINT64_MASK
   138→    return value ^ (value >> 31)
   139→
   140→
   141→def _coprime_stride(modulus: int, seed: int) -> int:
   142→    """Returns a permutation stride coprime with the modulus."""
   143→    if modulus <= 1:
   144→        return 0
   145→
   146→    candidate = int(seed % (modulus - 1)) + 1
   147→    while math.gcd(candidate, modulus) != 1:
   148→        candidate += 1
   149→        if candidate >= modulus:
   150→            candidate = 1


Let me check configuration and inference:


     1→import dataclasses
     2→import enum
     3→import json
     4→import os
     5→from typing import Any
     6→
     7→import cattr
     8→import wandb
     9→
    10→from tts.training import environment
    11→
    12→_REQUIRED_CONFIG_KEYS = [
    13→    "train_weighted_datasets",
    14→    "val_weighted_datasets",
    15→    "training",
    16→    "modeling",
    17→    "checkpointing",
    18→]
    19→
    20→
    21→class TrainingStrategy(enum.Enum):
    22→    """Training strategy."""
    23→
    24→    # Distributed Data Parallel (DDP). Used by default.
    25→    DDP = "ddp"
    26→
    27→    # Fully Sharded Data Parallel (FSDP).
    28→    # When using FSDP with model compilation (e.g., torch.compile),
    29→    # evaluation must run before training.
    30→    # To support this, make sure checkpoints are saved just before evaluation steps,
    31→    # so that after loading a checkpoint, evaluation runs immediately.
    32→    FSDP = "fsdp"
    33→
    34→    # DeepSpeed Stage 2 that uses ZeRO-2 (aka gradient partitioning).
    35→    DEEPSPEED = "deepspeed"
    36→
    37→
    38→@dataclasses.dataclass(frozen=False)
    39→class DatasetConfig:
    40→    """Defines variable dataset parameters computed upon training start."""
    41→
    42→    # List of allowed languages for dataset filtering.
    43→    # If empty or None, no filtering will be done.
    44→    allowed_languages: list[str]
    45→    # Minimum DNSMOS score for dataset filtering.
    46→    min_dnsmos_score: float
    47→    # Minimum sample rate for audio in the dataset filtering.
    48→    min_sample_rate: int
    49→    # Whether to enable RLHF (Reinforcement Learning from Human Feedback) training.
    50→    enable_rlhf_training: bool
    51→    # Minimum audio duration for dataset filtering.
    52→    min_audio_duration: float = 0.0
    53→    # List of allowed instruction finetuning annotations.
    54→    # If empty or None, no instruction finetuning will be done.
    55→    allowed_ift_annotations: list[str] | None = None
    56→    # The total number of training steps needed to traverse the dataset.
    57→    # To be set when loading the dataset.
    58→    total_training_steps: int | None = None
    59→    # The number of steps needed to traverse the dataset once.
    60→    # To be set when loading the dataset.
    61→    steps_per_epoch: int | None = None
    62→
    63→
    64→@dataclasses.dataclass(frozen=True)
    65→class ModelingConfig:
    66→    # Model architecture parameters.
    67→    parameters: dict[str, Any]
    68→
    69→
    70→@dataclasses.dataclass(frozen=False)
    71→class CheckpointingConfig:
    72→    # The number of steps to save the model for.
    73→    save_steps: int
    74→
    75→    # Checkpoint dir.
    76→    directory: str
    77→
    78→    # Collect model health statistics every |save_steps| steps.
    79→    collect_health_stats: bool
    80→
    81→    # Whether to save intermediate generations of voice given random
    82→    # validation dataset samples.
    83→    save_intermediate_generations: bool
    84→
    85→    # Type of validation to perform during checkpointing ("random_phrases" or
    86→    # "prompt_continuation")
    87→    validation_type: str = "random_phrases"
    88→
    89→    # Codec encoder checkpoint path for quality validation decoding.
    90→    codec_encoder_path: str | None = None
    91→
    92→    # Codec decoder checkpoint path for quality validation decoding.
    93→    codec_decoder_path: str | None = None
    94→
    95→    # Config-driven canary prompt WAVs for quality validation.
    96→    # Dict mapping wav_path -> transcript.  When None, built-in defaults are used.
    97→    canary_prompts: dict[str, str] | None = None
    98→
    99→    # Whether to keep only the last N checkpoints.
   100→    keep_only_last_n_checkpoints: int | None = None
   101→
   102→    # If set, the model will resume training from the given checkpoint.
   103→    checkpoint_file_to_resume_from: str | None = None
   104→
   105→    # If set, the model will only load the model weights from the given checkpoint.
   106→    only_load_model_weights: bool = False
   107→
   108→
   109→@dataclasses.dataclass(frozen=True)
   110→class TrainingConfig:
   111→    """Configuration for training."""
   112→
   113→    # The global seed for all ranks.
   114→    seed: int
   115→
   116→    # The number of steps to perform logging for.
   117→    logging_steps: int
   118→
   119→    # The number of steps to do periodic evaluation for.
   120→    eval_steps: int
   121→
   122→    # The number of steps to accumulate gradients for.
   123→    gradient_accumulation_steps: int
   124→
   125→    # The value to clip gradients to.
   126→    gradient_clip_value: float
   127→
   128→    # The initial learning rate.
   129→    learning_rate: float
   130→
   131→    # AdamW betas.
   132→    betas: tuple[float, float]
   133→
   134→    # The number of warmup steps as a fraction of the total number of steps.
   135→    warmup_ratio: float
   136→
   137→    # The batch size.
   138→    batch_size: int
   139→
   140→    # The weight decay.s
   141→    weight_decay: float
   142→
   143→    # The precision to use for training.
   144→    precision: str
   145→
   146→    # The training strategy to use.
   147→    strategy: TrainingStrategy
   148→
   149→    # If true, the model will use gradient checkpointing.
   150→    gradient_checkpointing: bool
   151→
   152→    # The number of workers to use for data loading.
   153→    num_workers: int = 0
   154→
   155→    # If set, use length-bucketed batching with this token budget per batch
   156→    # instead of fixed batch_size. Reduces padding waste for variable-length SFT.
   157→    max_tokens_per_batch: int | None = None
   158→
   159→    # Explicit step budget override. When set, takes precedence over the
   160→    # epoch-derived step count (useful for staged PT / SFT recipes).
   161→    max_training_steps: int | None = None
   162→
   163→    # Minimum LR as a fraction of peak LR for the cosine schedule floor.
   164→    min_lr_ratio: float = 0.1
   165→
   166→
   167→@dataclasses.dataclass(frozen=True)
   168→class RLHFConfig:
   169→    # The directory of the base model.
   170→    base_model_dir: str
   171→
   172→    # Float that controls the cumulative probability of the top tokens to consider.
   173→    top_p: float
   174→
   175→    # Integer that controls the number of top tokens to consider.
   176→    top_k: int
   177→
   178→    # Float that controls the repetition penalty.
   179→    repetition_penalty: float
   180→
   181→    # Float that controls the temperature of the logits.
   182→    temperature: float
   183→
   184→    # Number of generations per prompt to sample.
   185→    num_generations: int
   186→
   187→    # Maximum length of the prompt. If the prompt is longer than this value, it
   188→    # will be truncated left.
   189→    max_prompt_length: int
   190→
   191→    # Maximum length of the generated completion.
   192→    max_completion_length: int
   193→
   194→    # Minimum length of the generated completion.
   195→    min_completion_length: int
   196→
   197→    # Whether to use vllm for training.
   198→    use_vllm: bool
   199→
   200→    # The reward functions to use for training.


     1→"""Implements inference logic for local TTS models."""
     2→
     3→import dataclasses
     4→from typing import Any
     5→
     6→import torch
     7→import transformers
     8→from absl import logging
     9→
    10→from tts.core import constants, prompting
    11→from tts.core.codec import decoding, encoding
    12→from tts.utils import custom_logging
    13→
    14→
    15→class InferenceSettings:
    16→    """Settings for a TTS inference."""
    17→
    18→    def __init__(
    19→        self,
    20→        temperature=0.8,
    21→        max_tokens=1792,
    22→        min_tokens=10,
    23→        top_p=1.0,
    24→        top_k=50,
    25→        repetition_penalty=1.1,
    26→        frequency_penalty=0.3,
    27→        seed=42,
    28→    ):
    29→        """Initialize inference settings with default or provided values."""
    30→        self.temperature = temperature
    31→        self.max_tokens = max_tokens
    32→        self.min_tokens = min_tokens
    33→        self.top_p = top_p
    34→        self.top_k = top_k
    35→        self.repetition_penalty = repetition_penalty
    36→        self.frequency_penalty = frequency_penalty
    37→        self.seed = seed
    38→
    39→
    40→DEFAULT_INFERENCE_SETTINGS = InferenceSettings()
    41→
    42→
    43→@dataclasses.dataclass(frozen=True)
    44→class InferenceResult:
    45→    """Result of a TTS inference."""
    46→
    47→    wav: torch.Tensor
    48→    encoding_time: float
    49→    decoding_time: float
    50→    inference_time: float
    51→
    52→
    53→def extract_speech_ids(speech_tokens_str: list[str]) -> list[int]:
    54→    """Extracts speech ids from speech tokens strings."""
    55→    speech_ids = []
    56→    for token_str in speech_tokens_str:
    57→        if token_str.startswith("<|s_") and token_str.endswith("|>"):
    58→            num_str = token_str[4:-2]
    59→            num = int(num_str)
    60→            speech_ids.append(num)
    61→        else:
    62→            logging.error("Unexpected token: %s", token_str)
    63→    return speech_ids
    64→
    65→
    66→@torch.no_grad()
    67→def _generate_speech_tokens(
    68→    model: transformers.AutoModelForCausalLM | Any,
    69→    input_ids: torch.Tensor,
    70→    inference_settings: InferenceSettings,
    71→    speech_end_id: int,
    72→    use_vllm: bool = False,
    73→) -> torch.Tensor:
    74→    """Implements actual speech token generation logic."""
    75→    if use_vllm:
    76→        import vllm
    77→
    78→        sampling_params = vllm.SamplingParams(
    79→            max_tokens=inference_settings.max_tokens,
    80→            min_tokens=inference_settings.min_tokens,
    81→            stop_token_ids=[speech_end_id],
    82→            repetition_penalty=inference_settings.repetition_penalty,
    83→            top_p=inference_settings.top_p,
    84→            top_k=inference_settings.top_k,
    85→            frequency_penalty=inference_settings.frequency_penalty,
    86→            temperature=inference_settings.temperature,
    87→            detokenize=False,
    88→        )
    89→        outputs = model.generate(
    90→            prompt_token_ids=input_ids[0].tolist(), sampling_params=sampling_params
    91→        )
    92→        return outputs[0].outputs[0].token_ids
    93→
    94→    return (
    95→        model.generate(
    96→            input_ids=input_ids,
    97→            attention_mask=torch.ones_like(input_ids),
    98→            # `InferenceSettings.max_tokens` is intended as "new tokens to sample"
    99→            # (same semantics as the vLLM path above). Using `max_length` here
   100→            # causes severe truncation for long prompts because it includes the
   101→            # prompt length. Use `max_new_tokens` to keep behavior consistent.
   102→            max_new_tokens=inference_settings.max_tokens,
   103→            min_new_tokens=min(
   104→                inference_settings.min_tokens, inference_settings.max_tokens
   105→            ),
   106→            eos_token_id=speech_end_id,
   107→            do_sample=True if inference_settings.temperature > 0.0 else False,
   108→            repetition_penalty=inference_settings.repetition_penalty,
   109→            top_p=inference_settings.top_p,
   110→            top_k=inference_settings.top_k,
   111→            temperature=inference_settings.temperature,
   112→            pad_token_id=speech_end_id,
   113→        )
   114→        .cpu()
   115→        .squeeze(0)
   116→    )
   117→
   118→
   119→@torch.no_grad()
   120→def _synthesize_audio(
   121→    model: transformers.AutoModelForCausalLM,
   122→    tokenizer: transformers.AutoTokenizer,
   123→    audio_decoder: decoding.AudioDecoderInterface,
   124→    speech_ids: list[int],
   125→    prompt: str,
   126→    model_device: torch.device,
   127→    inference_settings: InferenceSettings,
   128→    use_vllm: bool,
   129→) -> tuple[torch.Tensor, float]:
   130→    """Synthesizes audio from text using a finetuned FinchTTS model."""
   131→    input_ids = tokenizer(prompt, add_special_tokens=True, return_tensors="pt")[
   132→        "input_ids"
   133→    ]
   134→    input_ids = input_ids.to(model_device)
   135→    speech_end_id = tokenizer.convert_tokens_to_ids(constants.SPEECH_END_TOKEN)
   136→
   137→    # Generate the speech autoregressively the classic way.
   138→    transformers.set_seed(inference_settings.seed)
   139→    generated_ids = _generate_speech_tokens(
   140→        model=model,
   141→        input_ids=input_ids,
   142→        inference_settings=inference_settings,
   143→        speech_end_id=speech_end_id,
   144→        use_vllm=use_vllm,
   145→    )
   146→
   147→    # Convert string speech tokens to speech token ids.
   148→    if use_vllm:
   149→        speech_tokens = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)
   150→        # append the prompt speech ids to the generated speech ids


Let me check example config files to understand training parameters:


     1→{
     2→    "training": {
     3→        "seed": 777,
     4→        "logging_steps": 150,
     5→        "eval_steps": 300,
     6→        "gradient_accumulation_steps": 1,
     7→        "gradient_clip_value": 1.0,
     8→        "learning_rate": 1e-04,
     9→        "betas": [
    10→            0.9,
    11→            0.95
    12→        ],
    13→        "warmup_ratio": 0.05,
    14→        "batch_size": 4,
    15→        "weight_decay": 0.1,
    16→        "precision": "bf16",
    17→        "strategy": "ddp",
    18→        "gradient_checkpointing": false,
    19→        "num_workers": 1
    20→    },
    21→    "modeling": {
    22→        "parameters": {
    23→            "codebook_size": 65536,
    24→            "max_seq_len": 2048,
    25→            "model_name": "meta-llama/Llama-3.2-1B-Instruct",
    26→            "enable_text_normalization": true
    27→        }
    28→    },
    29→    "checkpointing": {
    30→        "save_steps": 100,
    31→        "collect_health_stats": false,
    32→        "save_intermediate_generations": false,
    33→        "only_load_model_weights": true,
    34→        "keep_only_last_n_checkpoints": 10
    35→    },
    36→    "train_weighted_datasets": {
    37→        "/path/to/your/vectorized_dataset": 1.0
    38→    },
    39→    "val_weighted_datasets": {
    40→        "/path/to/your/vectorized_dataset": 1.0
    41→    },
    42→    "dataset": {
    43→        "randomize_slider": 0,
    44→        "allowed_languages": [],
    45→        "min_dnsmos_score": 0,
    46→        "min_sample_rate": 0,
    47→        "enable_rlhf_training": false
    48→    }
    49→}
    50→