data/dataset.py

#!/usr/bin/env python3
"""Turn the prepared CSV into tokenized PyTorch batches.

  1. Load ``data/processed/dataset.csv`` (written by ``data/download.py``).
  2. Read the class list from the ``labels.json`` sidecar and map names to ids.
  3. Split into train / validation / test, *stratified*.
  4. Wrap each split in a ``Dataset`` that tokenizes text.
  5. Hand back ``DataLoader``s producing padded, batched tensors.

VERIFY
------
    python -m data.dataset --inspect

Prints split sizes, per-split class balance, a decoded round-trip, and the exact
shape and dtype of every tensor in a batch. If those shapes are wrong, nothing
downstream will work.
"""

from __future__ import annotations

import argparse
import json
import logging
import sys
from pathlib import Path
from typing import Sequence

import pandas as pd
import torch
from torch.utils.data import DataLoader, Dataset
from transformers import AutoTokenizer, PreTrainedTokenizerBase

log = logging.getLogger("dataset")

DEFAULT_DATA = Path("data/processed/dataset.csv")
DEFAULT_MODEL = "distilbert-base-uncased"
LABELS_FILENAME = "labels.json"


# --------------------------------------------------------------------------
# Labels
# --------------------------------------------------------------------------


def load_labels(data_path: Path = DEFAULT_DATA) -> list[str]:
    """Read the ordered class list written alongside the CSV.

    The order defines which output position of the model means which class. It
    lives in one file, written once by ``download.py``, precisely so that
    training and inference cannot disagree. Re-deriving it independently in each
    script would mean a class absent from one split shifts every id, and the
    demo would confidently display the wrong class names.

    Falls back to sorted unique values in the CSV if the sidecar is missing.
    """
    sidecar = data_path.parent / LABELS_FILENAME
    if sidecar.exists():
        return list(json.loads(sidecar.read_text())["labels"])

    log.warning("%s missing — deriving labels from the CSV instead", sidecar)
    frame = pd.read_csv(data_path)
    column = "label" if "label" in frame.columns else "grade"
    return sorted(frame[column].astype(str).str.strip().unique())


def label_mappings(labels: Sequence[str]) -> tuple[dict[str, int], dict[int, str]]:
    """Build ``(name -> id, id -> name)`` from an ordered class list."""
    to_id = {name: i for i, name in enumerate(labels)}
    return to_id, {i: name for name, i in to_id.items()}


# --------------------------------------------------------------------------
# Loading and splitting
# --------------------------------------------------------------------------


def load_dataframe(path: Path = DEFAULT_DATA, labels: Sequence[str] | None = None) -> pd.DataFrame:
    """Read the prepared CSV and validate it.

    Accepts a legacy ``grade`` column name so the archived LendingClub sample in
    ``results/grade_null_result/`` still loads.
    """
    if not path.exists():
        raise FileNotFoundError(f"{path} not found. Run: python data/download.py")

    frame = pd.read_csv(path)
    if "label" not in frame.columns and "grade" in frame.columns:
        frame = frame.rename(columns={"grade": "label"})
    missing = {"text", "label"} - set(frame.columns)
    if missing:
        raise ValueError(f"{path} is missing column(s): {sorted(missing)}")

    frame = frame.dropna(subset=["text", "label"])
    frame["text"] = frame["text"].astype(str)
    frame["label"] = frame["label"].astype(str).str.strip()

    if labels is not None:
        unknown = set(frame["label"]) - set(labels)
        if unknown:
            raise ValueError(f"Labels in {path} not in labels.json: {sorted(unknown)[:8]}")

    log.info("Loaded %d rows, %d classes from %s", len(frame), frame["label"].nunique(), path)
    return frame.reset_index(drop=True)


