data.py

import os, io, re, tempfile, random
import numpy as np
import torch
import torchaudio
import soundfile as sf
import librosa

from typing import Dict, Any
from audiomentations import Compose, TimeStretch, Gain, PitchShift, OneOf, AddGaussianNoise
from datasets import load_dataset, concatenate_datasets, Audio, Dataset
from torch.utils.data import DataLoader, Dataset
from torch.nn.utils.rnn import pad_sequence

from transformers import WhisperFeatureExtractor
from processor import VLFMProcessor
from configs import VLFMConfig, build_tokenizer
from constants import IGNORE_INDEX


def remove_brackets(text: str) -> str:
    return re.sub(r"<.*?>", "", text or "").strip()


AUG = Compose([
    TimeStretch(min_rate=0.95, max_rate=1.05, p=0.10, leave_length_unchanged=False),
    Gain(min_gain_db=-3, max_gain_db=3, p=0.10),
    PitchShift(min_semitones=-2, max_semitones=2, p=0.10),
    OneOf([AddGaussianNoise(min_amplitude=0.003, max_amplitude=0.010, p=0.5)], p=0.10),
])


class VLFMDataset(Dataset):
    """
    Simple S2T:
      inputs  = user_text + <AUDIO_SPAN> (expanded) + [TGT] + target_text
      labels  = mask BEFORE [TGT] (keep [TGT] and after); EOS appended in processor
      KD alt  = user_text + source_text + [TGT] + target_text  (text-only)
    """
    def __init__(
        self,
        config: VLFMConfig,
        dataset: Dataset,
        kd: bool = False,
        train_on_inputs: bool = False,
        max_response_tokens: int | None = None,
        use_augmentation: bool = True,
    ):
        self.dataset = dataset
        self.config = config
        self.kd = kd
        self.train_on_inputs = train_on_inputs
        self.max_response_tokens = max_response_tokens
        self.use_aug = use_augmentation

        self.feature_extractor = WhisperFeatureExtractor.from_pretrained(config.audio_model_id)
        self.tokenizer, self.audio_token_id = build_tokenizer(config.text_model_id, config.tokenizer_padding_side)
        self.processor = VLFMProcessor(self.feature_extractor, self.tokenizer, config)

        self.TGT_ID = self.tokenizer.convert_tokens_to_ids("[TGT]")


    @staticmethod
    def _ensure_mono_16k_float32(audio, sr, target_sr=16000):
        if audio.ndim == 2:
            audio = audio.mean(axis=0 if audio.shape[0] < audio.shape[1] else 1)
        elif audio.ndim != 1:
            raise ValueError(f"Expected 1-D mono waveform, got shape {audio.shape}")
        if audio.size == 0 or not np.isfinite(audio).all():
            raise ValueError("Decoded audio is empty or contains non-finite values")
        if sr != target_sr:
            audio = librosa.resample(audio, orig_sr=sr, target_sr=target_sr)
        return audio.astype(np.float32, copy=False)

    def process_audio(self, audio_dict: Dict[str, Any]):
        if "bytes" in audio_dict and audio_dict["bytes"] is not None:
            bio = io.BytesIO(audio_dict["bytes"])
            try:
                y, sr = sf.read(bio, dtype="float32", always_2d=False)
                return self._ensure_mono_16k_float32(y, int(sr))
            except Exception:
                pass
            try:
                bio.seek(0)
                with tempfile.NamedTemporaryFile(suffix=".bin", delete=False) as tmp:
                    tmp.write(bio.read()); tmp_path = tmp.name
                yt, sr = torchaudio.load(tmp_path)
                y = yt.mean(0).numpy() if yt.ndim == 2 else yt.numpy()
                return self._ensure_mono_16k_float32(y, int(sr))
            finally:
                try: os.remove(tmp_path)
                except Exception: pass

        if "path" in audio_dict and audio_dict["path"]:
            path = audio_dict["path"]
            try:
                y, sr = sf.read(path, dtype="float32", always_2d=False)
                return self._ensure_mono_16k_float32(y, int(sr))
            except Exception:
                yt, sr = torchaudio.load(path)
                y = yt.mean(0).numpy() if yt.ndim == 2 else yt.numpy()
                return self._ensure_mono_16k_float32(y, int(sr))

        raise ValueError("Audio object has neither usable 'bytes' nor 'path'")


    def _mask_before_tgt(self, input_ids_1d: torch.Tensor) -> torch.Tensor:
        labels = input_ids_1d.clone()
        if self.train_on_inputs:
            return labels
        pos = (input_ids_1d == self.TGT_ID).nonzero(as_tuple=True)
        if len(pos[0]) == 0:
            # If [TGT] somehow missing, be conservative: mask nothing
            return labels
        tgt_idx = int(pos[0][0].item())
        labels[:tgt_idx] = IGNORE_INDEX
        return labels

    def __len__(self): 
        return len(self.dataset)

    def __getitem__(self, idx):
        retries, last_err = 10, None
        for _ in range(retries):
            try:
                row = self.dataset[idx]

                # Required fields
                audio_dict   = row["audio"]
                source_text  = remove_brackets(row.get("source_text", ""))
                target_text  = remove_brackets(row.get("target_text", ""))
                user_text    = row.get("user_prompt", "Translate from English to Yoruba:")

                # audio
                audio = self.process_audio(audio_dict)
                if self.use_aug:
                    audio = AUG(samples=audio, sample_rate=16000)

                # Build STUDENT (audio) input
                full = self.processor.build_inputs(
                    user_text=user_text,
                    target_text=target_text,
                    audio_array=audio,
                    sampling_rate=16000,
                    add_eos=True,
                    max_length=None,   
                )

                # Labels (mask BEFORE [TGT], keep [TGT] + target + EOS)
                labels = self._mask_before_tgt(full["input_ids"].squeeze(0))

                # Optional response window cap
                if self.max_response_tokens is not None:
                    ids_1d = full["input_ids"].squeeze(0)
                    pos = (ids_1d == self.TGT_ID).nonzero(as_tuple=True)
                    if len(pos[0]) > 0:
                        start = int(pos[0][0].item())
                        max_len = min(ids_1d.size(0), start + self.max_response_tokens)
                        for k in ("input_ids", "attention_mask"):
                            full[k] = full[k][:, :max_len]
                        labels = labels[:max_len]

                # Build item
                item = {
                    "input_ids":      full["input_ids"].squeeze(0),
                    "attention_mask": full["attention_mask"].squeeze(0),
                    "labels":         labels,
                    "input_features": full["input_features"].squeeze(0),  # [80, T]
                    "audio_token_start_idx": full["audio_token_start_idx"].squeeze(0),
                    "audio_token_len":       full["audio_token_len"].squeeze(0),
                    "audio_lens":            full["audio_lens"].squeeze(0),
                    "audio_batch_size":      full["audio_batch_size"].squeeze(0),
                }

                # KD / Teacher path (text-only)
                if self.kd:
                    alt = self.processor.build_text_only(
                        user_text=user_text,
                        source_text=source_text,
                        target_text=target_text,
                        add_eos=True,
                        max_length=None,
                    )
                    alt_ids = alt["input_ids"]
                    alt_labels = self._mask_before_tgt(alt_ids)
                    item.update({
                        "alt_input_ids": alt_ids,
                        "alt_attention_mask": alt["attention_mask"],
                        "alt_labels": alt_labels,
                    })

                return item

            except Exception as e:
                last_err = e
                idx = np.random.randint(0, len(self.dataset))
        raise last_err


