dataset.py

"""
MMT-JEPA datasets — one class per JEPA objective.

    ObjA  Audio -> Text   (both languages)
          English: openslr/librispeech_asr  validation  (default; was train.100)
          Twi:     ghananlpcommunity/twi-speech-text-multispeaker-16k
                   + BibleTTS local (openslr.org/129)

    ObjB  Text -> Text    (both languages, translation)
          ghananlpcommunity/twi-english-paragraph-dataset_news

    ObjC  Text -> Audio   (both languages)
          Same audio sources as ObjA, direction flipped

Quick start (full data):
    tok = MyTokenizer()
    for batch in ObjA(tok).loader(): ...
    for batch in ObjB(tok).loader(): ...
    for batch in ObjC(tok).loader(): ...

Quick start (proof-of-concept, small):
    for batch in ObjA(tok, max_samples=500).loader(batch_size=16): ...
    for batch in ObjB(tok, max_samples=1000).loader(batch_size=16): ...
    for batch in ObjC(tok, max_samples=500).loader(batch_size=16): ...
"""

from __future__ import annotations
import io
import random
from functools import lru_cache

import numpy as np
import soundfile as sf
import torch
import torchaudio
from torch import Tensor
from torch.nn.utils.rnn import pad_sequence
from torch.utils.data import ConcatDataset, DataLoader, Dataset


def _loader_persistent_kw(num_workers: int) -> dict:
    if num_workers > 0:
        return {"persistent_workers": True}
    return {}

ENG, TWI    = 0, 1
TEXT, AUDIO = 0, 1
PAD = 0


def _mel_width_to_stem_steps(T: int) -> int:
    """
    Time bins in mel (n_mels, T) → sequence length after AudioStem (conv×2 + AvgPool1d 2).
    Must stay in sync with ``AudioStem`` in model.py (not the old [::2] shortcut).
    """
    if T <= 0:
        return 0
    # Two Conv1d stride=1 pad=1 preserve length; AvgPool1d(2,2), pad=0
    return (T - 2) // 2 + 1


@lru_cache(maxsize=16)
def _mel_transforms(target_sr: int, n_mels: int):
    """Cached mel + dB transforms (same params as former librosa pipeline)."""
    return (
        torchaudio.transforms.MelSpectrogram(
            sample_rate=target_sr,
            n_fft=400,
            hop_length=160,
            n_mels=n_mels,
            power=2.0,
        ),
        torchaudio.transforms.AmplitudeToDB(stype="power", top_db=80.0),
    )


def make_mel(audio, src_sr: int, cfg=None,
             target_sr: int = 16_000, n_mels: int = 80) -> Tensor | None:
    """
    audio: numpy array or torch Tensor, any shape — mono or stereo.
    Returns (n_mels, T) float32 tensor, z-score normalised. None if too long/short.
    Accepts an optional ModelConfig; if provided, uses cfg.sample_rate and cfg.n_mels.

    Uses torchaudio (not librosa) so DataLoader workers do not import numba, which
    breaks under NumPy 2.3+ with older numba builds.
    """
    if cfg is not None:
        target_sr = cfg.sample_rate
        n_mels    = cfg.n_mels
        # AudioStem downsamples by 2, so we allow 2x max_seq_len frames
        max_frames = cfg.max_seq_len * 2
    else:
        max_frames = 3000

    if isinstance(audio, Tensor):
        y = audio.detach().cpu().float()
    else:
        y = torch.from_numpy(np.asarray(audio, dtype=np.float32))

    if y.ndim > 1:
        y = y.mean(dim=0)

    n = y.numel()
    if n == 0 or n > target_sr * 60:
        return None

    if src_sr != target_sr:
        y = torchaudio.functional.resample(y.unsqueeze(0), src_sr, target_sr).squeeze(0)

    mel_spec, to_db = _mel_transforms(target_sr, n_mels)
    mel_db = to_db(mel_spec(y).float())
    mel_t = mel_db[:, :max_frames]

    return (mel_t - mel_t.mean()) / (mel_t.std() + 1e-6)



