training/interleaving.py

"""
Interleaved sequence creation for IST-LM training.

Single Responsibility: Only handles interleaved sequence generation.
"""

import random
import torch
import torch.nn.functional as F
from typing import List, Dict, Tuple, Optional, Any
from .config import SNAC_BASE_OFFSET, SNAC_LAYERS_PER_FRAME, SNAC_LAYER_OFFSET, EOS_TOKEN


def apply_snac_offset(token_idx: int, position: int) -> int:
    """
    Apply position-based offset to SNAC token.

    SNAC uses 7 tokens per frame with position-based offsets:
    offset = SNAC_BASE_OFFSET + (position % 7) * 4096

    Args:
        token_idx: Raw SNAC token index
        position: Position in the sequence

    Returns:
        Offset-adjusted token index
    """
    if int(token_idx) >= SNAC_BASE_OFFSET:
        # Already has offset applied
        return int(token_idx)
    offset = SNAC_BASE_OFFSET + (position % SNAC_LAYERS_PER_FRAME) * SNAC_LAYER_OFFSET
    return int(token_idx) + offset


def get_text_ratio(
    global_step: int,
    decay_steps: int = 300,
    initial_ratio: float = 0.9,
    min_ratio: float = 0.0
) -> float:
    """
    Calculate text ratio based on training step (IST-LM schedule).

    Schedule: Start at 90% text, decrease by 10% every decay_steps.
    - Step 0-299: 0.9
    - Step 300-599: 0.8
    - Step 600-899: 0.7
    - ...
    - Step 2700+: 0.0 (pure audio)

    Args:
        global_step: Current training step
        decay_steps: Steps between each 10% decay
        initial_ratio: Starting text ratio
        min_ratio: Minimum text ratio

    Returns:
        Current text ratio
    """
    num_decays = global_step // decay_steps
    text_ratio = initial_ratio - (num_decays * 0.1)
    return max(min_ratio, text_ratio)