def vlfm_collate_fn(batch, pad_token_id):
    input_ids = pad_sequence([b["input_ids"] for b in batch], batch_first=True, padding_value=pad_token_id)
    attention_mask = pad_sequence([b["attention_mask"] for b in batch], batch_first=True, padding_value=0)
    labels = pad_sequence([b["labels"] for b in batch], batch_first=True, padding_value=IGNORE_INDEX)

    feats = [b["input_features"] for b in batch]  # each [80, T]
    T_max = max(f.shape[-1] for f in feats)
    feats_pad = [torch.nn.functional.pad(f, (0, T_max - f.shape[-1])) for f in feats]
    input_features = torch.stack(feats_pad, dim=0)

    audio_token_start_idx = torch.stack([b["audio_token_start_idx"] for b in batch])
    audio_token_len       = torch.stack([b["audio_token_len"]       for b in batch])
    audio_lens            = torch.stack([b["audio_lens"]            for b in batch])
    audio_batch_size      = torch.stack([b["audio_batch_size"]      for b in batch]) if "audio_batch_size" in batch[0] else torch.ones(len(batch), dtype=torch.long)

    out = {
        "input_ids": input_ids,
        "attention_mask": attention_mask,
        "labels": labels,
        "input_features": input_features,
        "audio_token_start_idx": audio_token_start_idx,
        "audio_token_len": audio_token_len,
        "audio_lens": audio_lens,
        "audio_batch_size": audio_batch_size,
    }

    if "alt_input_ids" in batch[0]:
        alt_input_ids = pad_sequence([b["alt_input_ids"] for b in batch], batch_first=True, padding_value=pad_token_id)
        alt_attention_mask = pad_sequence([b["alt_attention_mask"] for b in batch], batch_first=True, padding_value=0)
        alt_labels = pad_sequence([b["alt_labels"] for b in batch], batch_first=True, padding_value=IGNORE_INDEX)
        out.update({
            "alt_input_ids": alt_input_ids,
            "alt_attention_mask": alt_attention_mask,
            "alt_labels": alt_labels,
        })
    return out