class _LibriSpeech(Dataset):
    """
    English audio + transcript from openslr/librispeech_asr.

    Parameters
    ----------
    split       : HuggingFace split name.
                  'validation' (~2.7k utterances) is recommended for POC runs.
                  'train.100' (~28k) for full training.
    max_samples : if set, only the first N utterances are kept after loading.
    """

    def __init__(
        self,
        tokenizer,
        cfg=None,
        max_text: int = 256,
        split: str = "validation",
        max_samples: int | None = None,
    ) -> None:
        from datasets import Audio, load_dataset

        ds = load_dataset("openslr/librispeech_asr", "clean", split=split)

        if max_samples is not None:
            ds = ds.select(range(min(max_samples, len(ds))))

        self.ds      = ds.cast_column("audio", Audio(decode=False))
        self.tok     = tokenizer
        self.cfg     = cfg
        self.max_txt = cfg.max_seq_len if cfg is not None else max_text
        print(f"  LibriSpeech ({split}): {len(self.ds):,} utterances"
              + (f"  [capped at {max_samples}]" if max_samples else ""))

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

    def __getitem__(self, i):
        row = self.ds[i]
        try:
            info = row["audio"]
            raw  = info.get("bytes") or open(info["path"], "rb").read()
            audio, sr = sf.read(io.BytesIO(raw))
            audio = audio.astype(np.float32)
        except Exception:
            return None
        mel = make_mel(audio, sr, self.cfg)
        if mel is None:
            return None
        ids = self.tok.encode(row["text"])[:self.max_txt]
        return {"mel": mel, "ids": torch.tensor(ids, dtype=torch.long), "lang": ENG}


class _TwiAudio(Dataset):
    """
    Twi audio + transcript.
    Sources:
      1. BibleTTS local  (set bibletts_dir, expects train/<stem>.flac + .txt)
      2. ghananlpcommunity/twi-speech-text-multispeaker-16k  (HuggingFace)

    Parameters
    ----------
    max_samples : if set, records list is trimmed to N after loading all sources.
    """

    def __init__(
        self,
        tokenizer,
        cfg=None,
        split: str = "train",
        max_text: int = 256,
        max_samples: int | None = None,
    ) -> None:
        self.tok     = tokenizer
        self.cfg     = cfg
        self.max_txt = cfg.max_seq_len if cfg is not None else max_text
        self.records: list[dict] = []

        try:
            from datasets import Audio, load_dataset
            # decode=False avoids torchcodec; we decode bytes with soundfile
            ds = load_dataset("ghananlpcommunity/twi-speech-text-multispeaker-16k", split=split)
            ds = ds.cast_column("audio", Audio(decode=False))
            before = len(self.records)
            for row in ds:
                text = (
                    row.get("text")
                    or row.get("sentence")
                    or row.get("transcription")
                    or ""
                ).strip()
                if not text:
                    continue
                info = row["audio"]
                raw  = info.get("bytes") or open(info["path"], "rb").read()
                self.records.append({"bytes": raw, "text": text})
            print(f"  Twi HuggingFace ({split}): +{len(self.records) - before:,} utterances")
        except Exception as e:
            print(f"  Twi HuggingFace skipped ({e})")

        if max_samples is not None and len(self.records) > max_samples:
            self.records = self.records[:max_samples]
            print(f"  Twi records capped at {max_samples}")

        assert self.records, (
            "No Twi audio found.\n"
            "  Ensure internet access for the HuggingFace dataset, or set bibletts_dir."
        )

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

    def __getitem__(self, i):
        rec = self.records[i]
        try:
            audio, sr = sf.read(io.BytesIO(rec["bytes"]))
            audio = audio.astype(np.float32)
        except Exception:
            return None
        mel = make_mel(audio, sr, self.cfg)
        if mel is None:
            return None
        ids = self.tok.encode(rec["text"])[:self.max_txt]
        if not ids:
            return None
        return {"mel": mel, "ids": torch.tensor(ids, dtype=torch.long), "lang": TWI}



