cas9/data/data.py

# Copyright (c) Meta Platforms, Inc.
# All rights reserved.
#
# This source code is licensed under the CC-by-NC license found in the
# LICENSE file in the root directory of this source tree.
# Part of this implementation is adapted from https://github.com/louaaron/Score-Entropy-Discrete-Diffusion
# which is released under MIT license

from dataclasses import dataclass, field
from typing import Dict, Iterable, List, Tuple, Optional

import os

from datasets import Dataset as HFDataset, DatasetDict, load_dataset
from omegaconf import OmegaConf

import torch
from torch.utils.data import DataLoader, Dataset
from transformers import AutoTokenizer
try:
    from transformers import EsmTokenizer
    _HAVE_ESM = True
except Exception:
    _HAVE_ESM = False

from cas9.data.tokenizer import wt_detokenizer
from cas9.data.utils import cycle_loader, StatefulDistributedSampler


# -------------------------
# Tokenizer helper
# -------------------------

def _get_tokenizer(name: Optional[str]):
    """
    Returns a HF tokenizer by name. Supports ESM when available.
    """
    tok_name = name or "gpt2"
    if _HAVE_ESM and tok_name.startswith("facebook/esm"):
        return EsmTokenizer.from_pretrained(tok_name)
    return AutoTokenizer.from_pretrained(tok_name)


# -------------------------
# HuggingFace datasets (ragged)
# -------------------------

def _get_hf_dataset_ragged(
    name: str,
    mode: str,
    cache_dir: Optional[str],
    num_proc: int,
    tokenizer_name: Optional[str],
) -> HFDataset:
    """
    Loads an HF dataset split and tokenizes into variable-length sequences.
    No chunking to fixed block_size; each example is one tokenized text (+ EOS if available).
    Returns a HuggingFace Dataset with a single column "input_ids" (list[int]) per row.
    """
    detokenizer = None
    if name == "wikitext103":
        ds = load_dataset("wikitext", name="wikitext-103-raw-v1", cache_dir=cache_dir)[mode]
        detokenizer = wt_detokenizer
    elif name == "fineweb-edu":
        ds = load_dataset("HuggingFaceFW/fineweb-edu", name="CC-MAIN-2024-10", cache_dir=cache_dir)[mode]
    else:
        ds = load_dataset(name, cache_dir=cache_dir)[mode]

    tok = _get_tokenizer(tokenizer_name)
    eos_id = None
    if hasattr(tok, "eos_token") and tok.eos_token is not None:
        eos_id = tok.convert_tokens_to_ids(tok.eos_token)
        if isinstance(eos_id, list):
            eos_id = eos_id[0]

    def _apply_detok(batch: Dict):
        texts = batch["text"]
        if detokenizer is not None:
            texts = [detokenizer(t) for t in texts]
        return {"text": texts}

    if "text" in ds.column_names:
        ds = ds.map(_apply_detok, batched=True, num_proc=num_proc, load_from_cache_file=True)

    def _tokenize(batch: Dict):
        # Use fast tokenizer; do not return attention masks
        ids = tok(batch["text"], add_special_tokens=True, return_attention_mask=False)["input_ids"]
        # if eos_id is not None:
        #     for arr in ids:
        #         arr.append(eos_id)
        return {"input_ids": ids}

    # Tokenize
    tokenized = ds.map(
        _tokenize,
        batched=True,
        num_proc=num_proc,
        remove_columns=[c for c in ds.column_names if c != "text"],
        load_from_cache_file=True,
    )

    # Keep only input_ids
    if "text" in tokenized.column_names:
        tokenized = tokenized.remove_columns("text")

    # For fineweb there can be extra meta columns; keep only input_ids
    extra_cols = [c for c in tokenized.column_names if c != "input_ids"]
    if extra_cols:
        tokenized = tokenized.remove_columns(extra_cols)

    tokenized = tokenized.with_format("python")  # return plain python objects; collate will tensorize
    return tokenized


# -------------------------
# Local FASTA dataset (ragged)
# -------------------------

class FASTADataset(Dataset):
    """
    Loads a local FASTA file and tokenizes each sequence to variable-length ids.
    One HF-like item per sequence: {"input_ids": List[int]}.
    """

    def __init__(self, fasta_path: str, tokenizer_name: Optional[str]):
        assert os.path.isfile(fasta_path), f"FASTA file not found: {fasta_path}"
        self.fasta_path = fasta_path
        self.tok = _get_tokenizer(tokenizer_name)
        self.eos_id = None
        if hasattr(self.tok, "eos_token") and self.tok.eos_token is not None:
            eid = self.tok.convert_tokens_to_ids(self.tok.eos_token)
            self.eos_id = eid[0] if isinstance(eid, list) else eid

        # Parse FASTA (simple parser; no Biopython dependency)
        self._seqs: List[str] = []
        with open(fasta_path, "r") as f:
            cur = []
            for line in f:
                line = line.strip()
                if not line:
                    continue
                if line.startswith(">"):
                    if cur:
                        self._seqs.append("".join(cur))
                        cur = []
                else:
                    cur.append(line)
            if cur:
                self._seqs.append("".join(cur))

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

    def __getitem__(self, idx: int) -> Dict[str, List[int]]:
        seq = self._seqs[idx]
        ids = self.tok(seq, add_special_tokens=True, return_attention_mask=False)["input_ids"]
        if isinstance(ids, list) and len(ids) > 0 and isinstance(ids[0], list):
            # Some tokenizers may return nested lists for per-char tokenization; flatten
            ids = [t for sub in ids for t in (sub if isinstance(sub, list) else [sub])]
        # if self.eos_id is not None:
        #     ids.append(self.eos_id)
        return {"input_ids": ids}