def get_loaders(config: VLFMConfig, batch_size: int, num_workers: int = 0):
    ds_2 = load_dataset(
    "parquet",
    data_files={"train": "hf://datasets/babs/yoruba-speech-translation/**/*.parquet"})
    ds_2 = ds_2["train"]
    ds_1 = load_dataset("babs/s2s-translated-yoruba", split="train")

    ds_1 = ds_1.remove_columns([c for c in ds_1.column_names if c in {"yoruba_translation"}])
    ds_1 = ds_1.rename_columns({
        "english_speech": "audio",
        "english_text": "source_text",
        "yoruba_text": "target_text"
    })

    ds_1 = ds_1.cast_column("audio", Audio(sampling_rate=16000, decode=False))
    ds_2 = ds_2.cast_column("audio", Audio(sampling_rate=16000, decode=False))

    all_dataset = concatenate_datasets([ds_1, ds_2])
    split = all_dataset.shuffle(seed=42).train_test_split(test_size=0.01, seed=42)
    train_data, valid_data = split["train"], split["test"]
    print(f"train:{len(train_data)}, valid:{len(valid_data)}")

    train_ds = VLFMDataset(config, train_data, kd=False, train_on_inputs=False, max_response_tokens=None, use_augmentation=False)
    val_ds   = VLFMDataset(config, valid_data, kd=False, train_on_inputs=False, max_response_tokens=None, use_augmentation=False)

    pad_id = train_ds.tokenizer.pad_token_id

    train_loader = DataLoader(
        train_ds, batch_size=batch_size, shuffle=True, num_workers=num_workers,
        collate_fn=lambda b: vlfm_collate_fn(b, pad_id), pin_memory=True
    )
    valid_loader = DataLoader(
        val_ds, batch_size=batch_size, shuffle=False, num_workers=num_workers,
        collate_fn=lambda b: vlfm_collate_fn(b, pad_id), pin_memory=True
    )
    return train_loader, valid_loader


if __name__ == "__main__":
    config = VLFMConfig(audio_model_id="openai/whisper-large-v3", text_model_id="babs/t2t-improved")
    train_loader, valid_loader = get_loaders(config=config, batch_size=1)
    batch = next(iter(train_loader))
    print(batch)
    print({k: v.shape if hasattr(v, "shape") else type(v) for k, v in batch.items()})