def calculate_dynamic_decay_steps(
    total_steps: int,
    steps_per_epoch: int = None,
    initial_ratio: float = 0.9,
    final_audio_portion: float = 0.2
) -> int:
    """
    Calculate decay_steps for scheduled interleaving.

    Two modes:
    1. If steps_per_epoch provided: Complete decay in first epoch only,
       remaining epochs are pure audio.
    2. Otherwise: Use final_audio_portion to spread decay across training.

    Args:
        total_steps: Total training steps
        steps_per_epoch: Steps per epoch (if provided, decay completes in epoch 1)
        initial_ratio: Starting text ratio (default 0.9)
        final_audio_portion: Portion of training with p=0 (only used if steps_per_epoch=None)

    Returns:
        Calculated decay_steps
    """
    # Number of decay stages: 0.9 -> 0.0 = 9 steps (not 10)
    num_decay_stages = int(initial_ratio / 0.1)

    if steps_per_epoch is not None:
        # Complete decay in first epoch - remaining epochs are pure audio
        return max(1, steps_per_epoch // num_decay_stages)
    else:
        # Original behavior: spread across training
        steps_until_pure_audio = int(total_steps * (1 - final_audio_portion))
        return max(1, steps_until_pure_audio // num_decay_stages)


class InterleavingStrategy:
    """
    Base class for interleaving strategies.

    Open/Closed: Can create new strategies without modifying this class.
    """

    def create_sequence(
        self,
        text_tokens: List[int],
        snac_tokens: List[int],
        text_ratio: float,
        **kwargs
    ) -> Tuple[List[int], List[bool]]:
        """
        Create interleaved sequence.

        Args:
            text_tokens: Text token IDs
            snac_tokens: SNAC audio token IDs
            text_ratio: Ratio of text vs audio (0.0 = pure audio)

        Returns:
            Tuple of (interleaved_tokens, is_audio_mask)
        """
        raise NotImplementedError


class PositionalInterleaving(InterleavingStrategy):
    """
    Positional interleaving strategy (fallback when no alignments).

    Interleaves text and audio tokens based on fixed patterns
    determined by text_ratio.
    """

    # Pattern lookup: text_ratio -> (text_per_chunk, frames_per_chunk)
    PATTERNS = {
        0.9: (1, 3),   # 1 text token + 3 audio frames
        0.7: (1, 5),   # 1 text token + 5 audio frames
        0.5: (1, 7),   # 1 text token + 7 audio frames
        0.3: (1, 10),  # 1 text token + 10 audio frames
        0.0: (0, 1),   # Pure audio
    }

    def create_sequence(
        self,
        text_tokens: List[int],
        snac_tokens: List[int],
        text_ratio: float,
        **kwargs
    ) -> Tuple[List[int], List[bool]]:
        interleaved = []
        is_audio_mask = []

        if len(snac_tokens) == 0:
            return text_tokens + [EOS_TOKEN], [False] * (len(text_tokens) + 1)

        # Group SNAC into frames of 7
        frames = self._group_into_frames(snac_tokens)
        if len(frames) == 0:
            return text_tokens + [EOS_TOKEN], [False] * (len(text_tokens) + 1)

        # Get interleaving pattern
        text_per_chunk, frames_per_chunk = self._get_pattern(text_ratio)

        total_text = len(text_tokens)
        total_frames = len(frames)

        text_idx = 0
        frame_idx = 0
        snac_position = 0

        while frame_idx < total_frames:
            # Add text tokens
            if text_per_chunk > 0 and text_idx < total_text:
                for _ in range(text_per_chunk):
                    if text_idx < total_text:
                        interleaved.append(text_tokens[text_idx])
                        is_audio_mask.append(False)
                        text_idx += 1

            # Add audio frames
            for _ in range(frames_per_chunk):
                if frame_idx < total_frames:
                    frame = frames[frame_idx]
                    for tok in frame:
                        interleaved.append(apply_snac_offset(tok, snac_position))
                        is_audio_mask.append(True)
                        snac_position += 1
                    frame_idx += 1

        # Add remaining text (only if not pure audio mode)
        if text_per_chunk > 0:
            while text_idx < total_text:
                interleaved.append(text_tokens[text_idx])
                is_audio_mask.append(False)
                text_idx += 1

        # Add EOS
        interleaved.append(EOS_TOKEN)
        is_audio_mask.append(False)

        return interleaved, is_audio_mask

    def _group_into_frames(self, snac_tokens: List[int]) -> List[List[int]]:
        """Group SNAC tokens into frames of 7."""
        frames = []
        for i in range(0, len(snac_tokens), SNAC_LAYERS_PER_FRAME):
            frame = snac_tokens[i:i + SNAC_LAYERS_PER_FRAME]
            if len(frame) == SNAC_LAYERS_PER_FRAME:
                frames.append(frame)
        return frames

    def _get_pattern(self, text_ratio: float) -> Tuple[int, int]:
        """Get interleaving pattern for given text_ratio."""
        if text_ratio >= 0.9:
            return self.PATTERNS[0.9]
        elif text_ratio >= 0.7:
            return self.PATTERNS[0.7]
        elif text_ratio >= 0.5:
            return self.PATTERNS[0.5]
        elif text_ratio >= 0.3:
            return self.PATTERNS[0.3]
        else:
            return self.PATTERNS[0.0]


class AlignedInterleaving(InterleavingStrategy):
    """
    Word-aligned interleaving strategy.

    Uses word alignments to semantically replace audio spans
    with corresponding text tokens.
    """

    def create_sequence(
        self,
        text_tokens: List[int],
        snac_tokens: List[int],
        text_ratio: float,
        word_alignments: Optional[List[Dict]] = None,
        tokenizer=None,
        answer_text: str = "",
        **kwargs
    ) -> Tuple[List[int], List[bool]]:
        # Fall back to positional if no alignments
        if not word_alignments or text_ratio <= 0:
            return PositionalInterleaving().create_sequence(
                text_tokens, snac_tokens, text_ratio
            )

        # Check for pre-computed tokens
        has_precomputed = (
            len(word_alignments) > 0 and
            'tokens' in word_alignments[0] and
            word_alignments[0]['tokens']
        )

        if not has_precomputed and not tokenizer:
            return PositionalInterleaving().create_sequence(
                text_tokens, snac_tokens, text_ratio
            )

        interleaved = []
        is_audio_mask = []

        # Group SNAC into frames
        frames = []
        for i in range(0, len(snac_tokens), SNAC_LAYERS_PER_FRAME):
            frame = snac_tokens[i:i + SNAC_LAYERS_PER_FRAME]
            if len(frame) == SNAC_LAYERS_PER_FRAME:
                frames.append(frame)

        if len(frames) == 0:
            return text_tokens + [EOS_TOKEN], [False] * (len(text_tokens) + 1)

        total_frames = len(frames)

        # Randomly select words to replace with text
        num_words = len(word_alignments)
        num_text_words = int(num_words * text_ratio)
        word_indices = list(range(num_words))
        random.shuffle(word_indices)
        text_word_indices = set(word_indices[:num_text_words])

        # Frame rate conversion
        max_alignment_frame = max(wa['end_frame'] for wa in word_alignments)
        frame_ratio = total_frames / max_alignment_frame if max_alignment_frame > total_frames * 2 else 1.0

        snac_position = 0

        for word_idx, alignment in enumerate(word_alignments):
            word = alignment['word']
            start_frame = int(alignment['start_frame'] * frame_ratio)
            end_frame = min(int(alignment['end_frame'] * frame_ratio), total_frames)

            if word_idx in text_word_indices:
                # Replace audio with text
                word_tokens = alignment.get('tokens', [])
                if not word_tokens and tokenizer:
                    word_tokens = tokenizer.encode(word, add_special_tokens=False)

                for tok in word_tokens:
                    interleaved.append(tok)
                    is_audio_mask.append(False)
                snac_position = end_frame * SNAC_LAYERS_PER_FRAME
            else:
                # Keep audio
                for f_idx in range(start_frame, end_frame):
                    if f_idx < total_frames:
                        frame = frames[f_idx]
                        for tok in frame:
                            interleaved.append(apply_snac_offset(tok, snac_position))
                            is_audio_mask.append(True)
                            snac_position += 1

        # Add remaining frames
        remaining_start = max(wa['end_frame'] for wa in word_alignments)
        remaining_start = min(int(remaining_start * frame_ratio), total_frames)
        for f_idx in range(remaining_start, total_frames):
            frame = frames[f_idx]
            for tok in frame:
                interleaved.append(apply_snac_offset(tok, snac_position))
                is_audio_mask.append(True)
                snac_position += 1

        # Add EOS
        interleaved.append(EOS_TOKEN)
        is_audio_mask.append(False)

        return interleaved, is_audio_mask


def create_interleaved_sequence(
    text_tokens: List[int],
    snac_tokens: List[int],
    text_ratio: float = 0.9,
    word_alignments: Optional[List[Dict]] = None,
    tokenizer=None,
    answer_text: str = ""
) -> Tuple[List[int], List[bool]]:
    """
    Create interleaved sequence (convenience function).

    Automatically selects the best strategy based on available data.
    """
    if word_alignments and text_ratio > 0:
        strategy = AlignedInterleaving()
    else:
        strategy = PositionalInterleaving()

    return strategy.create_sequence(
        text_tokens=text_tokens,
        snac_tokens=snac_tokens,
        text_ratio=text_ratio,
        word_alignments=word_alignments,
        tokenizer=tokenizer,
        answer_text=answer_text,
    )


def apply_interleaving(
    batch: Dict[str, Any],
    text_ratio: float,
    tokenizer=None,
    max_seq_len: int = 2048
) -> Dict[str, torch.Tensor]:
    """
    Apply interleaving to a batch of samples.

    Args:
        batch: Batch from DataLoader with 'whisper' and 'raw_data'
        text_ratio: Current text ratio
        tokenizer: Tokenizer for on-the-fly encoding
        max_seq_len: Maximum sequence length

    Returns:
        Batch with 'whisper', 'interleaved', 'is_audio_mask'
    """
    raw_data = batch["raw_data"]
    sequences = []
    max_seq = 0

    for item in raw_data:
        interleaved, is_audio = create_interleaved_sequence(
            item["text_tokens"],
            item["snac_tokens"],
            text_ratio,
            word_alignments=item.get("word_alignments"),
            tokenizer=tokenizer,
            answer_text=item.get("answer_text", "")
        )

        # Truncate if needed
        if len(interleaved) > max_seq_len:
            interleaved = interleaved[:max_seq_len]
            is_audio = is_audio[:max_seq_len]

        sequences.append((interleaved, is_audio))
        max_seq = max(max_seq, len(interleaved))

    # Pad and stack
    interleaved_batch = []
    is_audio_batch = []

    for interleaved, is_audio in sequences:
        seq_tensor = torch.tensor(interleaved, dtype=torch.long)
        mask_tensor = torch.tensor(is_audio, dtype=torch.bool)

        seq_pad = F.pad(seq_tensor, (0, max_seq - len(interleaved)), value=-100)
        mask_pad = F.pad(mask_tensor, (0, max_seq - len(is_audio)), value=False)

        interleaved_batch.append(seq_pad)
        is_audio_batch.append(mask_pad)

    return {
        "whisper": batch["whisper"],
        "interleaved": torch.stack(interleaved_batch),
        "is_audio_mask": torch.stack(is_audio_batch)
    }