def stratified_split(
    frame: pd.DataFrame,
    val_fraction: float = 0.15,
    test_fraction: float = 0.15,
    seed: int = 42,
) -> tuple[pd.DataFrame, pd.DataFrame, pd.DataFrame]:
    """Split into train/val/test keeping each class's proportion in every split.

    With 77 classes averaging ~170 rows each, a plain random split can easily
    leave a class with zero test examples, making its metrics undefined and the
    run non-comparable. Splitting *within* each class guarantees representation.

    Slices by position inside each class group rather than calling sklearn, so
    the logic is visible and it degrades gracefully for tiny classes.
    """
    if not 0 < val_fraction + test_fraction < 1:
        raise ValueError("val_fraction + test_fraction must be between 0 and 1")

    train_parts, val_parts, test_parts = [], [], []
    for _, group in frame.groupby("label"):
        group = group.sample(frac=1.0, random_state=seed)  # shuffle within class
        n = len(group)
        n_val = int(round(n * val_fraction))
        n_test = int(round(n * test_fraction))
        if n >= 3:  # guarantee representation when rounding would zero it out
            n_val, n_test = max(n_val, 1), max(n_test, 1)
        if n_val + n_test >= n:  # pathological: 1-2 rows in this class
            n_val, n_test = 0, min(1, n - 1)

        val_parts.append(group.iloc[:n_val])
        test_parts.append(group.iloc[n_val : n_val + n_test])
        train_parts.append(group.iloc[n_val + n_test :])

    def finish(parts: list[pd.DataFrame]) -> pd.DataFrame:
        # Reshuffle after concatenating, or the frame is ordered by class and
        # every batch would contain a single label.
        return pd.concat(parts).sample(frac=1.0, random_state=seed).reset_index(drop=True)

    train, val, test = finish(train_parts), finish(val_parts), finish(test_parts)
    log.info("Split: train=%d  val=%d  test=%d", len(train), len(val), len(test))
    return train, val, test


# --------------------------------------------------------------------------
# Dataset
# --------------------------------------------------------------------------


class TextClassificationDataset(Dataset):
    """One split, tokenized and ready for a ``DataLoader``.

    Tokenization happens once in ``__init__`` rather than per ``__getitem__``:
    for 13k short strings that costs well under a second and avoids re-running
    the tokenizer every epoch.

    No padding here on purpose. Each example keeps its true length and
    ``collate_batch`` pads per batch. Padding everything to a fixed length up
    front would leave most batches mostly padding, and the model would burn
    compute attending to nothing.
    """

    def __init__(
        self,
        texts: Sequence[str],
        labels: Sequence[str],
        label_to_id: dict[str, int],
        tokenizer: PreTrainedTokenizerBase,
        max_length: int = 64,
    ) -> None:
        """
        Args:
            texts: Raw input strings.
            labels: Class names, parallel to ``texts``.
            label_to_id: Name -> integer id mapping from ``label_mappings``.
            tokenizer: A HuggingFace tokenizer (DistilBERT WordPiece here).
            max_length: Hard truncation limit in tokens. Attention is quadratic
                in sequence length, so do not set this to 512 "just in case" —
                check the p99 printed by ``download.py``.
        """
        if len(texts) != len(labels):
            raise ValueError(f"texts ({len(texts)}) and labels ({len(labels)}) differ")

        self.texts = list(texts)
        self.labels = [label_to_id[name] for name in labels]
        self.max_length = max_length

        # Batch-encode the whole split at once. Each sequence already carries
        # DistilBERT's special tokens: [CLS] at position 0 and [SEP] at the end.
        # [CLS] matters later — the classifier reads its hidden state.
        encoded = tokenizer(
            self.texts, truncation=True, max_length=max_length,
            padding=False, add_special_tokens=True,
        )
        self.input_ids: list[list[int]] = encoded["input_ids"]
        self.attention_mask: list[list[int]] = encoded["attention_mask"]

        truncated = sum(1 for ids in self.input_ids if len(ids) >= max_length)
        if truncated:
            log.warning(
                "%d/%d examples (%.1f%%) hit the %d-token limit and were truncated",
                truncated, len(self.input_ids),
                100 * truncated / len(self.input_ids), max_length,
            )

    def __len__(self) -> int:
        return len(self.labels)

    def __getitem__(self, index: int) -> dict[str, torch.Tensor]:
        """Return one example as tensors.

        Shapes (``L`` = this example's own token count, varies per example):
            input_ids      (L,)  int64 — vocabulary indices
            attention_mask (L,)  int64 — all 1s here; collate adds the 0s
            labels         ()    int64 — a *scalar* tensor, the class id
        """
        return {
            "input_ids": torch.tensor(self.input_ids[index], dtype=torch.long),
            "attention_mask": torch.tensor(self.attention_mask[index], dtype=torch.long),
            "labels": torch.tensor(self.labels[index], dtype=torch.long),
        }