# -------------------------
# Ragged collate
# -------------------------

def ragged_collate(batch: List[Dict[str, List[int]]]) -> Dict[str, List[torch.Tensor]]:
    """
    Collate a list of {"input_ids": List[int]} into {"input_ids": List[LongTensor]}.
    No padding; tensorization only.
    """
    out: List[torch.Tensor] = []
    for item in batch:
        ids = item["input_ids"]
        out.append(torch.tensor(ids, dtype=torch.long))
    return {"input_ids": out}


# -------------------------
# Public dataclasses
# -------------------------

@dataclass
class DatasetWrap:
    dataset: Dataset  # HF Dataset (python format) OR torch Dataset
    sampler: StatefulDistributedSampler


@dataclass
class DataState:
    train: DatasetWrap
    test: DatasetWrap


# -------------------------
# Builders
# -------------------------

def _get_dataset(
    name: str,
    mode: str,
    cache_dir: Optional[str],
    num_proc: int,
    batch_size: int,
    ngpus: int,
    tokenizer_name: Optional[str],
    fasta_path: Optional[str] = None,
) -> DatasetWrap:
    """
    Build either an HF dataset (ragged) or a local FASTA dataset.
    - name == "fasta" -> requires fasta_path
    - otherwise -> HF dataset by name
    """
    assert batch_size % ngpus == 0, f"{mode} batch size must be divisible by number of gpus."

    if name.lower() == "fasta":
        assert fasta_path is not None, "For name='fasta', you must set data.fasta_path in the config."
        dataset = FASTADataset(fasta_path=fasta_path, tokenizer_name=tokenizer_name)
    else:
        dataset = _get_hf_dataset_ragged(
            name=name,
            mode=mode,
            cache_dir=cache_dir,
            num_proc=num_proc,
            tokenizer_name=tokenizer_name,
        )

    sampler = StatefulDistributedSampler(dataset=dataset)
    return DatasetWrap(dataset=dataset, sampler=sampler)


def get_data_state(config: OmegaConf) -> DataState:
    """
    Expects (typical):
      config.data.train / config.data.valid : dataset name (e.g., "wikitext103", "fineweb-edu", "fasta", or any HF hub dataset id)
      config.data.cache_dir
      config.data.num_workers
      config.data.tokenizer (string, e.g., "gpt2" or "facebook/esm2_t33_650M_UR50S")
      config.data.fasta_path (optional, used when name == "fasta")
      config.training.batch_size
      config.eval.batch_size
      config.compute.ngpus
    """
    tok_name = getattr(config.data, "tokenizer", None)

    train = _get_dataset(
        name=config.data.train,
        mode="train",
        cache_dir=config.data.cache_dir,
        num_proc=config.data.num_workers,
        batch_size=config.training.batch_size,
        ngpus=config.compute.ngpus,
        tokenizer_name=tok_name,
        fasta_path=getattr(config.data, "fasta_path", None),
    )

    valid = _get_dataset(
        name=config.data.valid,
        mode="validation",
        cache_dir=config.data.cache_dir,
        num_proc=config.data.num_workers,
        batch_size=config.eval.batch_size,
        ngpus=config.compute.ngpus,
        tokenizer_name=tok_name,
        fasta_path=getattr(config.data, "fasta_path", None),
    )

    return DataState(train=train, test=valid)


def get_data_loaders(
    config: OmegaConf,
    data_state: DataState,
) -> Tuple[Iterable, Iterable]:
    """
    Returns infinite iterators over train/valid using ragged_collate.
    Each batch item is a dict: {"input_ids": List[LongTensor]}.
    """
    per_gpu_train_bs = config.training.batch_size // config.compute.ngpus
    per_gpu_valid_bs = config.eval.batch_size // config.compute.ngpus

    train_loader = cycle_loader(
        DataLoader(
            data_state.train.dataset,
            batch_size=per_gpu_train_bs,
            sampler=data_state.train.sampler,
            num_workers=config.data.num_workers,
            pin_memory=True,
            shuffle=(data_state.train.sampler is None),
            persistent_workers=True,
            collate_fn=ragged_collate,
        )
    )

    valid_loader = cycle_loader(
        DataLoader(
            data_state.test.dataset,
            batch_size=per_gpu_valid_bs,
            sampler=data_state.test.sampler,
            num_workers=config.data.num_workers,
            pin_memory=True,
            shuffle=(data_state.test.sampler is None),
            collate_fn=ragged_collate,
        )
    )

    return iter(train_loader), iter(valid_loader)