class ObjA(Dataset):
    """
    Context : audio  (English or Twi)
    Target  : text   (same language)

    Batch keys
    ----------
    ctx_audio    (B, 80, T)  |  ctx_text     None
    tgt_text     (B, L)      |  tgt_audio    None
    ctx_pad_mask (B, stem_T) |  tgt_pad_mask (B, L)   # stem_T = f(mel width), see _mel_width_to_stem_steps
    src_lang / tgt_lang / src_mod (=AUDIO) / tgt_mod (=TEXT)

    Parameters
    ----------
    max_samples : cap applied independently to each source (English + Twi).
    libri_split : LibriSpeech split; 'validation' for POC, 'train.100' for full.
    """

    def __init__(
        self,
        tokenizer,
        cfg=None,
        max_text: int = 256,
        max_samples: int | None = None,
        libri_split: str = "validation",
    ) -> None:
        print("ObjA sources:")
        eng = _LibriSpeech(tokenizer, cfg, max_text,
                           split=libri_split, max_samples=max_samples)
        twi = _TwiAudio(tokenizer, cfg, max_text=max_text,
                        max_samples=max_samples)
        self.ds = ConcatDataset([eng, twi])

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

    def __getitem__(self, i):
        item = self.ds[i]
        if item is None:
            return None
        return {
            "mel":      item["mel"],    # audio is context
            "ctx_ids":  None,
            "tgt_ids":  item["ids"],    # text is target
            "src_lang": item["lang"],
            "tgt_lang": item["lang"],
            "src_mod":  AUDIO,
            "tgt_mod":  TEXT,
        }

    def loader(self, batch_size: int = 32, num_workers: int = 0) -> DataLoader:
        return DataLoader(
            self, batch_size, shuffle=True,
            num_workers=num_workers,
            collate_fn=_collate_audio_text,
            drop_last=True,
            **_loader_persistent_kw(num_workers),
        )


class ObjB(Dataset):
    """
    Context : text  (English or Twi)
    Target  : text  (other language — translation)
    Direction randomly flipped each call so model sees both eng->twi and twi->eng.

    Batch keys
    ----------
    ctx_audio    None        |  ctx_text     (B, L)
    tgt_audio    None        |  tgt_text     (B, L)
    ctx_pad_mask (B, L)      |  tgt_pad_mask (B, L)
    src_lang / tgt_lang / src_mod (=TEXT) / tgt_mod (=TEXT)

    Parameters
    ----------
    max_samples  : total pairs kept across all sources (applied after loading).
    reverse_prob : probability of flipping to twi->eng direction per sample.
    """

    SOURCES = [
        ("ghananlpcommunity/twi-english-paragraph-dataset_news", {}, "ENGLISH", "TWI"),
        ("ghananlpcommunity/english-twi-sentences-non-nouns",    {}, "english", "twi"),
        ("ghananlpcommunity/english-twi-nouns-v2",               {}, "english", "twi"),
        # add more (repo, kwargs, eng_col, twi_col) here as datasets become available
    ]

    def __init__(
        self,
        tokenizer,
        cfg=None,
        max_text: int = 256,
        max_samples: int | None = None,
        reverse_prob: float = 0.5,
    ) -> None:
        from datasets import load_dataset

        self.tok          = tokenizer
        self.max_txt      = cfg.max_seq_len if cfg is not None else max_text
        self.reverse_prob = reverse_prob
        self.pairs: list[tuple[str, str]] = []

        print("ObjB sources:")
        for repo, kwargs, eng_col, twi_col in self.SOURCES:
            try:
                ds = load_dataset(repo, split="train", **kwargs)
                before = len(self.pairs)
                for row in ds:
                    eng = str(row.get(eng_col) or "").strip()
                    twi = str(row.get(twi_col) or "").strip()
                    if eng and twi:
                        self.pairs.append((eng, twi))
                print(f"  {repo.split('/')[1]}: +{len(self.pairs) - before:,} pairs")
            except Exception as e:
                print(f"  {repo.split('/')[1]} skipped ({e})")

        if max_samples is not None and len(self.pairs) > max_samples:
            random.shuffle(self.pairs)           # shuffle before capping for variety
            self.pairs = self.pairs[:max_samples]
            print(f"  ObjB pairs capped at {max_samples}")

        assert self.pairs, "ObjB: no data loaded — check internet or sources list"

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

    def __getitem__(self, i):
        eng, twi = self.pairs[i]
        if random.random() < self.reverse_prob:
            ctx, tgt, sl, tl = twi, eng, TWI, ENG
        else:
            ctx, tgt, sl, tl = eng, twi, ENG, TWI
        return {
            "ctx_ids":  torch.tensor(self.tok.encode(ctx)[:self.max_txt], dtype=torch.long),
            "tgt_ids":  torch.tensor(self.tok.encode(tgt)[:self.max_txt], dtype=torch.long),
            "src_lang": sl,
            "tgt_lang": tl,
            "src_mod":  TEXT,
            "tgt_mod":  TEXT,
        }

    def loader(self, batch_size: int = 32, num_workers: int = 0) -> DataLoader:
        return DataLoader(
            self, batch_size, shuffle=True,
            num_workers=num_workers,
            collate_fn=_collate_text_text,
            drop_last=True,
            **_loader_persistent_kw(num_workers),
        )