def collate_batch(
    batch: list[dict[str, torch.Tensor]], pad_token_id: int = 0
) -> dict[str, torch.Tensor]:
    """Stack variable-length examples into one rectangular batch.

    Pads every sequence to the longest one *in this batch* (dynamic padding).

    The attention mask is the important part: 1 for real tokens, 0 for padding.
    DistilBERT uses it to push padded positions to -inf before the attention
    softmax so they contribute nothing. Get this wrong and the model quietly
    attends to padding.

    Shapes out (``B`` = batch size, ``L`` = longest sequence in this batch):
        input_ids      (B, L)  int64
        attention_mask (B, L)  int64
        labels         (B,)    int64
    """
    lengths = [len(item["input_ids"]) for item in batch]
    max_len = max(lengths)

    input_ids = torch.full((len(batch), max_len), pad_token_id, dtype=torch.long)
    attention_mask = torch.zeros((len(batch), max_len), dtype=torch.long)

    for i, item in enumerate(batch):
        length = lengths[i]
        input_ids[i, :length] = item["input_ids"]
        attention_mask[i, :length] = item["attention_mask"]

    return {
        "input_ids": input_ids,
        "attention_mask": attention_mask,
        "labels": torch.stack([item["labels"] for item in batch]),
    }


# --------------------------------------------------------------------------
# Public entry point
# --------------------------------------------------------------------------


def build_dataloaders(
    data_path: Path = DEFAULT_DATA,
    model_name: str = DEFAULT_MODEL,
    batch_size: int = 32,
    max_length: int = 64,
    val_fraction: float = 0.15,
    test_fraction: float = 0.15,
    seed: int = 42,
    num_workers: int = 0,
) -> tuple[DataLoader, DataLoader, DataLoader, PreTrainedTokenizerBase, list[str]]:
    """Build train/val/test ``DataLoader``s from the prepared CSV.

    The one function ``train.py`` and ``evaluate.py`` call.

    Args:
        data_path: CSV produced by ``download.py``.
        model_name: Pretrained tokenizer to use; must match the model.
        batch_size: 32 fits a free Colab T4 easily; drop to 16 on CUDA OOM.
        max_length: Truncation limit in tokens.
        val_fraction: Validation share.
        test_fraction: Test share.
        seed: Shuffle/split seed. Same seed across modes = fair comparison.
        num_workers: Leave at 0 on Colab — with a dataset this small, spawning
            workers costs more than it saves.

    Returns:
        ``(train_loader, val_loader, test_loader, tokenizer, labels)``.
    """
    labels = load_labels(data_path)
    label_to_id, _ = label_mappings(labels)

    tokenizer = AutoTokenizer.from_pretrained(model_name)
    frame = load_dataframe(data_path, labels)
    train_df, val_df, test_df = stratified_split(frame, val_fraction, test_fraction, seed)

    def make_loader(split: pd.DataFrame, *, shuffle: bool) -> DataLoader:
        dataset = TextClassificationDataset(
            split["text"].tolist(), split["label"].tolist(),
            label_to_id, tokenizer, max_length,
        )
        return DataLoader(
            dataset,
            batch_size=batch_size,
            # Shuffle train so gradient steps see varied class mixes; never
            # shuffle val/test, so metrics are reproducible run to run.
            shuffle=shuffle,
            collate_fn=lambda b: collate_batch(b, tokenizer.pad_token_id or 0),
            num_workers=num_workers,
            drop_last=False,
        )

    return (
        make_loader(train_df, shuffle=True),
        make_loader(val_df, shuffle=False),
        make_loader(test_df, shuffle=False),
        tokenizer,
        labels,
    )


