#!/usr/bin/env python3 import argparse import numpy as np import torch from torch.utils.data import Dataset, DataLoader, Sampler, BatchSampler import torch.distributed as dist from lightning import LightningDataModule from pathlib import Path from multiprocessing import cpu_count import random import pandas as pd import shelve from torch.nn.utils.rnn import pad_sequence from typing import List, Iterable, Sequence import sys import rootutils import logging import math from dpacman.utils import pylogger root = rootutils.setup_root(__file__, indicator=".project-root", pythonpath=True) logger = pylogger.RankedLogger(__name__, rank_zero_only=True) class PreBatchedDistributedBatchSampler(BatchSampler): """ Accepts a precomputed list of batches (list[list[int]]) and shards them across DDP ranks. - shuffle_batch_order: shuffle order of batches each epoch (deterministic via set_epoch) - drop_last: drop remainder so each rank gets same #steps """ def __init__(self, batches, shuffle_batch_order=False, drop_last=False, seed: int = 0): # super expects attributes batch_size, drop_last, sampler – but we don't need them. # We only need to subclass BatchSampler to satisfy Lightning's check. self.batches = [list(b) for b in batches] self.shuffle = shuffle_batch_order self.drop_last = drop_last self.seed = int(seed) self.epoch = 0 if dist.is_available() and dist.is_initialized(): self.world_size = dist.get_world_size() self.rank = dist.get_rank() else: self.world_size = 1 self.rank = 0 def __iter__(self): n_batches = len(self.batches) order = list(range(n_batches)) if self.shuffle: g = torch.Generator() g.manual_seed(self.seed + self.epoch) order = torch.randperm(n_batches, generator=g).tolist() # make divisible across ranks if self.drop_last: total = (len(order) // self.world_size) * self.world_size order = order[:total] else: pad = (-len(order)) % self.world_size if pad: order = order + order[:pad] # shard by rank for i in order[self.rank::self.world_size]: yield self.batches[i] def __len__(self): n = len(self.batches) if self.drop_last: return (n // self.world_size) return math.ceil(n / self.world_size) # Lightning will call this if present via its epoch hooks def set_epoch(self, epoch: int): self.epoch = int(epoch) class PreBatchedSampler(Sampler[List[int]]): """ Yields precomputed batches of indices, e.g. [[3,7,9], [0,1,2], ...]. Useful when you've already formed batches by length. """ def __init__( self, batches: Sequence[Sequence[int]], shuffle_batch_order: bool = False, generator=None, ): self.batches = [list(b) for b in batches] self.shuffle_batch_order = shuffle_batch_order self.generator = generator def __iter__(self) -> Iterable[List[int]]: if self.shuffle_batch_order: # local copy we can shuffle without touching the original idxs = list(range(len(self.batches))) g = self.generator if self.generator is not None else torch.Generator() perm = torch.randperm(len(idxs), generator=g).tolist() for i in perm: yield self.batches[i] else: for b in self.batches: yield b def __len__(self) -> int: return len(self.batches) def compute_tr_lengths_from_shelf( tr_shelf_path: str, tr_sequences: list[str] ) -> list[int]: """ Opens the TR shelf once and returns length for each sequence. 2D array -> length = shape[0]; 1D array (pooled) -> length = 1. """ lengths = [] with shelve.open(tr_shelf_path, flag="r") as db: for s in tr_sequences: arr = np.asarray(db[str(s)]) if arr.ndim == 1: lengths.append(1) else: lengths.append(int(arr.shape[0])) return lengths def make_length_batches( dataset_records: list[dict], tr_shelf_path: str, batch_size: int, drop_last: bool = False, ) -> list[list[int]]: """ dataset_records: output of PairDataset._load_and_normalize(...), i.e. list of dicts with keys: "dna_sequence", "tr_sequence", "scores", ... Returns a list of batches, each a list of indices, sorted by (dna_len, tr_len). """ # DNA length comes from label length dna_lens = [len(r["scores"]) for r in dataset_records] tr_seqs = [r["tr_sequence"] for r in dataset_records] # TR length requires a quick shelf lookup (done once here) tr_lens = compute_tr_lengths_from_shelf(tr_shelf_path, tr_seqs) # sort indices by (dna_len, tr_len) idxs = list(range(len(dataset_records))) idxs.sort(key=lambda i: (dna_lens[i], tr_lens[i])) # chunk into fixed-size batches batches = [idxs[i : i + batch_size] for i in range(0, len(idxs), batch_size)] if drop_last and len(batches) and len(batches[-1]) < batch_size: batches.pop() return batches # ---- dataset --------------------------------------------------------- class PairDataset(Dataset): def __init__( self, dataset: pd.DataFrame, norm_value: int = 1333, round_to: int = 4, score_col="scores", target_col="dna_sequence", binder_col="tr_sequence" ): """ Args: - dataset: a dataset with the needed information: ID, dna_sequence, tr_sequence, scores - norm_value: max score, which we'll use to divide all the integer scores in "scores" - round_to: how many decimal places for the numerical score values """ self.fake_scores=False self.score_col = score_col self.target_col = target_col self.binder_col = binder_col self.norm_value = norm_value self.round_to = round_to self.dataset = self._load_and_normalize(dataset) def _load_and_normalize(self, dataset): """ Labels come in looking like "0,0,0,100,100,133,133,100,100,0,0," This method turns the labels from strings into floats out to 4 decimal places """ if self.score_col not in dataset.columns: logger.info(f"Scores not provided. Adding placeholder scores where all positions are considered binding") dataset[self.score_col] = dataset[self.target_col].str.len() dataset[self.score_col] = dataset[self.score_col].apply(lambda x: ",".join([str(self.norm_value)]*x)) self.fake_scores=True # split string into list of strings dataset[self.score_col] = dataset[self.score_col].apply(lambda x: x.split(",")) dataset["copycol"] = dataset[self.score_col] # turn list of strings into list of normalized, rounded floats dataset[self.score_col] = dataset[self.score_col].apply( lambda x: [round(int(y) / self.norm_value, self.round_to) for y in x] ) # convert to records for ease of loading dataset = dataset.to_dict(orient="records") return dataset def __len__(self): return len(self.dataset) def __getitem__(self, idx): item = self.dataset[idx] return {**(item if isinstance(item, dict) else {})} class PairDataModule(LightningDataModule): def __init__( self, train_file: Path | str = "../data_files/splits/train.csv", val_file: Path | str = "../data_files/splits/val.csv", test_file: Path | str = "../data_files/splits/test.csv", tr_shelf_path: ( Path | str ) = "../data_files/processed/embeddings/fimo_hits_only/trs_esm.shelf", dna_shelf_path: ( Path | str ) = "../data_files/processed/embeddings/fimo_hits_only/baby_peaks_segmentnt_pernuc_with_onehot.shelf", batch_size: int = 1, num_workers=8, maximize_num_workers=False, debug_run: bool = False, pin_memory: bool = False, shuffle_train_batch_order: bool = True, score_col: str = "scores", target_col: str = "dna_sequence", binder_col: str = "tr_sequence", norm_value: int = 1333 ): super().__init__() self.save_hyperparameters() self.debug_run = debug_run self.norm_value = norm_value # Initialize the data files self.train_data_file = train_file self.val_data_file = val_file self.test_data_file = test_file self.target_col = target_col self.binder_col = binder_col self.score_col = score_col # Initialize hyperparameters like batch size self.batch_size = batch_size self.num_workers = ( cpu_count() if maximize_num_workers else min(num_workers, cpu_count()) ) # Set up ShelfCollator self.collate = ShelfCollator( tr_shelf_path=str(tr_shelf_path), dna_shelf_path=str(dna_shelf_path), tr_key=self.binder_col, dna_key=self.target_col, dtype=torch.float32, pad_value=-1.0, debug_run =self.debug_run, score_col = self.score_col ) self.drop_last = False # or True, your choice self.shuffle_batch_order = shuffle_train_batch_order # False keep batches deterministic per epoch; set True if you want to shuffle batch order logger.info(f"num_workers={self.num_workers}") logger.info("Initialized BinderDecoyDataModule constants") def load_file(self, file_path, lim=None): """ Load and unpack an input csv whose columns are binder_path,glm_path,label """ try: df = pd.read_csv(file_path) if lim is not None: df = df[:lim].reset_index(drop=True) return df except: raise Exception(f"{file_path} is not a valid file") def setup(self, stage: str | None = None): lim = 5 if self.debug_run else None # FIT: build train & val (so val exists during training) if stage in (None, "fit"): if not hasattr(self, "train_dataset"): train_df = self.load_file(self.train_data_file, lim=lim) self.train_dataset = PairDataset(train_df, norm_value = self.norm_value, score_col = self.score_col, target_col = self.target_col, binder_col = self.binder_col) self.train_batches = make_length_batches( dataset_records=self.train_dataset.dataset, tr_shelf_path=str(self.hparams.tr_shelf_path), batch_size=self.batch_size, drop_last=self.drop_last, ) self.train_batch_sampler = PreBatchedDistributedBatchSampler( self.train_batches, shuffle_batch_order=self.shuffle_batch_order, drop_last=self.drop_last, seed=0, ) if not hasattr(self, "val_dataset"): val_df = self.load_file(self.val_data_file, lim=lim) self.val_dataset = PairDataset(val_df, norm_value = self.norm_value, score_col = self.score_col, target_col = self.target_col, binder_col = self.binder_col) self.val_batches = make_length_batches( dataset_records=self.val_dataset.dataset, tr_shelf_path=str(self.hparams.tr_shelf_path), batch_size=self.batch_size, drop_last=False, ) self.val_batch_sampler = PreBatchedDistributedBatchSampler( self.val_batches, shuffle_batch_order=False, drop_last=False, seed=0 ) # VALIDATE called standalone: ensure val is built if stage in (None, "validate"): if not hasattr(self, "val_dataset"): val_df = self.load_file(self.val_data_file, lim=lim) self.val_dataset = PairDataset(val_df, norm_value = self.norm_value, score_col = self.score_col, target_col = self.target_col, binder_col = self.binder_col) self.val_batches = make_length_batches( dataset_records=self.val_dataset.dataset, tr_shelf_path=str(self.hparams.tr_shelf_path), batch_size=self.batch_size, drop_last=False, ) self.val_batch_sampler = PreBatchedDistributedBatchSampler( self.val_batches, shuffle_batch_order=False, drop_last=False, seed=0 ) # TEST phase if stage in (None, "test"): if not hasattr(self, "test_dataset"): test_df = self.load_file(self.test_data_file, lim=lim) self.test_dataset = PairDataset(test_df, norm_value = self.norm_value, score_col = self.score_col, target_col = self.target_col, binder_col = self.binder_col) self.test_batches = make_length_batches( dataset_records=self.test_dataset.dataset, tr_shelf_path=str(self.hparams.tr_shelf_path), batch_size=self.batch_size, drop_last=False, ) self.test_batch_sampler = PreBatchedSampler( self.test_batches, shuffle_batch_order=False ) def train_dataloader(self): return DataLoader( self.train_dataset, batch_sampler=self.train_batch_sampler, collate_fn=self.collate, num_workers=self.num_workers, persistent_workers=(self.num_workers > 0), pin_memory=self.hparams.pin_memory, ) def val_dataloader(self): return DataLoader( self.val_dataset, batch_sampler=self.val_batch_sampler, collate_fn=self.collate, num_workers=self.num_workers, persistent_workers=(self.num_workers > 0), pin_memory=self.hparams.pin_memory, ) def test_dataloader(self): return DataLoader( self.test_dataset, batch_sampler=self.test_batch_sampler, collate_fn=self.collate, num_workers=self.num_workers, persistent_workers=(self.num_workers > 0), pin_memory=self.hparams.pin_memory, ) def predict_dataloader(self): # Same as test return DataLoader( self.test_dataset, batch_sampler=self.test_batch_sampler, collate_fn=self.collate, num_workers=self.num_workers, persistent_workers=(self.num_workers > 0), pin_memory=self.hparams.pin_memory, ) class ShelfCollator: """ Lazily opens TR (binder) and DNA shelves the first time each worker calls __call__. Expects each item to contain keys: - "tr_sequence": str (key for TR shelf) - "dna_sequence": str (key for DNA shelf) - "scores": list[float] (per-base labels for DNA) - optional "ID" Returns a dict with: - binder_emb: FloatTensor [B, Lb_max, Db] (padded) - binder_mask: BoolTensor [B, Lb_max] - glm_emb: FloatTensor [B, Lg_max, Dg] (padded) - glm_mask: BoolTensor [B, Lg_max] - labels: FloatTensor [B, Lg_max] (padded, zeros where masked) - ids, tr_sequences, dna_sequences: lists """ def __init__( self, tr_shelf_path: str, dna_shelf_path: str, tr_key: str = "tr_sequence", dna_key: str = "dna_sequence", dtype: torch.dtype = torch.float32, pad_value: float = -1.0, debug_run: bool = False, score_col = "scores" ): self.tr_path = tr_shelf_path self.dna_path = dna_shelf_path self.score_col = score_col self.tr_key = tr_key self.dna_key = dna_key self.dtype = dtype self.pad_value = pad_value self.debug_run = debug_run # opened lazily per worker: self._tr_db = None self._dna_db = None def _ensure_open(self): if self._tr_db is None: self._tr_db = shelve.open(self.tr_path, flag="r") # read-only if self._dna_db is None: self._dna_db = shelve.open(self.dna_path, flag="r") def __call__(self, batch): """ batch: list[dict] from Dataset.__getitem__ """ self._ensure_open() ids = [b.get("ID", None) for b in batch] tr_seqs = [b[self.tr_key] for b in batch] dna_seqs = [b[self.dna_key] for b in batch] scores_list = [b[self.score_col] for b in batch] # 1) Fetch embeddings lazily from shelves binder_list = [] glm_list = [] binder_lens = [] glm_lens = [] for tr, dna, scores in zip(tr_seqs, dna_seqs, scores_list): # ----- binder/TR ----- tr_arr = np.asarray(self._tr_db[str(tr)]) # ensure 2D: [Lb, Db] (if pooled 1D, make length=1) if tr_arr.ndim == 1: tr_arr = tr_arr[None, :] binder_list.append(torch.from_numpy(tr_arr).to(self.dtype)) binder_lens.append(tr_arr.shape[0]) # ----- DNA / GLM ----- dna_arr = np.asarray(self._dna_db[str(dna)]) if dna_arr.ndim == 1: dna_arr = dna_arr[None, :] glm_list.append(torch.from_numpy(dna_arr).to(self.dtype)) glm_lens.append(dna_arr.shape[0]) # sanity: scores length should match dna length if len(scores) != dna_arr.shape[0]: raise ValueError( f"Length mismatch for DNA seq: shelf length={dna_arr.shape[0]} " f"but scores length={len(scores)}" ) # 2) Pad sequences to batch max length binder_emb = pad_sequence( binder_list, batch_first=True, padding_value=self.pad_value ) # [B, Lb_max, Db] glm_emb = pad_sequence( glm_list, batch_first=True, padding_value=self.pad_value ) # [B, Lg_max, Dg] binder_lens = torch.as_tensor(binder_lens, dtype=torch.int64) glm_lens = torch.as_tensor(glm_lens, dtype=torch.int64) binder_mask = torch.arange(binder_emb.size(1)).unsqueeze( 0 ) < binder_lens.unsqueeze( 1 ) # [B, Lb_max] glm_mask = torch.arange(glm_emb.size(1)).unsqueeze(0) < glm_lens.unsqueeze( 1 ) # [B, Lg_max] # True = PAD (what MHA expects) binder_kpm = ~binder_mask glm_kpm = ~glm_mask # 3) Collate labels for DNA and pad labels_list = [torch.tensor(s, dtype=torch.float32) for s in scores_list] labels = pad_sequence( labels_list, batch_first=True, padding_value=self.pad_value ) # [B, Lg_max] if self.debug_run: max_binder_len = max(binder_lens) max_glm_len = max(glm_lens) binder_expected_false = sum(max_binder_len-binder_lens).item() binder_expected_true = sum(binder_lens) binder_expected_total = binder_expected_true + binder_expected_false glm_expected_false = sum(max_glm_len-glm_lens).item() glm_expected_true = sum(glm_lens).item() glm_expected_total = glm_expected_true + glm_expected_false labels_neg1 = sum(sum(labels==-1)).item() expected_labels_neg1 = glm_expected_false logger.info(f" Max binder length: {max_binder_len}, original lengths: {binder_lens}, ultimate dimensions: {binder_emb.shape}") logger.info(f" Binder expect: true/total = {binder_expected_true}/{binder_expected_total}") logger.info(f" Max DNA length: {max_glm_len}, original lengths: {glm_lens}, ultimate dimensions: {glm_emb.shape}") logger.info(f" DNA expect: true/total = {glm_expected_true}/{glm_expected_total}") logger.info(f" Labels expect -1: -1/total = {expected_labels_neg1}/{glm_expected_total}. True: {labels_neg1}/{labels.numel()}") return { "binder_emb": binder_emb, # [B, Lb_max, Db] "binder_mask": binder_mask, # [B, Lb_max] "binder_kpm": binder_kpm.bool(), # True = PAD ← pass to MHA "glm_emb": glm_emb, # [B, Lg_max, Dg] "glm_mask": glm_mask, # [B, Lg_max] "glm_kpm": glm_kpm.bool(), # True = PAD ← pass to MHA "labels": labels, # [B, Lg_max] "ID": ids, "tr_sequence": tr_seqs, "dna_sequence": dna_seqs, } # ------------------------ Helpers for main method debugging only ------------------------------------------# def _peek_batches(dl, n_batches: int = 2, tag: str = "train"): logger.info(f"\n=== Peek {n_batches} batch(es) from {tag} loader ===") for i, batch in enumerate(dl): be = batch["binder_emb"] bm = batch["binder_mask"] ge = batch["glm_emb"] gm = batch["glm_mask"] y = batch["labels"] ids = batch.get("ID", [""] * be.size(0)) logger.info(f"\n[{tag}] batch {i+1}") logger.info(f" binder_emb: {tuple(be.shape)} dtype={be.dtype}") logger.info(f" binder_emb: {tuple(bm.shape)} dtype={bm.dtype}") logger.info(f" binder_mask true count: {bm.sum().item()} / {bm.numel()}") logger.info(f" glm_emb: {tuple(ge.shape)} dtype={ge.dtype}") logger.info(f" glm_mask true count: {gm.sum().item()} / {gm.numel()}") logger.info(f" glm_mask: {tuple(gm.shape)} dtype={gm.dtype}") logger.info( f" labels: {tuple(y.shape)} min={y.min().item():.4f} max={y.max().item():.4f}, total -1 = {sum(sum(y==-1)).item()}" ) logger.info(f" IDs (first 5): {ids[:5]}") # should make sure that the number of labels that are -1 equals the number of padding tokens if i + 1 >= n_batches: break def _warn_on_paths(args): import os for p, label in [ (args.train_file, "train_file"), (args.val_file, "val_file"), (args.test_file, "test_file"), (args.tr_shelf_path, "tr_shelf_path"), (args.dna_shelf_path, "dna_shelf_path"), ]: if p and not os.path.exists(p): logger.info(f"{label} does not exist: {p}") if str(args.tr_shelf_path).endswith(".pkl"): logger.info( "Warning: tr_shelf_path ends with .pkl but ShelfCollator expects a shelve DB " "(e.g., `.shelf`). Pass the correct path via --tr_shelf_path." ) def main(): parser = argparse.ArgumentParser( description="Peek pre-batched, shelf-backed dataloaders" ) parser.add_argument( "--train_file", type=str, default="../data_files/processed/splits/by_dna/babytrain.csv", ) parser.add_argument( "--val_file", type=str, default="../data_files/processed/splits/by_dna/babyval.csv", ) parser.add_argument( "--test_file", type=str, default="../data_files/processed/splits/by_dna/babytest.csv", ) parser.add_argument( "--tr_shelf_path", type=str, default="../data_files/processed/embeddings/fimo_hits_only/trs_esm.shelf", ) parser.add_argument( "--dna_shelf_path", type=str, default="../data_files/processed/embeddings/fimo_hits_only/peaks_caduceus.shelf", ) parser.add_argument("--batch_size", type=int, default=4) parser.add_argument("--num_workers", type=int, default=4) parser.add_argument( "--debug_run", default=True, action="store_true", help="limit dataset to a few rows" ) parser.add_argument( "--n_batches", type=int, default=2, help="how many batches to print per split" ) parser.add_argument("--shuffle_train_batch_order", action="store_true") args = parser.parse_args() _warn_on_paths(args) dm = PairDataModule( train_file=args.train_file, val_file=args.val_file, test_file=args.test_file, tr_shelf_path=args.tr_shelf_path, dna_shelf_path=args.dna_shelf_path, batch_size=args.batch_size, num_workers=args.num_workers, debug_run=args.debug_run, shuffle_train_batch_order=args.shuffle_train_batch_order, pin_memory=False, score_col="binary_scores", norm_value=1 ) # ---- Train ---- dm.setup(stage="fit") train_dl = dm.train_dataloader() _peek_batches(train_dl, n_batches=args.n_batches, tag="fit") # ---- Val ---- dm.setup(stage="validate") val_dl = dm.val_dataloader() _peek_batches(val_dl, n_batches=1, tag="val") # usually enough to sanity-check # ---- Test ---- dm.setup(stage="test") test_dl = dm.test_dataloader() _peek_batches(test_dl, n_batches=1, tag="test") logger.info("\nAll good") if __name__ == "__main__": # (Optional) set a deterministic seed for batch order shuffling torch.manual_seed(0) random.seed(0) np.random.seed(0) logging.basicConfig( level=logging.INFO, format="%(asctime)s | %(levelname)s | %(name)s:%(lineno)d | %(message)s", datefmt="%H:%M:%S", handlers=[logging.StreamHandler(sys.stdout)], # stdout, not stderr force=True, # override any prior config from imported libs ) main()