class ObjC(Dataset):
    """
    Context : text   (English or Twi)
    Target  : audio  (same language)
    Same underlying data as ObjA, direction flipped.

    Batch keys
    ----------
    ctx_audio    None        |  ctx_text     (B, L)
    tgt_audio    (B, 80, T)  |  tgt_text     None
    ctx_pad_mask (B, L)      |  tgt_pad_mask (B, stem_T)
    src_lang / tgt_lang / src_mod (=TEXT) / tgt_mod (=AUDIO)

    Parameters
    ----------
    max_samples : cap applied independently to each source (English + Twi).
    libri_split : LibriSpeech split; 'validation' for POC, 'train.100' for full.
    """

    def __init__(
        self,
        tokenizer,
        cfg=None,
        max_text: int = 256,
        max_samples: int | None = None,
        libri_split: str = "validation",
    ) -> None:
        print("ObjC sources:")
        eng = _LibriSpeech(tokenizer, cfg, max_text,
                           split=libri_split, max_samples=max_samples)
        twi = _TwiAudio(tokenizer, cfg, max_text=max_text,
                        max_samples=max_samples)
        self.ds = ConcatDataset([eng, twi])

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

    def __getitem__(self, i):
        item = self.ds[i]
        if item is None:
            return None
        return {
            "mel":      item["mel"],    # audio is target
            "ctx_ids":  item["ids"],    # text is context
            "tgt_ids":  None,
            "src_lang": item["lang"],
            "tgt_lang": item["lang"],
            "src_mod":  TEXT,
            "tgt_mod":  AUDIO,
        }

    def loader(self, batch_size: int = 32, num_workers: int = 0) -> DataLoader:
        return DataLoader(
            self, batch_size, shuffle=True,
            num_workers=num_workers,
            collate_fn=_collate_text_audio,
            drop_last=True,
            **_loader_persistent_kw(num_workers),
        )



def _collate_audio_text(batch):
    """Audio context, text target  (ObjA)."""
    batch = [b for b in batch if b is not None]
    if not batch:
        return {}
    max_T  = max(b["mel"].shape[1] for b in batch)
    n_mels = batch[0]["mel"].shape[0]
    ctx_audio = torch.zeros(len(batch), n_mels, max_T)
    for i, b in enumerate(batch):
        t = b["mel"].shape[1]
        ctx_audio[i, :, :t] = b["mel"]
    tgt_text = pad_sequence(
        [b["tgt_ids"] for b in batch], batch_first=True, padding_value=PAD
    )
    stem_T = _mel_width_to_stem_steps(max_T)
    ctx_mask_ds = torch.ones(len(batch), stem_T, dtype=torch.bool)
    for i, b in enumerate(batch):
        t = b["mel"].shape[1]
        ctx_mask_ds[i, : _mel_width_to_stem_steps(t)] = False
    return {
        "ctx_audio":    ctx_audio,
        "ctx_text":     None,
        "tgt_audio":    None,
        "tgt_text":     tgt_text,
        "ctx_pad_mask": ctx_mask_ds,
        "tgt_pad_mask": tgt_text == PAD,
        "src_lang": torch.tensor([b["src_lang"] for b in batch]),
        "tgt_lang": torch.tensor([b["tgt_lang"] for b in batch]),
        "src_mod":  torch.tensor([b["src_mod"]  for b in batch]),
        "tgt_mod":  torch.tensor([b["tgt_mod"]  for b in batch]),
    }


def _collate_text_text(batch):
    """Text context, text target  (ObjB)."""
    batch = [b for b in batch if b is not None]
    if not batch:
        return {}
    ctx = pad_sequence(
        [b["ctx_ids"] for b in batch], batch_first=True, padding_value=PAD
    )
    tgt = pad_sequence(
        [b["tgt_ids"] for b in batch], batch_first=True, padding_value=PAD
    )
    return {
        "ctx_audio":    None,
        "ctx_text":     ctx,
        "tgt_audio":    None,
        "tgt_text":     tgt,
        "ctx_pad_mask": ctx == PAD,
        "tgt_pad_mask": tgt == PAD,
        "src_lang": torch.tensor([b["src_lang"] for b in batch]),
        "tgt_lang": torch.tensor([b["tgt_lang"] for b in batch]),
        "src_mod":  torch.tensor([b["src_mod"]  for b in batch]),
        "tgt_mod":  torch.tensor([b["tgt_mod"]  for b in batch]),
    }