def class_weights(frame: pd.DataFrame, labels: Sequence[str]) -> torch.Tensor:
    """Inverse-frequency weights per class, shape ``(num_labels,)``.

    Optional, for ``nn.CrossEntropyLoss(weight=...)``. Raises macro F1 and lowers
    raw accuracy. Which you want depends on whether rare-class mistakes cost more.
    """
    counts = frame["label"].value_counts().reindex(list(labels), fill_value=0)
    weights = torch.tensor(
        [len(frame) / (len(labels) * (c + 1)) for c in counts], dtype=torch.float
    )
    return weights / weights.mean()


# --------------------------------------------------------------------------
# CLI — verification
# --------------------------------------------------------------------------


def inspect(args: argparse.Namespace) -> int:
    """Print everything needed to eyeball whether the data step is correct."""
    train_loader, val_loader, test_loader, tokenizer, labels = build_dataloaders(
        data_path=args.data, model_name=args.model, batch_size=args.batch_size,
        max_length=args.max_length, seed=args.seed,
    )
    _, id_to_label = label_mappings(labels)

    print("\n" + "=" * 68)
    print("  DATA PIPELINE INSPECTION")
    print("=" * 68)
    print(f"\n  Classes: {len(labels)}  (first 5: {labels[:5]})")
    print(f"  Tokenizer: {args.model}  (vocab {tokenizer.vocab_size:,}, "
          f"pad id {tokenizer.pad_token_id})")

    print("\n  Split sizes and class balance")
    print("  " + "-" * 58)
    for name, loader in (("train", train_loader), ("val", val_loader), ("test", test_loader)):
        ids = pd.Series(loader.dataset.labels)
        share = ids.value_counts(normalize=True)
        print(f"    {name:<6} {len(ids):>6} rows | {ids.nunique():>3} classes present | "
              f"largest class {share.max():.1%} | smallest {share.min():.2%}")
    print("  " + "-" * 58)
    print("    (largest-class share should match across splits — stratification working)")

    # Round-trip: raw text -> ids -> back. If the decoded version does not
    # resemble the original, tokenization is misconfigured.
    example = train_loader.dataset[0]
    print("\n  One example, round-tripped")
    print("  " + "-" * 58)
    print(f"    raw text : {train_loader.dataset.texts[0][:64]!r}")
    print(f"    input_ids: {example['input_ids'].tolist()[:14]}"
          f"{' ...' if len(example['input_ids']) > 14 else ''}")
    print(f"    decoded  : {tokenizer.decode(example['input_ids'])[:64]!r}")
    print(f"    label    : {example['labels'].item()} -> "
          f"{id_to_label[example['labels'].item()]}")

    batch = next(iter(train_loader))
    print("\n  One collated batch")
    print("  " + "-" * 58)
    for key, tensor in batch.items():
        print(f"    {key:<16} {str(tuple(tensor.shape)):<12} {tensor.dtype}")
    real = int(batch["attention_mask"].sum())
    total = batch["attention_mask"].numel()
    print(f"\n    Real tokens {real}/{total} ({100 * real / total:.0f}%) — "
          f"rest is padding masked to 0.")
    print("\n" + "=" * 68 + "\n")
    return 0


def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace:
    """Define and parse the command-line interface."""
    p = argparse.ArgumentParser(description="Inspect the tokenized data pipeline.")
    p.add_argument("--inspect", action="store_true", help="Print the verification report.")
    p.add_argument("--data", type=Path, default=DEFAULT_DATA)
    p.add_argument("--model", default=DEFAULT_MODEL)
    p.add_argument("--batch-size", type=int, default=32)
    p.add_argument("--max-length", type=int, default=64)
    p.add_argument("--seed", type=int, default=42)
    return p.parse_args(argv)


def main(argv: Sequence[str] | None = None) -> int:
    """Entry point. Returns a process exit code."""
    logging.basicConfig(level=logging.INFO, format="%(levelname)-7s %(message)s")
    for noisy in ("httpx", "urllib3", "filelock", "huggingface_hub"):
        logging.getLogger(noisy).setLevel(logging.WARNING)
    args = parse_args(argv)
    if not args.inspect:
        print("Nothing to do. Try: python -m data.dataset --inspect")
        return 0
    try:
        return inspect(args)
    except (FileNotFoundError, ValueError) as exc:
        log.error("%s", exc)
        return 1


if __name__ == "__main__":
    sys.exit(main())