Merge branch 'hillst/block-sampling3' of /home/skothenhill/bionemo-framework-fresh/bionemo-framework/. into hillst/block-sampling3

Signed-off-by: Steven <skothenhill@nvidia.com>

Author: skothenhill-nv

sub-packages/bionemo-core/src/bionemo/core/data/multi_epoch_dataset.py

@@ -17,7 +17,7 @@
 import math
 from abc import ABC, abstractmethod
 from dataclasses import dataclass
-from typing import Generic, NamedTuple, Protocol, Sequence, TypeVar
+from typing import Any, Generic, NamedTuple, Protocol, Sequence, TypeVar
 
 import numpy as np
 from torch.utils.data import Dataset
@@ -130,10 +130,17 @@ def __post_init__(self):
 
     def __getitem__(self, index: int) -> T_co:
         """Get the sample at the given index."""
-        if index not in range(len(self)):
+        if index < 0 or index >= len(self):
             raise IndexError(f"Index {index} out of bounds for dataset of length {len(self)}.")
         return self.dataset[self._global_index_to_permuted_local_index(index)]
 
+    def __getitems__(self, indices: list[int]) -> Any:
+        """Get the samples at the given indices."""
+        if hasattr(self.dataset, '__getitems__'):
+            return self.dataset.__getitems__([self[i] for i in indices])
+        else:
+            return [self[i] for i in indices]
+
     def __len__(self) -> int:
         """Return the length of the resampled dataset."""
         return self.num_samples  # type: ignore

sub-packages/bionemo-geneformer/src/bionemo/geneformer/data/block_sampling.py

@@ -145,8 +145,11 @@ def __getitems__(self, indices: List[int]) -> Any:
             _sorted_order = np.argsort(shuffled_ids)
             _sorted_idxs = np.sort(shuffled_ids)
 
-            # Sort for I/O locality as we use blocked fetches.
-            sorted_data = self.dataset[_sorted_idxs]
+            # Turn it back into a list so torch does the right things.
+            if hasattr(self.dataset, '__getitems__'):
+                sorted_data = self.dataset.__getitems__(_sorted_idxs.tolist())
+            else:
+                sorted_data = [self.dataset[idx] for idx in _sorted_idxs.tolist()]
 
             # Reverse the sorting to return the args in the original state.
             data = np.array(sorted_data)[np.argsort(_sorted_order)]
@@ -361,7 +364,6 @@ def __iter__(self):
                     # Other workers get the base number of fetches
                     start = worker_info.id * per_worker + remainder
                     end = start + per_worker
-
                 fetches = fetches[start:end]
 
             if self.sort_before_fetch:
@@ -374,15 +376,14 @@ def __iter__(self):
                 if self.fetch_callback is not None:
                     data = self.fetch_callback(self.collection, fetch_ids)
                 else:
-                    data = self.collection[fetch_ids]
+                    data = list(self.collection[i] for i in fetch_ids)
 
                 if not isinstance(data, np.ndarray):
                     data = np.array(data)
 
                 # Call fetch transform if provided
                 if self.fetch_transform is not None:
                     data = self.fetch_transform(data)
-
                 if self.shuffle_before_yield:
                     # Shuffle the indices
                     if bionemo_permute:
@@ -408,7 +409,6 @@ def __iter__(self):
                     # Call batch transform if provided
                     if self.batch_transform is not None:
                         batch_data = self.batch_transform(batch_data)
-
                     yield batch_data
 
         else:  # Not shuffling indices before fetching

sub-packages/bionemo-geneformer/src/bionemo/geneformer/data/singlecell/datamodule.py

@@ -19,6 +19,7 @@
 from typing import List, Literal, Optional, Sequence
 
 import numpy as np
+from bionemo.geneformer.data.block_sampling import MapStyleScDataset
 from nemo.lightning.data import WrappedDataLoader
 from nemo.lightning.pytorch.plugins import MegatronDataSampler
 from nemo.utils import logging
@@ -84,6 +85,8 @@ def __init__(  # noqa: D107
         persistent_workers: bool = True,
         pin_memory: bool = True,
         include_unrecognized_vocab_in_dataset: bool = False,
+        block_size: Optional[int] = None,
+        fetch_factor: Optional[int] = None,
     ) -> None:
         super().__init__()
         if predict_dataset_path is None:
@@ -111,10 +114,16 @@ def __init__(  # noqa: D107
         self.num_workers = num_workers
         self.persistent_workers = persistent_workers
         self.pin_memory = pin_memory
+        self.global_batch_size = global_batch_size
+        # Block sampling parameters
+        self.block_size = block_size
+        self.fetch_factor = fetch_factor
+        self.block_sampling = block_size and fetch_factor
 
         rng = np.random.default_rng(seed)
         if self.data_path_train is not None:
             assert self.data_path_val is not None and self.data_path_test is not None
+
             self._train_dataset_ori = SingleCellDataset(
                 self.data_path_train,
                 self.tokenizer,
@@ -201,12 +210,31 @@ def setup(self, stage: str = "") -> None:  # noqa: D102
             num_train_samples = int(max_train_steps * self.data_sampler.global_batch_size)
 
             # This happens exactly once during setup.
-            self._train_ds = MultiEpochDatasetResampler(
-                self._train_dataset_ori,
-                num_samples=num_train_samples,
-                shuffle=True,
-                seed=self.seed,
-            )
+            if self.block_sampling:
+                # We also need associated block sampling parameters.
+
+                # dataset size must divide block size * batch_size
+                if num_train_samples % (self.global_batch_size * self.block_size) != 0:
+                    # Warning
+                    num_train_samples -=  num_train_samples % (self.global_batch_size * self.block_size)
+
+
+                from bionemo.geneformer.data.block_sampling import MapStyleScDataset
+                from bionemo.core.data.multi_epoch_dataset import MultiEpochDatasetResampler
+
+                self._train_ds = MultiEpochDatasetResampler(
+                    self._train_dataset_ori,
+                    num_samples=num_train_samples,
+                    shuffle=False,
+                    seed=self.seed,
+                )
+                self._train_ds = MapStyleScDataset(
+                    self._train_ds,
+                    block_size=self.block_size,
+                    batch_size=self.block_sampling,
+                    fetch_factor=self.fetch_factor,
+                    seed=self.seed * 2,
+                )
             if self.trainer.limit_val_batches == 0:  # disable validation
                 logging.info("Skip creating validation dataset because trainer.limit_val_batches=0.")
             else:

sub-packages/bionemo-geneformer/src/bionemo/geneformer/data/singlecell/dataset.py

@@ -119,6 +119,8 @@ def __len__(self):  # noqa: D105
 
     def __getitem__(self, index: EpochIndex) -> types.BertSample:
         """Performs a lookup and the required transformation for the model."""
+        if not isinstance(index, EpochIndex):
+            index = EpochIndex(idx=index, epoch=0)
         rng = np.random.default_rng([self._seed, index.epoch, index.idx])
         values, feature_ids = self.scdl.get_row(index.idx, return_features=True, feature_vars=["feature_id"])
         assert (
@@ -145,7 +147,6 @@ def __getitem__(self, index: EpochIndex) -> types.BertSample:
             include_unrecognized_vocab_in_dataset=self.include_unrecognized_vocab_in_dataset,
         )
 
-
 def _gather_medians(
     gene_names: np.ndarray,
     gene_data: np.ndarray,
@@ -155,16 +156,12 @@ def _gather_medians(
     include_unrecognized_vocab_in_dataset: bool = False,
 ) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
     """Filter out genes that are not in the provided tokenizer vocab, and tokenize the gene names."""
-    genes, tokens, medians = [], [], []
-    for tok, gene in zip(gene_names, gene_data):
-        if tok in vocab:
-            tokens.append(vocab[tok])
-            genes.append(gene)
-            if normalize:
-                med = gene_median[tok]  # If not in the dictionary we default to no normalization (1)
-                medians.append(med)
-        elif include_unrecognized_vocab_in_dataset:
-            raise ValueError(f"Provided gene identifier, {str(tok)}, is not in the tokenizer vocab.")
+    tok_genes = filter(
+        lambda x: x[0] is not None, 
+        ((vocab.get(tok), gene, gene_median.get(tok, 1.0)) for tok, gene in zip(gene_names, gene_data))
+    )
+
+    tokens, genes, medians = zip(*tok_genes)
     return np.asarray(genes), np.asarray(tokens), np.asarray(medians)
 
 

sub-packages/bionemo-geneformer/src/bionemo/geneformer/scripts/scdataset/scdataset_perf_tests.py

@@ -0,0 +1,128 @@
+from pathlib import Path
+from bionemo.core.data.load import load
+from bionemo.geneformer.data.block_sampling import MapStyleScDataset, scDataset
+from bionemo.geneformer.data.singlecell.dataset import SingleCellDataset
+from torch.utils.data import DataLoader
+from bionemo.core.data.multi_epoch_dataset import MultiEpochDatasetResampler
+import time
+import tqdm
+import functools
+from bionemo.llm.data import collate
+
+from bionemo.geneformer.data.singlecell.preprocess import GeneformerPreprocess
+from bionemo.geneformer.tokenizer.gene_tokenizer import GeneTokenizer
+def make_dataset():
+    data_path: Path = load("single_cell/testdata-20241203") / "cellxgene_2023-12-15_small_processed_scdl" / "train"
+
+    train_data_path = Path("/home/ubuntu/data/cellxgene_2023-12-15/train")
+
+    preprocessor = GeneformerPreprocess(
+        download_directory=train_data_path,
+        medians_file_path=train_data_path / "medians.json",
+        tokenizer_vocab_path=train_data_path / "geneformer.vocab",
+    )
+    match preprocessor.preprocess():
+        case {"tokenizer": tokenizer, "median_dict": median_dict}:
+            tokenizer, median_dict = tokenizer, median_dict
+        case _:
+                raise ValueError("Preprocessing must have failed.")
+
+    dataset = SingleCellDataset(train_data_path, tokenizer=tokenizer, median_dict=median_dict, max_len=2048)
+    print("done loading ds")
+    return dataset
+
+def get_configs():
+    return [
+        {
+            "block_size": 64,
+            "batch_size": 128 * 8,
+            "fetch_factor": 8,
+            "seed": 42
+        }
+    ]
+
+def mapstyle_throughput():
+    dataset = make_dataset()
+    tokenizer = dataset.tokenizer
+
+    configs = get_configs()
+    for config in configs:
+        factor = config["fetch_factor"] * config["batch_size"]
+        extra = len(dataset) % factor
+        to_add = factor - extra 
+        num_samples = (len(dataset) + to_add) 
+
+        dataset = MultiEpochDatasetResampler(
+            dataset, 
+            num_samples=num_samples,
+            shuffle=False,
+        )
+        '''
+        When we stack the datasets this way, a whole vector is passed into getitem for 
+        MultiEpochDatasetResampler
+        '''
+        dataset = MapStyleScDataset(dataset, **config)
+
+        start = time.time()
+        dataloader = DataLoader(dataset, batch_size=config["batch_size"], num_workers=16,
+            collate_fn=functools.partial(
+                collate.bert_padding_collate_fn,
+                padding_value=tokenizer.token_to_id(GeneTokenizer.pad_token),
+                min_length=2048,
+                max_length=2048,
+            ),
+        
+        )
+
+
+        for i, batch in enumerate(tqdm.tqdm(dataloader)):
+            if i > 100 * config["fetch_factor"]:
+                break
+            pass
+
+        end = time.time()
+        print(f"MapStyleScDataset: {end - start} seconds")
+        print(f"MapStyleScDataset: { 800 * config['batch_size'] / ( end - start)} samples per second")
+
+def iterstyle_throughput():
+    dataset = make_dataset()
+    tokenizer = dataset.tokenizer
+
+    configs = get_configs()
+
+    for config in configs:
+        num_samples = (len(dataset)  - (len(dataset) % (config["batch_size"] * config["block_size"]))) * 2
+        dataset = MultiEpochDatasetResampler(
+            dataset, 
+            num_samples=num_samples,
+            shuffle=False,
+        )
+        # TODO get some intermediate metrics
+        dataset = scDataset(dataset, bionemo_permute=False, **config)
+
+        start = time.time()
+        dataloader = DataLoader(dataset, batch_size=None, num_workers=16, shuffle=False, 
+            collate_fn=functools.partial(
+                collate.bert_padding_collate_fn,
+                padding_value=tokenizer.token_to_id(GeneTokenizer.pad_token),
+                min_length=2048,
+                max_length=2048,
+            ),
+        )
+
+        # I think this just happens if its not an even multiple of the batch size
+        try:
+            for i, batch in enumerate(tqdm.tqdm(dataloader)):
+                if i > 100 * config["fetch_factor"]:
+                    break
+                pass
+        except RuntimeError as e:
+            print(e)
+
+        end = time.time()
+        print(f"IterStyleDataset: {end - start} seconds")
+        print(f"IterStyleDataset: { 800 * config['batch_size']/ ( end - start)} samples per second")
+
+if __name__ == "__main__":
+    mapstyle_throughput() 
+    iterstyle_throughput()