def _collate_text_audio(batch):
    """Text context, audio target  (ObjC)."""
    batch = [b for b in batch if b is not None]
    if not batch:
        return {}
    ctx_text  = pad_sequence(
        [b["ctx_ids"] for b in batch], batch_first=True, padding_value=PAD
    )
    max_T     = max(b["mel"].shape[1] for b in batch)
    n_mels    = batch[0]["mel"].shape[0]
    tgt_audio = torch.zeros(len(batch), n_mels, max_T)
    for i, b in enumerate(batch):
        t = b["mel"].shape[1]
        tgt_audio[i, :, :t] = b["mel"]
    stem_T = _mel_width_to_stem_steps(max_T)
    tgt_mask_ds = torch.ones(len(batch), stem_T, dtype=torch.bool)
    for i, b in enumerate(batch):
        t = b["mel"].shape[1]
        tgt_mask_ds[i, : _mel_width_to_stem_steps(t)] = False
    return {
        "ctx_audio":    None,
        "ctx_text":     ctx_text,
        "tgt_audio":    tgt_audio,
        "tgt_text":     None,
        "ctx_pad_mask": ctx_text == PAD,
        "tgt_pad_mask": tgt_mask_ds,
        "src_lang": torch.tensor([b["src_lang"] for b in batch]),
        "tgt_lang": torch.tensor([b["tgt_lang"] for b in batch]),
        "src_mod":  torch.tensor([b["src_mod"]  for b in batch]),
        "tgt_mod":  torch.tensor([b["tgt_mod"]  for b in batch]),
    }



if __name__ == "__main__":

    class _Tok:
        def encode(self, text: str) -> list[int]:
            return [ord(c) % 20_000 for c in text[:64]]

    tok = _Tok()

    # --- change these to None / "train.100" for a full run ---
    MAX = 50          # samples per source
    BS  = 8           # batch size for the smoke test

    def _grab(ds, n=12):
        items = []
        for i in range(min(len(ds), 500)):
            x = ds[i]
            if x is not None:
                items.append(x)
            if len(items) == n:
                break
        return items

    print("\n── Obj A: Audio -> Text (both languages) ──")
    try:
        ds = ObjA(tok, max_samples=MAX)
        b  = _collate_audio_text(_grab(ds, BS))
        print(f"  ctx_audio    : {b['ctx_audio'].shape}")
        print(f"  tgt_text     : {b['tgt_text'].shape}")
        print(f"  src_mod      : {b['src_mod'].tolist()}  (1=audio)")
        print(f"  tgt_mod      : {b['tgt_mod'].tolist()}  (0=text)")
        print(f"  src_lang mix : {b['src_lang'].tolist()}  (0=eng 1=twi)")
        print("  PASS")
    except Exception as e:
        print(f"  FAIL — {e}")

    print("\n── Obj B: Text -> Text (translation) ──")
    try:
        ds = ObjB(tok, max_samples=MAX * 2)
        b  = _collate_text_text(_grab(ds, BS))
        print(f"  ctx_text     : {b['ctx_text'].shape}")
        print(f"  tgt_text     : {b['tgt_text'].shape}")
        print(f"  src_mod      : {b['src_mod'].tolist()}  (0=text)")
        print(f"  src_lang mix : {b['src_lang'].tolist()}  (0=eng 1=twi)")
        print("  PASS")
    except Exception as e:
        print(f"  FAIL — {e}")

    print("\n── Obj C: Text -> Audio (both languages) ──")
    try:
        ds = ObjC(tok, max_samples=MAX)
        b  = _collate_text_audio(_grab(ds, BS))
        print(f"  ctx_text     : {b['ctx_text'].shape}")
        print(f"  tgt_audio    : {b['tgt_audio'].shape}")
        print(f"  src_mod      : {b['src_mod'].tolist()}  (0=text)")
        print(f"  tgt_mod      : {b['tgt_mod'].tolist()}  (1=audio)")
        print(f"  src_lang mix : {b['src_lang'].tolist()}  (0=eng 1=twi)")
        print("  PASS")
    except Exception as e:
        print(f"  FAIL — {e}")