omnivoice/training/builder.py

#!/usr/bin/env python3
# Copyright    2026  Xiaomi Corp.        (authors:  Han Zhu)
#
# See ../../LICENSE for clarification regarding multiple authors
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Builders for constructing training components.

Provides factory functions to assemble the model, tokenizer, and data loaders
from a ``TrainingConfig``. Called by ``omnivoice.cli.train`` to set up training.

Key functions:
- ``build_model_and_tokenizer()``: Loads the model and text tokenizer.
- ``build_dataloaders()``: Builds train/eval data loaders from a data config JSON.
  The batching strategy is chosen based on ``TrainingConfig.attn_implementation``:

  - ``"flex_attention"``: sequence packing via ``PackingIterableDataset`` +
    ``PackingDataCollator``. Batch shape is ``[1, C, batch_tokens]``.
  - other (e.g. ``"sdpa"``): length-grouped padding via
    ``StreamLengthGroupDataset`` + ``PaddingDataCollator``. Batch shape
    is ``[B, C, max_len]`` where B ≥ 1 and max_len ≤ batch_tokens.
"""

import logging
from functools import partial
from typing import Tuple

import torch
from torch.utils.data import DataLoader
from transformers import AutoConfig, AutoModel, AutoTokenizer
from transformers import logging as hf_logging
from transformers.trainer_utils import seed_worker

from omnivoice.data.batching import PackingIterableDataset, StreamLengthGroupDataset
from omnivoice.data.collator import PackingDataCollator, PaddingDataCollator
from omnivoice.data.dataset import WebDatasetReader, prepare_data_manifests_from_json
from omnivoice.data.processor import OmniVoiceSampleProcessor
from omnivoice.models.omnivoice import OmniVoice, OmniVoiceConfig, _resolve_model_path
from omnivoice.training.config import TrainingConfig

logger = logging.getLogger(__name__)


def build_model_and_tokenizer(
    config: TrainingConfig,
) -> Tuple[OmniVoice, AutoTokenizer]:
    """Load Tokenizer and Model, handle resizing and special tokens."""
    logger.info("Initializing Model & Tokenizer...")

    # 1. Tokenizer
    tokenizer_path = (
        config.init_from_checkpoint
        if config.init_from_checkpoint
        else config.llm_name_or_path
    )
    tokenizer_path = _resolve_model_path(tokenizer_path)
    tokenizer = AutoTokenizer.from_pretrained(tokenizer_path)
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token

    new_tokens = [
        "<|denoise|>",
        "<|lang_start|>",
        "<|lang_end|>",
        "<|instruct_start|>",
        "<|instruct_end|>",
        "<|text_start|>",
        "<|text_end|>",
    ]

    tokens_to_add = [t for t in new_tokens if t not in tokenizer.get_vocab()]
    if tokens_to_add:
        tokenizer.add_special_tokens({"additional_special_tokens": tokens_to_add})

    if config.init_from_checkpoint:
        logger.info(f"Loading weights from {config.init_from_checkpoint}")
        model = OmniVoice.from_pretrained(
            config.init_from_checkpoint,
            attn_implementation=config.attn_implementation,
            dtype=torch.float32,
            train=True,
        )
    else:
        resolved_llm = _resolve_model_path(config.llm_name_or_path)
        llm_config = AutoConfig.from_pretrained(resolved_llm)

        ov_config = OmniVoiceConfig(
            audio_vocab_size=config.audio_vocab_size,
            audio_mask_id=config.audio_mask_id,
            num_audio_codebook=config.num_audio_codebook,
            audio_codebook_weights=config.audio_codebook_weights,
            llm_config=llm_config,
        )

        original_level = hf_logging.get_verbosity()
        hf_logging.set_verbosity_error()  # suppress expected lm_head.weight warnings

        llm = AutoModel.from_pretrained(
            resolved_llm,
            attn_implementation=config.attn_implementation,
            dtype=torch.float32,
        )

        hf_logging.set_verbosity(original_level)
        model = OmniVoice(config=ov_config, llm=llm)

    # 3. Resize Embeddings
    if len(tokenizer) != model.config.llm_config.vocab_size:
        model.llm.resize_token_embeddings(len(tokenizer))
        model.config.llm_config.vocab_size = len(tokenizer)

    # 4. Config IDs
    model.config.pad_token_id = tokenizer.pad_token_id
    model.config.bos_token_id = tokenizer.bos_token_id
    model.config.eos_token_id = tokenizer.eos_token_id

    return model, tokenizer


def build_dataloaders(
    config: TrainingConfig, tokenizer: AutoTokenizer
) -> Tuple[DataLoader, DataLoader]:
    """Setup Data Pipeline: Manifests -> WDS -> Batching -> Loaders.

    Batching strategy depends on ``config.attn_implementation``:
    - ``"flex_attention"``: sequence packing (PackingIterableDataset +
      PackingDataCollator). All samples are concatenated into one long sequence.
    - other (e.g. ``"sdpa"``): length-grouped padding
      (LengthGroupedIterableDataset + PaddingDataCollator). Samples with
      similar token lengths are batched together and padded to the same length.
    """
    logger.info("Initializing Data Readers...")

    processor = OmniVoiceSampleProcessor(
        text_tokenizer=tokenizer,
        num_channels=config.num_audio_codebook,
        audio_mask_id=config.audio_mask_id,
        prompt_ratio_range=config.prompt_ratio_range,
        mask_ratio_range=config.mask_ratio_range,
        drop_cond_ratio=config.drop_cond_ratio,
        language_ratio=config.language_ratio,
        use_pinyin_ratio=config.use_pinyin_ratio,
        instruct_ratio=config.instruct_ratio,
        only_instruct_ratio=config.only_instruct_ratio,
    )

    train_manifests, dev_manifests = prepare_data_manifests_from_json(
        config.data_config
    )
    raw_train_ds = WebDatasetReader(manifests=train_manifests, evaluation=False)

    use_packing = config.attn_implementation == "flex_attention"

    if use_packing:
        train_dataset = PackingIterableDataset(
            raw_train_ds, processor, config.batch_tokens
        )
        collate_fn = PackingDataCollator(processor, config.batch_tokens)
    else:
        train_dataset = StreamLengthGroupDataset(
            raw_train_ds,
            batch_duration=config.batch_tokens,
            min_length=config.min_sample_tokens,
            max_length=config.max_sample_tokens,
            max_sample=config.max_batch_size,
            processor=processor,
            length_fn=lambda s: s["length"],
        )
        collate_fn = PaddingDataCollator(processor, config.batch_tokens)

    logger.info(
        "Using %s (attn_implementation=%s)",
        "sequence packing" if use_packing else "length-grouped padding",
        config.attn_implementation,
    )

    init_fn = partial(
        seed_worker,
        num_workers=config.num_workers,
        rank=(
            torch.distributed.get_rank()
            if torch.distributed.is_initialized()
            else 0
        ),
    )

    train_loader = DataLoader(
        train_dataset,
        batch_size=None,
        num_workers=config.num_workers,
        collate_fn=collate_fn,
        worker_init_fn=init_fn,
        pin_memory=True,
        prefetch_factor=4,
    )

    eval_loader = None
    if dev_manifests:
        raw_dev_ds = WebDatasetReader(
            manifests=dev_manifests, evaluation=True
        )
        if use_packing:
            dev_dataset = PackingIterableDataset(
                raw_dev_ds, processor, config.batch_tokens
            )
        else:
            dev_dataset = StreamLengthGroupDataset(
                raw_dev_ds,
                batch_duration=config.batch_tokens,
                min_length=config.min_sample_tokens,
                max_length=config.max_sample_tokens,
                max_sample=config.max_batch_size,
                processor=processor,
                length_fn=lambda s: s["length"],
            )
        eval_loader = DataLoader(
            dev_dataset,
            batch_size=None,  # Each item is already a collated batch
            num_workers=1,
            collate_fn=collate_fn,
            pin_memory=True,
            prefetch_factor=2,
        )

    return train_loader, eval_loader