Add BERT-style masking function

.vscode/settings.json

@@ -1,8 +1,14 @@
 {
     "cSpell.words": [
-        // List of words to be added to the spell-checker dictionary.
-        // In vscode, use the "Add <word> to workspace settings" in the quick-fix menu to add words to this list.
-        "bionemo"
+        "allclose",
+        "bionemo",
+        "dtype",
+        "nemo",
+        "pretraining",
+        "rampup",
+        "resamplers",
+        "singlecell",
+        "uniref"
     ],
     "editor.rulers": [
         120

scripts/singlecell/geneformer/pretrain.py

@@ -169,7 +169,7 @@ def main(
         train_dataset_path=train_data_path,
         val_dataset_path=val_data_path,
         test_dataset_path=test_data_path,
-        random_token_prob=0.1,  # this is the incorrect setting we originally used.
+        random_token_prob=0.02,  # changed to represent the incorrect setting we originally used.
         median_dict=median_dict,
         micro_batch_size=micro_batch_size,
         global_batch_size=micro_batch_size * int(num_nodes * devices / pipeline_model_parallel_size),

sub-packages/bionemo-core/src/bionemo/core/utils/random_utils.py

@@ -44,3 +44,11 @@ def random_numpy_context(seed: int = 42) -> Iterator[None]:
         yield
     finally:
         np.random.set_state(state)
+
+
+def get_seed_from_rng(rng: np.random.Generator) -> int:
+    """Generates a deterministic random seed from an existing random generator.
+
+    Used to seed a torch random generator from a numpy random generator.
+    """
+    return rng.integers(np.iinfo(np.int64).max)

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

@@ -14,9 +14,11 @@
 # limitations under the License.
 
 
+import functools
 from pathlib import Path
 from typing import List, Optional, Sequence
 
+import numpy as np
 import pytorch_lightning as pl
 from nemo.lightning.pytorch.plugins import MegatronDataSampler
 from nemo.utils import logging
@@ -25,7 +27,10 @@
 from torch.utils.data import DataLoader
 
 from bionemo.core.data.resamplers import PRNGDatasetShuffler
+from bionemo.core.utils import random_utils
 from bionemo.geneformer.data.singlecell.dataset import SingleCellDataset
+from bionemo.geneformer.tokenizer.gene_tokenizer import GeneTokenizer
+from bionemo.llm.data import collate
 
 
 __all__: Sequence[str] = ("SingleCellDataModule",)
@@ -90,6 +95,8 @@ def __init__(  # noqa: D107
         self.persistent_workers = persistent_workers
         self.pin_memory = pin_memory
         self.index_mapping_dir = index_mapping_dir or str(Path(self.data_path_train).parent)
+
+        rng = np.random.default_rng(seed)
         self._train_dataset_ori = SingleCellDataset(
             self.data_path_train,
             self.tokenizer,
@@ -98,6 +105,7 @@ def __init__(  # noqa: D107
             mask_prob=self.mask_prob,
             mask_token_prob=self.mask_token_prob,
             random_token_prob=self.random_token_prob,
+            seed=random_utils.get_seed_from_rng(rng),
         )
         self._val_dataset_ori = SingleCellDataset(
             self.data_path_val,
@@ -107,6 +115,7 @@ def __init__(  # noqa: D107
             mask_prob=self.mask_prob,
             mask_token_prob=self.mask_token_prob,
             random_token_prob=self.random_token_prob,
+            seed=random_utils.get_seed_from_rng(rng),
         )
         self._test_dataset_ori = SingleCellDataset(
             self.data_path_test,
@@ -116,6 +125,7 @@ def __init__(  # noqa: D107
             mask_prob=self.mask_prob,
             mask_token_prob=self.mask_token_prob,
             random_token_prob=self.random_token_prob,
+            seed=random_utils.get_seed_from_rng(rng),
         )
 
         # This is needed here, or you need to specify it in the megatron adapter thing TODO name?
@@ -169,7 +179,12 @@ def _create_dataloader(self, dataset, **kwargs) -> DataLoader:
             num_workers=self.num_workers,
             pin_memory=self.pin_memory,
             persistent_workers=self.persistent_workers,
-            # collate_fn=dataset.collate_fn,  No special work happens in this dataloader outside of getitem
+            collate_fn=functools.partial(
+                collate.bert_padding_collate_fn,
+                padding_value=self.tokenizer.token_to_id(GeneTokenizer.pad_token),
+                min_length=None,
+                max_length=self.max_len,
+            ),
             **kwargs,
         )
 

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

@@ -16,32 +16,25 @@
 
 import json
 from pathlib import Path
-from typing import Any, Dict, Optional, Sequence, Tuple, TypedDict
+from typing import Any, Dict, Optional, Sequence, Tuple
 
 import numpy as np
+import torch
 from nemo.utils import logging
 from torch.utils.data import Dataset
 
-from bionemo.geneformer.data.singlecell.utils import sample_or_truncate_plus_pad
+from bionemo.core.utils import random_utils
+from bionemo.geneformer.data.singlecell.utils import sample_or_truncate
 from bionemo.geneformer.tokenizer.gene_tokenizer import GeneTokenizer
+from bionemo.llm.data import masking, types
 
 
 __all__: Sequence[str] = (
     "SingleCellDataset",
-    "Item",
     "process_item",
 )
 
 
-class Item(TypedDict):  # noqa: D101
-    text: np.ndarray
-    types: np.ndarray
-    padding_mask: np.ndarray
-    labels: np.ndarray
-    loss_mask: np.ndarray
-    is_random: np.ndarray
-
-
 class SingleCellDataset(Dataset):
     """A dataset class for single-cell pre-training. These can be generated using the sc_memmap.py script. Future
     updates will contain more comprehensive workflows for generating a Sparse Memmap from scRNA-seq.
@@ -94,6 +87,7 @@ def __init__(  # noqa: D107
         random_token_prob: float = 0.1,
         prepend_cls_token: bool = True,
         assert_increasing_columns: bool = True,
+        seed: int = np.random.SeedSequence().entropy,  # type: ignore
     ):
         super().__init__()
         self.data_path = data_path
@@ -102,6 +96,7 @@ def __init__(  # noqa: D107
         self.mask_token_prob = mask_token_prob
         self.mask_prob = mask_prob
         self.prepend_cls_token = prepend_cls_token
+        self._seed = seed
         # check if column indices are increasing for looking up genes. This is a way of spotting if the sc_memmap.py
         #  script produced properly strctured sparse files.
         self.assert_increasing_columns = assert_increasing_columns
@@ -198,15 +193,18 @@ def lookup_cell_by_idx(self, idx) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
         )
         return gene_data, col_idxs, feature_ids
 
-    def __getitem__(self, idx: int) -> Item:
-        """Performs a lookup and the required transformation for the model"""  # noqa: D415
+    def __getitem__(self, idx: int) -> types.BertSample:  # noqa: D105
+        rng = np.random.default_rng([self._seed, idx])
+
+        """Performs a lookup and the required transformation for the model"""
         gene_data, col_idxs, feature_ids = self.lookup_cell_by_idx(idx)
         return process_item(
             gene_data,
             col_idxs,
             feature_ids,
             self.tokenizer,
             gene_median=self.gene_medians,
+            rng=rng,
             max_len=self.max_len,
             mask_token_prob=self.mask_token_prob,
             mask_prob=self.mask_prob,
@@ -221,14 +219,15 @@ def process_item(  # noqa: D417
     feature_ids: np.ndarray,
     tokenizer: GeneTokenizer,
     gene_median: dict,
+    rng: np.random.Generator,
     max_len: int = 1024,
     mask_prob: float = 0.15,
     mask_token_prob: float = 0.8,
     random_token_prob: float = 0.1,
     target_sum: int = 10000,
     normalize: bool = True,
     prepend_cls_token: bool = True,
-) -> Item:
+) -> types.BertSample:
     """Process a single item in the dataset.
 
     Optionally performs median normalization and rank ordering. The tokenizers CLS token is added to the beginning
@@ -240,6 +239,7 @@ def process_item(  # noqa: D417
         feature_ids (list): Feature ids for the full dataset.
         tokenizer (Tokenizer): Tokenizer object.
         gene_median (optional(dict)): Dictionary of gene medians. Defaults to None. Expects ensembl IDs to be keys.
+        rng: Random number generator to ensure deterministic results.
         max_len (int): Maximum length of the item. Defaults to 1024. Applies padding to any sequence shorter than max_len and truncates any sequence longer than max_len.
         mask_prob (float): Probability of masking a token. Defaults to 0.15.
         target_sum (int): Target sum for normalization. Defaults to 10000.
@@ -259,14 +259,6 @@ def process_item(  # noqa: D417
     if max_len < 1:
         raise ValueError(f"max_len must be greater than 1, {max_len=}")
 
-    if random_token_prob + mask_token_prob > 1.0:
-        raise ValueError(
-            "Sum of random_token_prob and mask_token_prob must be less than or equal to 1.0, identity_token_prob is any remainder less than 1.0."
-        )
-
-    identity_token_prob = 1.0 - (random_token_prob + mask_token_prob)
-    assert identity_token_prob >= 0.0
-
     if gene_median is None:
         raise ValueError("gene_median must be provided for this tokenizer")
 
@@ -296,62 +288,34 @@ def process_item(  # noqa: D417
         token_ids = token_ids[idxs]
 
     # - select max_len subset, set sample to false so it doesnt permute the already rank ordered expression values.
-    token_ids = sample_or_truncate_plus_pad(
-        token_ids, max_len, tokenizer.token_to_id(tokenizer.pad_token), sample=False
+    token_ids = sample_or_truncate(token_ids, max_len, sample=False)
+
+    masked_tokens, labels, loss_mask = masking.apply_bert_pretraining_mask(
+        tokenized_sequence=torch.from_numpy(token_ids),
+        random_seed=random_utils.get_seed_from_rng(rng),
+        mask_config=masking.BertMaskConfig(
+            mask_token=tokenizer.token_to_id(tokenizer.mask_token),
+            random_tokens=range(5, len(tokenizer.vocab)),
+            mask_prob=mask_prob,
+            mask_token_prob=mask_token_prob,
+            random_token_prob=random_token_prob,
+        ),
     )
 
-    mask = None
-    mask_tokens_positions = None
-    random_tokens_positions = None
-
-    # - masked tokens
-    if mask_prob > 0.0:
-        probs = np.full(token_ids.shape[0], mask_prob)
-        probs[token_ids == tokenizer.token_to_id(tokenizer.pad_token)] = 0.0
-        mask = np.random.binomial(1, probs).astype(bool)
-        mask_tokens_positions = mask & np.random.binomial(1, mask_token_prob, mask.shape).astype(bool)
-        random_tokens_positions = (
-            mask & np.random.binomial(1, random_token_prob, mask.shape).astype(bool) & (~mask_tokens_positions)
-        )
-        # - ensure [CLS] token is masked from the loss. Note that we're dealing with 1d arrays so flattening isn't a problem here.
-        if prepend_cls_token:
-            mask = np.insert(mask, 0, False)
-            mask_tokens_positions = np.insert(mask_tokens_positions, 0, False)
-            random_tokens_positions = np.insert(random_tokens_positions, 0, False)
-
-    # - add [CLS] token, note that token_ids is a 1d array so flattening isn't a problem here.
     if prepend_cls_token:
-        token_ids = np.insert(token_ids, 0, tokenizer.token_to_id(tokenizer.cls_token))
-    attention_mask = token_ids != tokenizer.token_to_id(tokenizer.pad_token)
-
-    labels = np.ones(len(token_ids)) * -1
-
-    if mask is None:
-        # If prob is set to zero, we get None for our mask, which could have unintended side effects.
-        # We abuse the scenario where mask == None
-        labels[mask] = token_ids[mask]
-        mask = np.zeros(shape=token_ids.shape, dtype=bool)
-    else:
-        mask[~attention_mask] = False  # make sure that we aren't doing MLM on [PAD] tokens
-        labels[mask] = token_ids[mask]
-    if mask_tokens_positions is None:
-        mask_tokens_positions = np.zeros_like(mask)
-    if random_tokens_positions is None:
-        random_tokens_positions = np.zeros_like(mask)
-    # identity_tokens = mask & (~mask_tokens_positions) & (~random_tokens_positions), not needed because
-    token_ids[mask_tokens_positions] = tokenizer.token_to_id(tokenizer.mask_token)
-    # There are 5 special tokens in the tokenizer, so we start from 5. TODO make this a parameter of the tokenizer.
-    if random_tokens_positions.sum() > 0:
-        token_ids[random_tokens_positions] = np.random.randint(5, len(tokenizer.vocab), random_tokens_positions.sum())
+        masked_tokens, labels, loss_mask = masking.add_cls_and_eos_tokens(
+            sequence=masked_tokens,
+            labels=labels,
+            loss_mask=loss_mask,
+            cls_token=tokenizer.token_to_id(tokenizer.cls_token),
+        )
 
     # NeMo megatron assumes this return structure.
-    item = {
-        "text": token_ids.astype(np.int64),
-        "types": np.zeros_like(token_ids).astype(np.int64),
-        "attention_mask": attention_mask.astype(np.int64),
-        "labels": labels.astype(np.int64),
-        "loss_mask": mask,
-        "is_random": np.zeros_like(token_ids).astype(np.int64),
+    return {
+        "text": masked_tokens,
+        "types": torch.zeros_like(masked_tokens, dtype=torch.int64),
+        "attention_mask": torch.ones_like(masked_tokens, dtype=torch.int64),
+        "labels": labels,
+        "loss_mask": loss_mask,
+        "is_random": torch.zeros_like(masked_tokens, dtype=torch.int64),
     }
-
-    return item

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

@@ -12,41 +12,30 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-from typing import Sequence
 
 import numpy as np
 
 
-__all__: Sequence[str] = ("sample_or_truncate_plus_pad",)
-
-
-def sample_or_truncate_plus_pad(
-    gene_ids: np.array,
+def sample_or_truncate(
+    gene_ids: np.ndarray,
     max_length: int,
-    pad_token_id: int,
     sample: bool = True,
-) -> np.array:
+) -> np.ndarray:
     """Truncate and pad samples.
 
     Args:
         gene_ids (np.ndarray): Array of gene IDs.
         max_length (int): Maximum length of the samples.
-        pad_token_id (int): ID of the padding token.
         sample (bool, optional): Whether to sample or truncate the samples. Defaults to True.
 
     Returns:
         np.array: Tuple containing the truncated or padded gene IDs.
     """
-    if len(gene_ids) == max_length:
+    if len(gene_ids) <= max_length:
         return gene_ids
 
-    if len(gene_ids) > max_length:  # - sample or truncate
-        if sample:
-            indices = np.random.permutation(len(gene_ids))[:max_length]
-            return gene_ids[indices]
-        else:
-            return gene_ids[:max_length]
-    else:  # - pad
-        pad_tokens = np.full((max_length - len(gene_ids)), pad_token_id, dtype=np.int32)
-        gene_ids = np.concatenate([gene_ids, pad_tokens])
-        return gene_ids
+    if sample:
+        indices = np.random.permutation(len(gene_ids))[:max_length]
+        return gene_ids[indices]
+    else:
+        return gene_ids[:max_length]