From 9c04b233610413ec59b6562cd214b7c02d6dfa01 Mon Sep 17 00:00:00 2001
From: Shahul Alam <alam.shahul@gmail.com>
Date: Tue, 7 Apr 2026 03:41:53 -0400
Subject: [PATCH 01/10] Sync distributed counterparts.

---
 slaf/distributed/worker.py | 55 ++------------------------------------
 slaf/ml/distributed.py     |  5 +---
 2 files changed, 3 insertions(+), 57 deletions(-)

diff --git a/slaf/distributed/worker.py b/slaf/distributed/worker.py
index 74b012d..430685f 100644
--- a/slaf/distributed/worker.py
+++ b/slaf/distributed/worker.py
@@ -92,48 +92,29 @@ def prefetch_worker(
     # - Out of the box: uses Window/Shuffle from slaf.distributed if no factory config
     # - No hardcoded dependencies on slaf.ml in slaf.distributed code
 
-    # Tokenizer is passed as a factory function name (will be created in ml/distributed.py).
-    # Built before window when use_tokenizer_window is set so tokenizer_instance.window exists.
     tokenizer_instance = None
     tokenizer_fn = None
     if processor_config.get("tokenizer_factory"):
-        # Dynamic import and factory call
         tokenizer_config = processor_config["tokenizer_factory"]
         tokenizer_module = importlib.import_module(tokenizer_config["module"])
         tokenizer_class = getattr(tokenizer_module, tokenizer_config["class"])
 
-        # SLAFTokenizer needs a slaf_array, so we need to recreate it from the data source path
-        # Extract slaf_path from data_source_config (assumes Lance path is under slaf_path/expression.lance)
         if data_source_config["type"] == "lance":
             lance_path = data_source_config["path"]
-            # Assume lance_path is like "path/to/slaf/expression.lance"
             slaf_path = lance_path.replace("/expression.lance", "")
 
-            # Recreate SLAFArray in worker
             from slaf.core.slaf import SLAFArray
 
             slaf_array = SLAFArray(slaf_path, load_metadata=False)
-
-            # Create tokenizer instance
             tokenizer_instance = tokenizer_class(
                 slaf_array=slaf_array, **tokenizer_config["kwargs"]
             )
 
-            # Create tokenizer function that works with grouped DataFrame
-            # The grouped DataFrame has gene_sequence and optionally expr_sequence columns
             def tokenize_grouped(
                 grouped_df: pl.DataFrame, schema: DataSchema
             ) -> dict[str, Any]:
-                """Tokenize grouped DataFrame with gene/value sequences.
-
-                For scGPT tokenizers, this enforces the dual-stream contract:
-                tokenized output must include aligned ``input_ids`` and ``values``.
-                """
-                # Extract gene sequences and expression sequences
+                """Tokenize grouped DataFrame with gene/value sequences."""
                 gene_sequences = grouped_df[schema.item_list_key].to_list()
-                is_scgpt_tokenizer = hasattr(tokenizer_instance, "n_expression_bins")
-
-                # Check if we have expression sequences (for scGPT)
                 if (
                     schema.value_list_key
                     and schema.value_list_key in grouped_df.columns
@@ -146,24 +127,7 @@ def tokenize_grouped(
                     input_ids, attention_mask, values = tokenizer_instance.tokenize(
                         gene_sequences
                     )
-
-                if is_scgpt_tokenizer:
-                    if (
-                        not schema.value_list_key
-                        or schema.value_list_key not in grouped_df.columns
-                    ):
-                        raise ValueError(
-                            "scGPT distributed tokenization requires expression/value sequences; "
-                            f"missing grouped column '{schema.value_list_key}'."
-                        )
-                    if values is None:
-                        raise ValueError(
-                            "scGPT distributed tokenization requires dual-stream output; "
-                            "tokenizer returned values=None."
-                        )
-
-                # Return as dict (format expected by processor)
-                result = {
+                return {
                     "input_ids": input_ids,
                     "attention_mask": attention_mask,
                 }
@@ -193,21 +157,6 @@ def tokenize_grouped(
     use_tokenizer_window = processor_config.get("use_tokenizer_window", False)
     if use_tokenizer_window and tokenizer_instance is not None:
         window = tokenizer_instance.window
-        apply_fn = getattr(window, "apply", None)
-        if apply_fn is None:
-            raise TypeError(
-                "tokenizer window must implement apply(df, schema, max_items, **kwargs)"
-            )
-        try:
-            params = inspect.signature(apply_fn).parameters
-        except (TypeError, ValueError):
-            params = None
-        if params is None or "schema" not in params or "max_items" not in params:
-            raise TypeError(
-                "tokenizer window must use the signature "
-                "apply(df, schema, max_items, **kwargs); "
-                "please upgrade slafdb in the worker image."
-            )
     elif processor_config.get("window_factory"):
         # Dynamic import based on config - module path comes from config, not hardcoded
         factory_config = processor_config["window_factory"]
diff --git a/slaf/ml/distributed.py b/slaf/ml/distributed.py
index accec75..430b46c 100644
--- a/slaf/ml/distributed.py
+++ b/slaf/ml/distributed.py
@@ -317,10 +317,7 @@ def __init__(
                 tokenizer_cls = GeneformerTokenizer
             else:
                 tokenizer_cls = ScGPTTokenizer
-            tokenizer_factory_kwargs = {
-                "vocab_size": vocab_size,
-                "max_genes": max_genes,
-            }
+            tokenizer_factory_kwargs = {"vocab_size": vocab_size}
             if tokenizer_type == "scgpt":
                 tokenizer_factory_kwargs["n_expression_bins"] = n_expression_bins
             self.tokenizer = tokenizer_cls(

From 0af5926f2cbfe1fd883e53ea023f4996bcb2748b Mon Sep 17 00:00:00 2001
From: Shahul Alam <alam.shahul@gmail.com>
Date: Sat, 23 May 2026 13:33:02 -0400
Subject: [PATCH 02/10] Clean up configuration objects with OmegaConf.

---
 slaf/distributed/worker.py       |  15 ++-
 slaf/ml/distributed.py           |   4 +-
 tests/test_distributed_worker.py | 173 +++++++++----------------------
 3 files changed, 65 insertions(+), 127 deletions(-)

diff --git a/slaf/distributed/worker.py b/slaf/distributed/worker.py
index 430685f..6483c41 100644
--- a/slaf/distributed/worker.py
+++ b/slaf/distributed/worker.py
@@ -26,8 +26,8 @@
 def prefetch_worker(
     worker_id: str,
     partition_indices: list[int],
-    data_source_config: dict[str, Any],
-    processor_config: dict[str, Any],
+    data_source_config: DictConfig,
+    processor_config: DictConfig,
     queue: Any,  # Modal Queue or any queue-like object
     n_scanners: int = 8,
     prefetch_batch_count: int = 32,
@@ -56,6 +56,15 @@ def prefetch_worker(
     Returns:
         Dictionary with worker metrics
     """
+    tokenizer_config = processor_config.tokenizer_config
+    shuffle_factory_config = processor_config.shuffle_factory
+    window_factory_config = processor_config.window_factory
+    use_tokenizer_window = bool(processor_config.use_tokenizer_window)
+    continuity_check = str(processor_config.continuity_check)
+    max_items = int(processor_config.max_items)
+    seed_value = int(processor_config.seed)
+    window_kwargs = processor_config.window_kwargs or {}
+
     # Import data source (generic)
     from slaf.distributed.data_source import LanceDataSource
 
@@ -279,7 +288,7 @@ def read_from_partition(partition_idx: int, batches_to_read: int):
     # Process partitions using Mixture of Scanners (MoS) approach
     total_batches = 0
     total_rows = 0
-    epochs = processor_config.get("n_epochs", 1)
+    epochs = processor_config.n_epochs
 
     # Async writer: use Modal's .aio so we don't block on network I/O (see modal.com/docs/guide/queues)
     writer_queue: queue_module.Queue[list[dict[str, Any]] | None] = queue_module.Queue(
diff --git a/slaf/ml/distributed.py b/slaf/ml/distributed.py
index 430b46c..e244255 100644
--- a/slaf/ml/distributed.py
+++ b/slaf/ml/distributed.py
@@ -95,8 +95,8 @@ def create_app(
     def distributed_prefetch_worker(
         worker_id: str,
         partition_indices: list[int],
-        data_source_config: dict[str, Any],
-        processor_config: dict[str, Any],
+        data_source_config: DictConfig,
+        processor_config: DictConfig,
         queue_name: str,
         n_scanners: int = 8,
         prefetch_batch_count: int = 32,
diff --git a/tests/test_distributed_worker.py b/tests/test_distributed_worker.py
index c06df3b..32e2636 100644
--- a/tests/test_distributed_worker.py
+++ b/tests/test_distributed_worker.py
@@ -11,6 +11,8 @@
 import polars as pl
 
 from slaf.distributed.data_source import DataSource
+from omegaconf import OmegaConf
+
 from slaf.distributed.worker import prefetch_worker
 
 
@@ -64,6 +66,33 @@ def pop(self, key: str, default: Any = None) -> Any:
         return self._store.pop(key, default)
 
 
+def make_data_source_config():
+    return OmegaConf.create({"type": "lance", "path": "/fake/path"})
+
+
+def make_processor_config(**overrides):
+    config = {
+        "schema": {
+            "group_key": "group_id",
+            "item_key": "item_id",
+            "value_key": "value",
+            "item_list_key": "item_list",
+        },
+        "tokenizer_config": None,
+        "shuffle_factory": None,
+        "window_factory": None,
+        "use_tokenizer_window": False,
+        "continuity_check": "sequential",
+        "max_items": 5,
+        "seed": 42,
+        "n_epochs": 1,
+        "window_kwargs": {},
+        "enable_cross_worker_boundary_merging": False,
+    }
+    config.update(overrides)
+    return OmegaConf.create(config)
+
+
 def create_test_dataframe(
     group_key: str = "group_id",
     item_key: str = "item_id",
@@ -102,18 +131,8 @@ def test_worker_initialization(self, mock_lance_ds):
         queue = MockQueue()
 
         # Configs
-        data_source_config = {"type": "lance", "path": "/fake/path"}
-        processor_config = {
-            "schema": {
-                "group_key": "group_id",
-                "item_key": "item_id",
-                "value_key": "value",
-                "item_list_key": "item_list",
-            },
-            "max_items": 5,
-            "seed": 42,
-            "n_epochs": 1,
-        }
+        data_source_config = make_data_source_config()
+        processor_config = make_processor_config()
 
         # Call worker
         result = prefetch_worker(
@@ -147,18 +166,8 @@ def test_worker_single_partition(self, mock_lance_ds):
         queue = MockQueue()
 
         # Configs
-        data_source_config = {"type": "lance", "path": "/fake/path"}
-        processor_config = {
-            "schema": {
-                "group_key": "group_id",
-                "item_key": "item_id",
-                "value_key": "value",
-                "item_list_key": "item_list",
-            },
-            "max_items": 5,
-            "seed": 42,
-            "n_epochs": 1,
-        }
+        data_source_config = make_data_source_config()
+        processor_config = make_processor_config()
 
         # Call worker
         result = prefetch_worker(
@@ -199,18 +208,8 @@ def create_reader(partition_index, batch_size):
         queue = MockQueue()
 
         # Configs
-        data_source_config = {"type": "lance", "path": "/fake/path"}
-        processor_config = {
-            "schema": {
-                "group_key": "group_id",
-                "item_key": "item_id",
-                "value_key": "value",
-                "item_list_key": "item_list",
-            },
-            "max_items": 5,
-            "seed": 42,
-            "n_epochs": 1,
-        }
+        data_source_config = make_data_source_config()
+        processor_config = make_processor_config()
 
         # Call worker
         result = prefetch_worker(
@@ -243,18 +242,8 @@ def test_worker_prefetch_batch_count(self, mock_lance_ds):
         queue = MockQueue()
 
         # Configs
-        data_source_config = {"type": "lance", "path": "/fake/path"}
-        processor_config = {
-            "schema": {
-                "group_key": "group_id",
-                "item_key": "item_id",
-                "value_key": "value",
-                "item_list_key": "item_list",
-            },
-            "max_items": 5,
-            "seed": 42,
-            "n_epochs": 1,
-        }
+        data_source_config = make_data_source_config()
+        processor_config = make_processor_config()
 
         # Call worker with prefetch_batch_count
         result = prefetch_worker(
@@ -287,18 +276,8 @@ def test_worker_prefetch_batch_size(self, mock_lance_ds):
         queue = MockQueue()
 
         # Configs
-        data_source_config = {"type": "lance", "path": "/fake/path"}
-        processor_config = {
-            "schema": {
-                "group_key": "group_id",
-                "item_key": "item_id",
-                "value_key": "value",
-                "item_list_key": "item_list",
-            },
-            "max_items": 5,
-            "seed": 42,
-            "n_epochs": 1,
-        }
+        data_source_config = make_data_source_config()
+        processor_config = make_processor_config()
 
         # Call worker with prefetch_batch_size
         result = prefetch_worker(
@@ -331,18 +310,8 @@ def test_worker_max_batches(self, mock_lance_ds):
         queue = MockQueue()
 
         # Configs
-        data_source_config = {"type": "lance", "path": "/fake/path"}
-        processor_config = {
-            "schema": {
-                "group_key": "group_id",
-                "item_key": "item_id",
-                "value_key": "value",
-                "item_list_key": "item_list",
-            },
-            "max_items": 5,
-            "seed": 42,
-            "n_epochs": 1,
-        }
+        data_source_config = make_data_source_config()
+        processor_config = make_processor_config()
 
         # Call worker with max_batches
         result = prefetch_worker(
@@ -377,20 +346,10 @@ def test_worker_cross_worker_boundary_merging(self, mock_lance_ds):
         kv_store = MockKVStore()
 
         # Configs with cross-worker merging enabled
-        data_source_config = {"type": "lance", "path": "/fake/path"}
-        processor_config = {
-            "schema": {
-                "group_key": "group_id",
-                "item_key": "item_id",
-                "value_key": "value",
-                "item_list_key": "item_list",
-            },
-            "max_items": 5,
-            "seed": 42,
-            "n_epochs": 1,
-            "enable_cross_worker_boundary_merging": True,
-            "continuity_check": "sequential",
-        }
+        data_source_config = make_data_source_config()
+        processor_config = make_processor_config(
+            enable_cross_worker_boundary_merging=True
+        )
 
         # Call worker with KV store
         result = prefetch_worker(
@@ -424,18 +383,8 @@ def test_worker_partition_exhaustion(self, mock_lance_ds):
         queue = MockQueue()
 
         # Configs
-        data_source_config = {"type": "lance", "path": "/fake/path"}
-        processor_config = {
-            "schema": {
-                "group_key": "group_id",
-                "item_key": "item_id",
-                "value_key": "value",
-                "item_list_key": "item_list",
-            },
-            "max_items": 5,
-            "seed": 42,
-            "n_epochs": 1,
-        }
+        data_source_config = make_data_source_config()
+        processor_config = make_processor_config()
 
         # Call worker
         result = prefetch_worker(
@@ -469,18 +418,8 @@ def create_reader(partition_index, batch_size):
         queue = MockQueue()
 
         # Configs
-        data_source_config = {"type": "lance", "path": "/fake/path"}
-        processor_config = {
-            "schema": {
-                "group_key": "group_id",
-                "item_key": "item_id",
-                "value_key": "value",
-                "item_list_key": "item_list",
-            },
-            "max_items": 5,
-            "seed": 42,
-            "n_epochs": 1,
-        }
+        data_source_config = make_data_source_config()
+        processor_config = make_processor_config()
 
         # Call worker - should handle error gracefully
         result = prefetch_worker(
@@ -513,18 +452,8 @@ def test_worker_metrics(self, mock_lance_ds):
         queue = MockQueue()
 
         # Configs
-        data_source_config = {"type": "lance", "path": "/fake/path"}
-        processor_config = {
-            "schema": {
-                "group_key": "group_id",
-                "item_key": "item_id",
-                "value_key": "value",
-                "item_list_key": "item_list",
-            },
-            "max_items": 5,
-            "seed": 42,
-            "n_epochs": 1,
-        }
+        data_source_config = make_data_source_config()
+        processor_config = make_processor_config()
 
         # Call worker
         result = prefetch_worker(

From 6636e9d562fda7cc490275258a4922c88dc0f3bf Mon Sep 17 00:00:00 2001
From: Shahul Alam <alam.shahul@gmail.com>
Date: Mon, 20 Apr 2026 16:43:48 -0400
Subject: [PATCH 03/10] Refactor SLAF tokenizer interfaces

Rework SLAFDataLoader and DistributedSLAFDataLoader to accept instantiated tokenizers instead of selecting tokenizers by type.

Add tokenizer-owned apply(), tokenize_grouped(), tokenizer_name, and
  get_factory_kwargs() entrypoints so local and distributed preprocessing depend on the tokenizer interface rather than tokenizer.window internals.

Update dataset and worker tokenization paths to preserve scGPT dual-stream values output while keeping the shared
  schema rooted in slaf.core.tabular_schema.

Generalize tokenizer configuration
---
 pyproject.toml                   |   7 +-
 slaf/distributed/worker.py       | 101 +++++++++---------
 slaf/ml/dataloaders.py           | 130 +++++++++++------------
 slaf/ml/datasets.py              |  16 ++-
 slaf/ml/distributed.py           | 138 ++++++++++++------------
 slaf/ml/tokenizers.py            | 175 ++++++++++++++++++++++++++++++-
 tests/test_dataloaders.py        | 164 ++++++++++++++++-------------
 tests/test_distributed_worker.py |   2 +
 tests/test_tokenizers.py         | 100 ++++++++++--------
 uv.lock                          |  52 ++++++---
 10 files changed, 549 insertions(+), 336 deletions(-)

diff --git a/pyproject.toml b/pyproject.toml
index 17babad..7740fd1 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -37,7 +37,12 @@ dependencies = [
 
 [project.optional-dependencies]
 # Machine learning - for ML modules
-ml = ["torch>=2.5.0", "tiledb>=0.34.2", "tiledbsoma>=1.17.1"]
+ml = [
+    "torch>=2.5.0",
+    "tiledb>=0.34.2",
+    "tiledbsoma>=1.17.1",
+    "omegaconf>=2.3.0",
+]
 
 # Advanced single-cell tools
 advanced = ["igraph>=0.11.9", "leidenalg>=0.10.2"]
diff --git a/slaf/distributed/worker.py b/slaf/distributed/worker.py
index 6483c41..d340d04 100644
--- a/slaf/distributed/worker.py
+++ b/slaf/distributed/worker.py
@@ -5,7 +5,6 @@
 
 import asyncio
 import importlib
-import inspect
 import pickle
 import queue as queue_module
 import random
@@ -17,6 +16,7 @@
 
 import polars as pl
 from loguru import logger
+from omegaconf import DictConfig, OmegaConf
 
 # Modal queue item size limit (1 MiB); we compress samples to stay under this
 # See https://modal.com/docs/guide/queues and https://modal.com/docs/reference/modal.Queue
@@ -68,20 +68,20 @@ def prefetch_worker(
     # Import data source (generic)
     from slaf.distributed.data_source import LanceDataSource
 
-    if data_source_config["type"] == "lance":
+    if data_source_config.type == "lance":
         logger.info(
             "[{worker_id}] Creating LanceDataSource for: {path}",
             worker_id=worker_id,
-            path=data_source_config["path"],
+            path=data_source_config.path,
         )
-        data_source = LanceDataSource(data_source_config["path"])
+        data_source = LanceDataSource(data_source_config.path)
         logger.info(
             "[{worker_id}] DataSource created, partition count: {count}",
             worker_id=worker_id,
             count=data_source.get_partition_count(),
         )
     else:
-        raise ValueError(f"Unknown data source type: {data_source_config['type']}")
+        raise ValueError(f"Unknown data source type: {data_source_config.type}")
 
     # Import processor (generic)
     from slaf.core.tabular_schema import DataSchema
@@ -103,40 +103,38 @@ def prefetch_worker(
 
     tokenizer_instance = None
     tokenizer_fn = None
-    if processor_config.get("tokenizer_factory"):
-        tokenizer_config = processor_config["tokenizer_factory"]
-        tokenizer_module = importlib.import_module(tokenizer_config["module"])
-        tokenizer_class = getattr(tokenizer_module, tokenizer_config["class"])
-
-        if data_source_config["type"] == "lance":
-            lance_path = data_source_config["path"]
+    if tokenizer_config:
+        # Recreate dataset access in the worker so the tokenizer can be reconstructed locally.
+        if data_source_config.type == "lance":
+            lance_path = data_source_config.path
             slaf_path = lance_path.replace("/expression.lance", "")
 
             from slaf.core.slaf import SLAFArray
 
             slaf_array = SLAFArray(slaf_path, load_metadata=False)
-            tokenizer_instance = tokenizer_class(
-                slaf_array=slaf_array, **tokenizer_config["kwargs"]
+
+            target = tokenizer_config.type
+            module_name, class_name = target.rsplit(".", 1)
+            tokenizer_module = importlib.import_module(module_name)
+            tokenizer_class = getattr(tokenizer_module, class_name)
+            tokenizer_kwargs = (
+                dict(tokenizer_config.args)
+                if "args" in tokenizer_config and tokenizer_config.args is not None
+                else {}
             )
+            tokenizer_instance = tokenizer_class( slaf_array=slaf_array, **tokenizer_kwargs)
 
             def tokenize_grouped(
                 grouped_df: pl.DataFrame, schema: DataSchema
             ) -> dict[str, Any]:
-                """Tokenize grouped DataFrame with gene/value sequences."""
-                gene_sequences = grouped_df[schema.item_list_key].to_list()
-                if (
-                    schema.value_list_key
-                    and schema.value_list_key in grouped_df.columns
-                ):
-                    expr_sequences = grouped_df[schema.value_list_key].to_list()
-                    input_ids, attention_mask, values = tokenizer_instance.tokenize(
-                        gene_sequences, expr_sequences
-                    )
-                else:
-                    input_ids, attention_mask, values = tokenizer_instance.tokenize(
-                        gene_sequences
-                    )
-                return {
+                """Tokenize grouped DataFrame with tokenizer-owned grouping contract."""
+                input_ids, attention_mask, values = tokenizer_instance.tokenize_grouped(
+                    grouped_df,
+                    schema=schema,
+                )
+
+                # Return as dict (format expected by processor)
+                result = {
                     "input_ids": input_ids,
                     "attention_mask": attention_mask,
                 }
@@ -147,41 +145,46 @@ def tokenize_grouped(
             tokenizer_fn = tokenize_grouped
 
     # Shuffle: use factory if provided, otherwise use default
-    if processor_config.get("shuffle_factory"):
+    if shuffle_factory_config:
         # Dynamic import based on config - module path comes from config, not hardcoded
-        factory_config = processor_config["shuffle_factory"]
+        factory_config = shuffle_factory_config
         shuffle_module = __import__(
-            factory_config["module"], fromlist=[factory_config["function"]]
+            factory_config.module, fromlist=[factory_config.function]
         )
-        shuffle_factory = getattr(shuffle_module, factory_config["function"])
-        shuffle = shuffle_factory(
-            factory_config["type"], **factory_config.get("kwargs", {})
+        shuffle_factory = getattr(shuffle_module, factory_config.function)
+        shuffle_kwargs = (
+            dict(factory_config.kwargs)
+            if "kwargs" in factory_config and factory_config.kwargs is not None
+            else {}
         )
+        shuffle = shuffle_factory(factory_config.type, **shuffle_kwargs)
     else:
         # Use generic implementation (works out of the box)
         shuffle = Shuffle()
 
     # Window: use tokenizer.window when use_tokenizer_window (aligns with ml tokenizers owning Window);
     # else factory if provided; otherwise generic Window (works out of the box).
-    use_tokenizer_window = processor_config.get("use_tokenizer_window", False)
     if use_tokenizer_window and tokenizer_instance is not None:
-        window = tokenizer_instance.window
-    elif processor_config.get("window_factory"):
+        window = tokenizer_instance
+    elif window_factory_config:
         # Dynamic import based on config - module path comes from config, not hardcoded
-        factory_config = processor_config["window_factory"]
+        factory_config = window_factory_config
         window_module = __import__(
-            factory_config["module"], fromlist=[factory_config["function"]]
+            factory_config.module, fromlist=[factory_config.function]
         )
-        window_factory = getattr(window_module, factory_config["function"])
-        window = window_factory(
-            factory_config["type"], **factory_config.get("kwargs", {})
+        window_factory = getattr(window_module, factory_config.function)
+        window_factory_kwargs = (
+            dict(factory_config.kwargs)
+            if "kwargs" in factory_config and factory_config.kwargs is not None
+            else {}
         )
+        window = window_factory(factory_config.type, **window_factory_kwargs)
     else:
         # Use generic implementation (works out of the box)
         window = Window()
 
     # Create processor with data schema
-    schema = DataSchema(**processor_config["schema"])
+    schema = DataSchema(**OmegaConf.to_container(processor_config.schema, resolve=True))
 
     # Create boundary handler with KV store support
     # partial_groups_kv is passed as a parameter (created by caller)
@@ -189,7 +192,7 @@ def tokenize_grouped(
 
     boundary_handler = GroupBoundaryHandler(
         schema=schema,
-        continuity_check=processor_config.get("continuity_check", "sequential"),
+        continuity_check=continuity_check,
         partial_groups_kv=partial_groups_kv,
     )
 
@@ -199,10 +202,10 @@ def tokenize_grouped(
         shuffle=shuffle,
         tokenizer=tokenizer_fn,
         boundary_handler=boundary_handler,
-        max_items=processor_config.get("max_items", 1024),
-        seed=processor_config.get("seed", 42),
-        continuity_check=processor_config.get("continuity_check", "sequential"),
-        **processor_config.get("window_kwargs", {}),
+        max_items=max_items,
+        seed=seed_value,
+        continuity_check=continuity_check,
+        **dict(window_kwargs),
     )
 
     # Queue is passed as a parameter (created by caller)
diff --git a/slaf/ml/dataloaders.py b/slaf/ml/dataloaders.py
index a1356c3..866e8d7 100644
--- a/slaf/ml/dataloaders.py
+++ b/slaf/ml/dataloaders.py
@@ -5,7 +5,7 @@
 from slaf.core.slaf import SLAFArray
 
 from .expression_preprocessor import ExpressionPreprocessor
-from .tokenizers import GeneformerTokenizer, ScGPTTokenizer, SLAFTokenizer
+from .tokenizers import SLAFTokenizer
 
 # Try to import torch, but make it optional
 try:
@@ -188,7 +188,8 @@ class SLAFDataLoader:
     Examples:
         >>> # Basic usage with default settings (MoS loading)
         >>> slaf_array = SLAFArray("path/to/data.slaf")
-        >>> dataloader = SLAFDataLoader(slaf_array)
+        >>> tokenizer = GeneformerTokenizer(slaf_array)
+        >>> dataloader = SLAFDataLoader(slaf_array, tokenizer=tokenizer)
         >>> for batch in dataloader:
         ...     print(f"Batch shape: {batch['input_ids'].shape}")
         ...     print(f"Cell IDs: {batch['cell_ids']}")
@@ -235,14 +236,14 @@ class SLAFDataLoader:
         Number of epochs: 5
 
         >>> # Custom configuration for training
+        >>> tokenizer = ScGPTTokenizer(slaf_array=slaf_array, max_genes=1024)
         >>> dataloader = SLAFDataLoader(
         ...     slaf_array=slaf_array,
-        ...     tokenizer_type="scgpt",
+        ...     tokenizer=tokenizer,
         ...     batch_size=64,
-        ...     max_genes=1024
         ... )
         >>> print(f"Tokenizer type: {dataloader.tokenizer_type}")
-        Tokenizer type: scgpt
+        Tokenizer type: ScGPTTokenizer
 
         >>> # Training loop example
         >>> for batch_idx, batch in enumerate(dataloader):
@@ -255,12 +256,12 @@ class SLAFDataLoader:
         >>> print("Training loop completed")
         Training loop completed
 
-        >>> # Error handling for invalid tokenizer type
+        >>> # Error handling for missing tokenizer
         >>> try:
-        ...     dataloader = SLAFDataLoader(slaf_array, tokenizer_type="invalid")
+        ...     dataloader = SLAFDataLoader(slaf_array)
         ... except ValueError as e:
         ...     print(f"Error: {e}")
-        Error: Unsupported tokenizer type: invalid
+        Error: tokenizer must be provided unless raw_mode=True.
     """
 
     device: Optional["torch.device"]  # type: ignore
@@ -269,11 +270,8 @@ class SLAFDataLoader:
     def __init__(
         self,
         slaf_array: SLAFArray,
-        tokenizer_type: str = "geneformer",
+        tokenizer: SLAFTokenizer | None = None,
         batch_size: int = 32,
-        max_genes: int = 2048,
-        vocab_size: int = 50000,
-        n_expression_bins: int = 10,
         n_epochs: int = 1,  # Add n_epochs parameter
         raw_mode: bool = False,  # Add raw_mode parameter
         verbose: bool = True,  # Add verbose parameter
@@ -295,20 +293,8 @@ def __init__(
                        Must be a valid SLAFArray with proper Lance dataset structure.
 
             # Tokenization Configuration
-            tokenizer_type: Tokenization strategy to use. Options: "geneformer", "scgpt".
-                          Geneformer uses ranked gene sequences. scGPT uses ranked genes
-                          with a parallel expression list (binned or raw), then
-                          ``ScGPTTokenizer`` builds aligned dual-stream ``input_ids`` and
-                          ``values`` tensors. Ignored when raw_mode=True.
-            max_genes: Maximum number of genes to include in each cell's tokenization.
-                     For Geneformer: caps sequence length (CLS + genes + padding).
-                     For scGPT: top ``max_genes`` genes per cell; each stream has length
-                     ``max_genes + 2`` (CLS, genes, SEP) in the tokenizer.
-            vocab_size: Size of the tokenizer vocabulary. Higher values allow more
-                       genes but use more memory. Range: 1000-100000, default: 50000.
-            n_expression_bins: Number of expression level bins for scGPT discretization.
-                             Higher values provide finer expression resolution.
-                             Range: 1-1000, default: 10.
+            tokenizer: Instantiated tokenizer used for tokenized mode.
+                      Required unless raw_mode=True.
 
             # Training Configuration
             batch_size: Number of cells per batch. Larger batches use more memory
@@ -356,7 +342,7 @@ def __init__(
                     Default: True.
 
         Raises:
-            ValueError: If tokenizer_type is not supported or parameters are invalid.
+            ValueError: If tokenizer configuration or parameters are invalid.
             RuntimeError: If PyTorch is not available or datasets module is missing.
             TypeError: If slaf_array is not a valid SLAFArray instance.
             ImportError: If required dependencies are not available.
@@ -371,7 +357,8 @@ def __init__(
         Examples:
             >>> # Basic initialization (MoS is now default)
             >>> slaf_array = SLAFArray("path/to/data.slaf")
-            >>> dataloader = SLAFDataLoader(slaf_array)
+            >>> tokenizer = GeneformerTokenizer(slaf_array)
+            >>> dataloader = SLAFDataLoader(slaf_array, tokenizer=tokenizer)
             >>> print(f"Batch size: {dataloader.batch_size}")
             Batch size: 32
             >>> print(f"MoS enabled: {dataloader.use_mixture_of_scanners}")
@@ -421,7 +408,6 @@ def __init__(
         """
         self.slaf_array = slaf_array
         self.batch_size = batch_size
-        self.max_genes = max_genes
         self.n_epochs = n_epochs
         self.raw_mode = raw_mode  # Add raw_mode attribute
         self.verbose = verbose  # Add verbose attribute
@@ -436,6 +422,7 @@ def __init__(
         self.parallelize_fragment_reads = (
             parallelize_fragment_reads  # Add parallelize_fragment_reads attribute
         )
+        self.expression_preprocessor = expression_preprocessor
 
         # Validate MoS parameters
         if self.use_mixture_of_scanners:
@@ -462,48 +449,36 @@ def __init__(
         # Initialize tokenizer (only needed for non-raw mode)
 
         if not self.raw_mode:
-            if tokenizer_type == "geneformer":
-                self.tokenizer = GeneformerTokenizer(
-                    slaf_array=slaf_array,
-                    vocab_size=vocab_size,
-                    max_genes=max_genes,
-                )
-            elif tokenizer_type == "scgpt":
-                self.tokenizer = ScGPTTokenizer(
-                    slaf_array=slaf_array,
-                    vocab_size=vocab_size,
-                    n_expression_bins=n_expression_bins,
-                    max_genes=max_genes,
-                )
-            else:
-                raise ValueError(
-                    "tokenizer_type must be one of ['geneformer', 'scgpt']; "
-                    f"{tokenizer_type=} is not supported."
-                )
-
-            # Get special tokens from tokenizer
+            if tokenizer is None:
+                raise ValueError("tokenizer must be provided unless raw_mode=True.")
+            self.tokenizer = tokenizer
+            self.max_genes = self.tokenizer.max_genes
+            self.tokenizer_type = self.tokenizer.name
             self.special_tokens = self.tokenizer.special_tokens
         else:
-            # For raw mode, we don't need a tokenizer
+            if tokenizer is not None:
+                raise ValueError("raw_mode=True is incompatible with tokenizer.")
             self.tokenizer = None
             self.special_tokens = None
+            self.tokenizer_type = "raw"
+            self.max_genes = 0
 
         # Use IterableDataset
         self._dataset = SLAFIterableDataset(
-            slaf_array=slaf_array,
+            slaf_array=self.slaf_array,
             tokenizer=self.tokenizer,
-            batch_size=batch_size,
+            batch_size=self.batch_size,
             seed=42,  # TODO: make configurable
-            max_queue_size=max_queue_size,  # Pass max_queue_size to dataset
-            n_epochs=n_epochs,  # Pass n_epochs to dataset
-            raw_mode=raw_mode,  # Pass raw_mode to dataset
-            verbose=verbose,  # Pass verbose to dataset
-            batches_per_chunk=batches_per_chunk,  # Pass batches_per_chunk to dataset
-            by_fragment=by_fragment,  # Pass by_fragment to dataset
-            use_mixture_of_scanners=use_mixture_of_scanners,  # Pass MoS to dataset
-            n_scanners=n_scanners,  # Pass n_scanners to dataset
-            prefetch_batch_size=prefetch_batch_size,  # Pass prefetch_batch_size to dataset
-            parallelize_fragment_reads=parallelize_fragment_reads,  # Pass parallelize_fragment_reads
+            max_queue_size=self.max_queue_size,  # Pass max_queue_size to dataset
+            n_epochs=self.n_epochs,  # Pass n_epochs to dataset
+            raw_mode=self.raw_mode,  # Pass raw_mode to dataset
+            verbose=self.verbose,  # Pass verbose to dataset
+            batches_per_chunk=self.batches_per_chunk,  # Pass batches_per_chunk to dataset
+            by_fragment=self.by_fragment,  # Pass by_fragment to dataset
+            use_mixture_of_scanners=self.use_mixture_of_scanners,  # Pass MoS to dataset
+            n_scanners=self.n_scanners,  # Pass n_scanners to dataset
+            prefetch_batch_size=self.prefetch_batch_size,  # Pass prefetch_batch_size to dataset
+            parallelize_fragment_reads=self.parallelize_fragment_reads,  # Pass parallelize_fragment_reads
             prefetcher_ready_timeout=prefetcher_ready_timeout,  # Pass prefetcher_ready_timeout
             expression_preprocessor=expression_preprocessor,
         )
@@ -525,12 +500,13 @@ def __iter__(self):
         Yields:
             dict: Batch dictionary containing:
                 - **Tokenized mode** (raw_mode=False):
-                    - input_ids: Pre-tokenized gene expression data (torch.Tensor)
+                    - input_ids: Pre-tokenized gene identity data (torch.Tensor)
+                    - values: Aligned expression/value stream (torch.Tensor)
                     - attention_mask: Boolean mask indicating valid tokens (torch.Tensor)
                     - cell_ids: Integer IDs of cells in the batch (torch.Tensor)
                 - **Raw mode** (raw_mode=True):
-                    - x: Raw cell × gene data as Polars DataFrame
-                    - cell_ids: List of cell integer IDs in the batch
+                    - x: Raw sparse cell × gene data (polars.DataFrame)
+                    - cell_ids: Integer IDs of cells in the batch (list[int])
                 - **Multi-epoch** (when n_epochs > 1):
                     - epoch: Current epoch number (int)
 
@@ -539,19 +515,24 @@ def __iter__(self):
             The training loop should handle device transfer as needed.
 
         Raises:
-            ValueError: If the tokenizer type is not supported.
+            ValueError: If tokenizer configuration is invalid.
             RuntimeError: If batch processing fails.
 
         Examples:
             >>> # Basic iteration (tokenized mode)
             >>> slaf_array = SLAFArray("path/to/data.slaf")
-            >>> dataloader = SLAFDataLoader(slaf_array, batch_size=16)
+            >>> tokenizer = GeneformerTokenizer(slaf_array)
+            >>> dataloader = SLAFDataLoader(
+            ...     slaf_array,
+            ...     tokenizer=tokenizer,
+            ...     batch_size=16,
+            ... )
             >>> for batch in dataloader:
             ...     print(f"Batch keys: {list(batch.keys())}")
             ...     print(f"Input shape: {batch['input_ids'].shape}")
             ...     print(f"Cell IDs: {batch['cell_ids']}")
             ...     break
-            Batch keys: ['input_ids', 'attention_mask', 'cell_ids']
+            Batch keys: ['input_ids', 'values', 'attention_mask', 'cell_ids']
             Input shape: (16, 2048)
             Cell IDs: tensor([0, 1, 2, ..., 13, 14, 15])
 
@@ -581,6 +562,7 @@ def __iter__(self):
             ...         if 'input_ids' in batch:  # Tokenized mode
             ...             input_ids = batch["input_ids"]
             ...             attention_mask = batch["attention_mask"]
+            ...             values = batch["values"]
             ...             cell_ids = batch["cell_ids"]
             ...         else:  # Raw mode
             ...             x = batch["x"]
@@ -596,9 +578,15 @@ def __iter__(self):
             Processed batch 1
             Processed batch 2
 
-            >>> # Different tokenizer types
-            >>> dataloader_geneformer = SLAFDataLoader(slaf_array, tokenizer_type="geneformer")
-            >>> dataloader_scgpt = SLAFDataLoader(slaf_array, tokenizer_type="scgpt")
+            >>> # Different tokenizer instances
+            >>> dataloader_geneformer = SLAFDataLoader(
+            ...     slaf_array,
+            ...     tokenizer=GeneformerTokenizer(slaf_array),
+            ... )
+            >>> dataloader_scgpt = SLAFDataLoader(
+            ...     slaf_array,
+            ...     tokenizer=ScGPTTokenizer(slaf_array),
+            ... )
             >>>
             >>> # Compare batch shapes
             >>> for batch in dataloader_geneformer:
diff --git a/slaf/ml/datasets.py b/slaf/ml/datasets.py
index 857e54b..9600f4a 100644
--- a/slaf/ml/datasets.py
+++ b/slaf/ml/datasets.py
@@ -1042,10 +1042,10 @@ def load_prefetch_batch(self) -> PrefetchBatch:
                     if tokenizer is None:
                         raise RuntimeError("Tokenizer is required for tokenized mode")
 
-                    grouped = tokenizer.window.apply(
+                    grouped = tokenizer.apply(
                         shuffled_df,
-                        SLAF_LANCE_COO_SCHEMA,
-                        tokenizer.max_genes,
+                        schema=SLAF_LANCE_COO_SCHEMA,
+                        max_items=tokenizer.max_genes,
                         **window_params,
                     )
                     window_time = time.time() - window_start
@@ -1056,13 +1056,9 @@ def load_prefetch_batch(self) -> PrefetchBatch:
                     if self.tokenizer is None:
                         raise RuntimeError("Tokenizer is required for tokenized mode")
 
-                    input_ids, attention_mask, values = self.tokenizer.tokenize(
-                        gene_sequences=grouped["gene_sequence"].to_list(),
-                        expr_sequences=(
-                            grouped["expr_sequence"].to_list()
-                            if "expr_sequence" in grouped.columns
-                            else None
-                        ),
+                    input_ids, attention_mask, values = self.tokenizer.tokenize_grouped(
+                        grouped,
+                        schema=SLAF_LANCE_COO_SCHEMA,
                     )
 
                     tokenize_time = time.time() - tokenize_start
diff --git a/slaf/ml/distributed.py b/slaf/ml/distributed.py
index e244255..4ba2a34 100644
--- a/slaf/ml/distributed.py
+++ b/slaf/ml/distributed.py
@@ -10,6 +10,7 @@
 
 import modal
 from loguru import logger
+from omegaconf import DictConfig, OmegaConf
 
 from slaf.core.slaf import SLAFArray
 from slaf.core.tabular_schema import DataSchema
@@ -21,7 +22,7 @@
 
 # Import SLAF-specific components (for type hints and adapters)
 from slaf.ml.expression_preprocessor import ExpressionPreprocessor
-from slaf.ml.tokenizers import GeneformerTokenizer, ScGPTTokenizer
+from slaf.ml.tokenizers import SLAFTokenizer
 
 # Configure Modal image for SLAF workers.
 # Cache bust must run *before* ``uv_pip_install`` so the install layer is not
@@ -113,6 +114,11 @@ def distributed_prefetch_worker(
         """
         from slaf.distributed.worker import prefetch_worker
 
+        if not isinstance(data_source_config, DictConfig):
+            data_source_config = OmegaConf.create(data_source_config)
+        if not isinstance(processor_config, DictConfig):
+            processor_config = OmegaConf.create(processor_config)
+
         # Inline Queue/Dict open — do not call module helpers here. ``serialized=True`` workers
         # unpickle against site-packages slaf; a PyPI lag behind your deploy machine would
         # raise DeserializationError if this referenced new helpers on the training side only.
@@ -126,7 +132,7 @@ def distributed_prefetch_worker(
             queue = modal.Queue.from_name(queue_name, create_if_missing=True)
         partial_groups_kv = None
         if (
-            processor_config.get("enable_cross_worker_boundary_merging", False)
+            bool(processor_config.enable_cross_worker_boundary_merging)
             and partial_groups_kv_name
         ):
             if modal_queue_environment is not None:
@@ -209,12 +215,13 @@ class DistributedSLAFDataLoader:
     (raw_mode=True, wait for queue size >= 50, then iterate).
     """
 
-    tokenizer: GeneformerTokenizer | ScGPTTokenizer | None
+    tokenizer: SLAFTokenizer | None
 
     def __init__(
         self,
         slaf_array: SLAFArray,
-        tokenizer_type: str = "geneformer",
+        tokenizer: SLAFTokenizer | None = None,
+        tokenizer_config: DictConfig | dict[str, Any] | None = None,
         n_workers: int = 64,
         n_scanners: int = 16,
         cpu: float = 8,
@@ -222,9 +229,6 @@ def __init__(
         prefetch_batch_size: int = 16384,  # 16K rows per Lance batch (small default for testing; raise e.g. 262144 for prod)
         prefetch_batch_count: int = 32,
         batch_size: int = 32,
-        max_genes: int = 1024,
-        vocab_size: int = 50000,
-        n_expression_bins: int = 10,
         n_epochs: int = 1,
         raw_mode: bool = False,
         return_tensors: bool = True,
@@ -241,8 +245,10 @@ def __init__(
 
         Args:
             slaf_array: SLAFArray instance containing the data
-            tokenizer_type [PRODUCER]: Tokenization strategy ("geneformer", "scgpt", or "raw")
-                           If "raw", raw_mode is automatically enabled
+            tokenizer [PRODUCER]: Instantiated tokenizer for tokenized mode.
+                           Required unless raw_mode=True.
+            tokenizer_config [PRODUCER]: Config used to reconstruct the tokenizer
+                           in distributed workers. Required unless raw_mode=True.
             n_workers: [PRODUCER] Number of Modal workers (producer-side parallelism)
             n_scanners: [PRODUCER] Number of scanners per worker (for Mixture of Scanners)
             cpu: [PRODUCER] CPU cores per worker; must match deploy_dataloader_app(cpu=...).
@@ -253,9 +259,6 @@ def __init__(
                                  Chunk size ≤ prefetch_batch_size * prefetch_batch_count rows per partition.
                                  Lower = less memory; higher = fewer process_batch calls.
             batch_size: [CONSUMER] Training batch size (number of samples per batch)
-            max_genes [PRODUCER]: Maximum genes per cell after window function
-            vocab_size [PRODUCER]: Vocabulary size for tokenizer
-            n_expression_bins: Number of expression bins for scGPT
             n_epochs [PRODUCER]: Number of epochs to process
             raw_mode [PRODUCER]: If True, return raw data without tokenization
             return_tensors: [CONSUMER] If True, return torch.Tensor objects (matches SLAFDataLoader).
@@ -279,7 +282,6 @@ def __init__(
         """
         self.slaf_array = slaf_array
         self.batch_size = batch_size
-        self.max_genes = max_genes
         self.n_epochs = n_epochs
         self.raw_mode = raw_mode
         self.return_tensors = return_tensors
@@ -293,41 +295,36 @@ def __init__(
         self.memory = memory
         self.seed = seed
 
-        tokenizer_type = tokenizer_type.lower()
-        if tokenizer_type not in {"geneformer", "scgpt", "raw"}:
-            raise ValueError(
-                f"Unsupported tokenizer_type: {tokenizer_type!r}; expected 'geneformer', 'scgpt', or 'raw'."
-            )
-
-        if tokenizer_type == "raw":
-            self.raw_mode = True
-
-        self.tokenizer_type = "raw" if self.raw_mode else tokenizer_type
-
         window_kwargs = dict(window_kwargs)
-        window_kwargs.setdefault("n_expression_bins", n_expression_bins)
         window_kwargs.setdefault("use_binned_expressions", True)
         if expression_preprocessor is not None:
             window_kwargs["expression_preprocessor"] = expression_preprocessor
+        if tokenizer_config is not None and not isinstance(
+            tokenizer_config, DictConfig
+        ):
+            tokenizer_config = OmegaConf.create(tokenizer_config)
 
-        tokenizer_factory_kwargs: dict[str, Any] | None = None
-        tokenizer_cls: type[GeneformerTokenizer] | type[ScGPTTokenizer] | None = None
         if not self.raw_mode:
-            if tokenizer_type == "geneformer":
-                tokenizer_cls = GeneformerTokenizer
-            else:
-                tokenizer_cls = ScGPTTokenizer
-            tokenizer_factory_kwargs = {"vocab_size": vocab_size}
-            if tokenizer_type == "scgpt":
-                tokenizer_factory_kwargs["n_expression_bins"] = n_expression_bins
-            self.tokenizer = tokenizer_cls(
-                slaf_array=slaf_array,
-                **tokenizer_factory_kwargs,
-            )
+            if tokenizer is None:
+                raise ValueError("tokenizer must be provided unless raw_mode=True.")
+            if tokenizer_config is None:
+                raise ValueError(
+                    "tokenizer_config must be provided unless raw_mode=True.",
+                )
+            self.tokenizer = tokenizer
+            self.tokenizer_type = self.tokenizer.name
+            self.max_genes = self.tokenizer.max_genes
             self.special_tokens = self.tokenizer.special_tokens
+            window_kwargs.setdefault(
+                "n_expression_bins",
+                getattr(self.tokenizer, "n_expression_bins", 10),
+            )
         else:
-            tokenizer_factory_kwargs = None
-            tokenizer_cls = None
+            if tokenizer is not None:
+                raise ValueError("raw_mode=True is incompatible with tokenizer.")
+            self.tokenizer_type = "raw"
+            self.max_genes = 0
+            tokenizer_config = None
             self.tokenizer = None
             self.special_tokens = None
         # Create data source
@@ -352,10 +349,12 @@ def __init__(
         partial_groups_kv_name = f"{queue_name}-partial-groups"
 
         # Prepare configs for workers
-        data_source_config = {
-            "type": "lance",
-            "path": lance_path,
-        }
+        data_source_config = OmegaConf.create(
+            {
+                "type": "lance",
+                "path": lance_path,
+            }
+        )
 
         # Data schema for SLAF (maps generic schema to SLAF column names)
         schema = DataSchema(
@@ -368,42 +367,37 @@ def __init__(
             value_list_key="expr_sequence",
         )
 
-        processor_config = {
-            "schema": {
-                "group_key": schema.group_key,
-                "item_key": schema.item_key,
-                "value_key": schema.value_key,
-                "group_key_out": schema.group_key_out,
-                "item_list_key": schema.item_list_key,
-                "value_list_key": schema.value_list_key,
-            },
-            "window_factory": None,
-            "shuffle_factory": None,
-            "max_items": max_genes,
-            "seed": seed,
-            "n_epochs": n_epochs,
-            "window_kwargs": window_kwargs,
-            "continuity_check": "sequential",
-            "enable_cross_worker_boundary_merging": True,
-            "use_tokenizer_window": not self.raw_mode,
-        }
+        processor_config = OmegaConf.create(
+            {
+                "schema": {
+                    "group_key": schema.group_key,
+                    "item_key": schema.item_key,
+                    "value_key": schema.value_key,
+                    "group_key_out": schema.group_key_out,
+                    "item_list_key": schema.item_list_key,
+                    "value_list_key": schema.value_list_key,
+                },
+                "window_factory": None,
+                "shuffle_factory": None,
+                "max_items": self.max_genes,
+                "seed": seed,
+                "n_epochs": n_epochs,
+                "window_kwargs": window_kwargs,
+                "continuity_check": "sequential",
+                "enable_cross_worker_boundary_merging": True,
+                "use_tokenizer_window": not self.raw_mode,
+            }
+        )
 
         # Create the KV dict to ensure it exists before workers try to access it
-        if processor_config.get("enable_cross_worker_boundary_merging", True):
+        if bool(processor_config.enable_cross_worker_boundary_merging):
             _modal_dict_from_name(
                 partial_groups_kv_name,
                 create_if_missing=True,
                 environment_name=modal_queue_environment,
             )
 
-        if self.tokenizer is not None and tokenizer_cls is not None:
-            processor_config["tokenizer_factory"] = {
-                "module": tokenizer_cls.__module__,
-                "class": tokenizer_cls.__name__,
-                "kwargs": tokenizer_factory_kwargs or {},
-            }
-        else:
-            processor_config["tokenizer_factory"] = None
+        processor_config.tokenizer_config = tokenizer_config
 
         # Spawn workers
         # NOTE: The app must be deployed before spawning workers:
diff --git a/slaf/ml/tokenizers.py b/slaf/ml/tokenizers.py
index c71e940..59ee5a1 100644
--- a/slaf/ml/tokenizers.py
+++ b/slaf/ml/tokenizers.py
@@ -3,8 +3,10 @@
 from typing import Any
 
 import numpy as np
+import polars as pl
 import torch
 
+from slaf.core.tabular_schema import SLAF_LANCE_COO_SCHEMA, DataSchema
 from slaf.core.slaf import SLAFArray
 from slaf.ml.aggregators import GeneformerWindow, ScGPTWindow, Window
 
@@ -75,6 +77,11 @@ def __init__(
         self._build_gene_vocabulary()
         self._setup_special_tokens()
 
+    @property
+    def name(self) -> str:
+        """Stable tokenizer identifier for logging and worker reconstruction."""
+        return self.__class__.__name__
+
     def _build_gene_vocabulary(self):
         """Build gene vocabulary from SLAF var DataFrame or genes Lance table."""
         try:
@@ -237,6 +244,36 @@ def create_window(self) -> Window:
         Create a window function based on the tokenizer type.
         """
 
+    def apply(
+        self,
+        df: pl.DataFrame,
+        schema: DataSchema,
+        max_items: int,
+        **kwargs: Any,
+    ) -> pl.DataFrame:
+        """Group per-cell COO rows into tokenizer-ready sequences."""
+        return self.window.apply(
+            df,
+            schema=schema,
+            max_items=max_items,
+            **kwargs,
+        )
+
+    def tokenize_grouped(
+        self,
+        grouped_df: pl.DataFrame,
+        schema: DataSchema = SLAF_LANCE_COO_SCHEMA,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]:
+        """Tokenize grouped cell sequences emitted by ``apply``."""
+        return self.tokenize(
+            gene_sequences=grouped_df[schema.item_list_key].to_list(),
+            expr_sequences=(
+                grouped_df[schema.value_list_key].to_list()
+                if schema.value_list_key and schema.value_list_key in grouped_df.columns
+                else None
+            ),
+        )
+
     def get_vocab_info(self) -> dict[str, Any]:
         """
         Get vocabulary information for debugging and analysis.
@@ -366,12 +403,93 @@ def __init__(
 
         self.n_expression_bins = n_expression_bins
         super().__init__(
-            slaf_array=slaf_array, vocab_size=vocab_size, max_genes=max_genes
+            slaf_array=slaf_array,
+            vocab_size=vocab_size,
+            max_genes=max_genes,
         )
 
     def create_window(self) -> Window:
         return ScGPTWindow()
 
+    def apply(
+        self,
+        df: pl.DataFrame,
+        schema: DataSchema,
+        max_items: int,
+        **kwargs: Any,
+    ) -> pl.DataFrame:
+        kwargs.setdefault("special_token_offset", 4)
+        kwargs.setdefault("expr_bin_start", self.expr_bin_start)
+        kwargs.setdefault("n_expression_bins", self.n_expression_bins)
+        return super().apply(df, schema=schema, max_items=max_items, **kwargs)
+
+    def tokenize_grouped(
+        self,
+        grouped_df: pl.DataFrame,
+        schema: DataSchema = SLAF_LANCE_COO_SCHEMA,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]:
+        gene_sequences = grouped_df[schema.item_list_key].to_list()
+        expr_sequences = (
+            grouped_df[schema.value_list_key].to_list()
+            if schema.value_list_key and schema.value_list_key in grouped_df.columns
+            else None
+        )
+        if expr_sequences is None:
+            raise ValueError(
+                "scGPT grouped tokenization requires expression token sequences"
+            )
+
+        max_sequence_length = self.max_genes + 2
+        batch_size = len(gene_sequences)
+        gene_token_array = np.full(
+            (batch_size, max_sequence_length),
+            self.special_tokens["PAD"],
+            dtype=np.int64,
+        )
+        value_array = np.full(
+            (batch_size, max_sequence_length),
+            self.special_tokens["PAD"],
+            dtype=np.int64,
+        )
+
+        for i, (genes, exprs) in enumerate(
+            zip(gene_sequences, expr_sequences, strict=False)
+        ):
+            n_pairs = min(len(genes), len(exprs), self.max_genes)
+
+            if n_pairs > 0:
+                gene_ids = np.full(
+                    n_pairs + 2, self.special_tokens["PAD"], dtype=np.int64
+                )
+                value_tokens = np.full(
+                    n_pairs + 2, self.special_tokens["PAD"], dtype=np.int64
+                )
+                gene_ids[0] = self.special_tokens["CLS"]
+                gene_ids[1 : 1 + n_pairs] = np.asarray(genes[:n_pairs], dtype=np.int64)
+                gene_ids[1 + n_pairs] = self.special_tokens["SEP"]
+                value_tokens[1 : 1 + n_pairs] = np.asarray(
+                    exprs[:n_pairs], dtype=np.int64
+                )
+            else:
+                gene_ids = np.array(
+                    [self.special_tokens["CLS"], self.special_tokens["SEP"]],
+                    dtype=np.int64,
+                )
+                value_tokens = np.array(
+                    [self.special_tokens["PAD"], self.special_tokens["PAD"]],
+                    dtype=np.int64,
+                )
+
+            length = min(len(gene_ids), max_sequence_length)
+            gene_token_array[i, :length] = gene_ids[:length]
+            value_array[i, :length] = value_tokens[:length]
+
+        input_ids = torch.from_numpy(gene_token_array)
+        values_tensor = torch.from_numpy(value_array)
+        attention_mask = input_ids != self.special_tokens["PAD"]
+
+        return input_ids, attention_mask, values_tensor
+
     def tokenize(
         self,
         gene_sequences: list[list[int] | list[tuple[int, float]]],
@@ -626,12 +744,65 @@ def __init__(
         max_genes: int = 2048,
     ):
         super().__init__(
-            slaf_array=slaf_array, vocab_size=vocab_size, max_genes=max_genes
+            slaf_array=slaf_array,
+            vocab_size=vocab_size,
+            max_genes=max_genes,
         )
 
     def create_window(self) -> Window:
         return GeneformerWindow()
 
+    def apply(
+        self,
+        df: pl.DataFrame,
+        schema: DataSchema,
+        max_items: int,
+        **kwargs: Any,
+    ) -> pl.DataFrame:
+        kwargs.setdefault("special_token_offset", 4)
+        return self.window.apply(df, schema=schema, max_items=max_items, **kwargs)
+
+    def tokenize_grouped(
+        self,
+        grouped_df: pl.DataFrame,
+        schema: DataSchema = SLAF_LANCE_COO_SCHEMA,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]:
+        gene_sequences = grouped_df[schema.item_list_key].to_list()
+        batch_size = len(gene_sequences)
+        token_array = np.full(
+            (batch_size, self.max_genes), self.special_tokens["PAD"], dtype=np.int64
+        )
+
+        for i, genes in enumerate(gene_sequences):
+            gene_tokens = np.asarray(genes, dtype=np.int64)
+            if len(gene_tokens) > 0:
+                tokens = np.concatenate(
+                    [
+                        [self.special_tokens["CLS"]],
+                        gene_tokens,
+                        [self.special_tokens["SEP"]],
+                    ]
+                )
+            else:
+                tokens = np.array(
+                    [self.special_tokens["CLS"], self.special_tokens["SEP"]],
+                    dtype=np.int64,
+                )
+
+            tokens = tokens[: self.max_genes]
+            if len(tokens) < self.max_genes:
+                padding = np.full(
+                    self.max_genes - len(tokens),
+                    self.special_tokens["PAD"],
+                    dtype=np.int64,
+                )
+                tokens = np.concatenate([tokens, padding])
+            token_array[i, :] = tokens
+
+        input_ids = torch.from_numpy(token_array)
+        attention_mask = input_ids != self.special_tokens["PAD"]
+        return input_ids, attention_mask, None
+
     def tokenize(
         self,
         gene_sequences: list[list[int] | list[tuple[int, float]]],
diff --git a/tests/test_dataloaders.py b/tests/test_dataloaders.py
index 78695af..48045d2 100644
--- a/tests/test_dataloaders.py
+++ b/tests/test_dataloaders.py
@@ -2,13 +2,29 @@
 import pytest
 import torch
 
-from slaf.ml.dataloaders import (
-    GeneformerTokenizer,
-    ScGPTTokenizer,
-    SLAFDataLoader,
-    get_device_info,
-    get_optimal_device,
-)
+from slaf.ml.dataloaders import SLAFDataLoader, get_device_info, get_optimal_device
+from slaf.ml.tokenizers import GeneformerTokenizer, ScGPTTokenizer
+
+
+def build_dataloader(slaf_array, tokenizer_kind="geneformer", raw_mode=False, **kwargs):
+    tokenizer_kwargs = {}
+    for key in ("max_genes", "vocab_size", "n_expression_bins"):
+        if key in kwargs:
+            tokenizer_kwargs[key] = kwargs.pop(key)
+
+    if raw_mode:
+        return SLAFDataLoader(slaf_array, raw_mode=True, **kwargs)
+
+    if tokenizer_kind == "scgpt":
+        tokenizer = ScGPTTokenizer(slaf_array, **tokenizer_kwargs)
+    else:
+        tokenizer = GeneformerTokenizer(slaf_array, **tokenizer_kwargs)
+
+    return SLAFDataLoader(
+        slaf_array,
+        tokenizer=tokenizer,
+        **kwargs,
+    )
 
 
 class TestSLAFDataLoader:
@@ -16,7 +32,7 @@ class TestSLAFDataLoader:
 
     def test_dataloader_initialization(self, tiny_slaf):
         """Test SLAFDataLoader initialization with new architecture"""
-        dataloader = SLAFDataLoader(tiny_slaf)
+        dataloader = build_dataloader(tiny_slaf)
 
         # Check basic attributes
         assert dataloader.slaf_array is tiny_slaf
@@ -34,9 +50,9 @@ def test_dataloader_initialization(self, tiny_slaf):
 
     def test_dataloader_initialization_custom_params(self, tiny_slaf):
         """Test SLAFDataLoader initialization with custom parameters"""
-        dataloader = SLAFDataLoader(
+        dataloader = build_dataloader(
             tiny_slaf,
-            tokenizer_type="scgpt",
+            tokenizer_kind="scgpt",
             batch_size=16,
             max_genes=1024,
             vocab_size=1000,
@@ -50,11 +66,18 @@ def test_dataloader_initialization_custom_params(self, tiny_slaf):
         assert dataloader.tokenizer.vocab_size == 1000
         assert dataloader.tokenizer.n_expression_bins == 5
 
+    def test_dataloader_initialization_with_default_loading_mode(self, tiny_slaf):
+        """Test SLAFDataLoader default loading mode."""
+        dataloader = build_dataloader(tiny_slaf)
+
+        assert dataloader._dataset.batch_processor.use_mixture_of_scanners is True
+        assert dataloader._dataset.batch_processor.by_fragment is True
+
     def test_geneformer_iteration(self, tiny_slaf):
         """Test dataloader iteration with Geneformer tokenizer"""
-        dataloader = SLAFDataLoader(
+        dataloader = build_dataloader(
             tiny_slaf,
-            tokenizer_type="geneformer",
+            tokenizer_kind="geneformer",
             batch_size=5,
             max_genes=10,
         )
@@ -90,9 +113,9 @@ def test_geneformer_iteration(self, tiny_slaf):
 
     def test_scgpt_iteration(self, tiny_slaf):
         """Test dataloader iteration with scGPT tokenizer"""
-        dataloader = SLAFDataLoader(
+        dataloader = build_dataloader(
             tiny_slaf,
-            tokenizer_type="scgpt",
+            tokenizer_kind="scgpt",
             batch_size=5,
             max_genes=10,
         )
@@ -113,7 +136,7 @@ def test_scgpt_iteration(self, tiny_slaf):
 
             assert input_ids.shape[0] == attention_mask.shape[0]
             assert input_ids.shape[0] == cell_ids.shape[0]
-            assert input_ids.shape[1] == 12  # scGPT dual stream: max_genes + 2
+            assert input_ids.shape[1] == 12  # scGPT: max_genes + 2
             assert values.shape == input_ids.shape
 
             # Check data types
@@ -129,7 +152,7 @@ def test_scgpt_iteration(self, tiny_slaf):
 
     def test_consistent_batch_sizes(self, tiny_slaf):
         """Test that batches have consistent sizes"""
-        dataloader = SLAFDataLoader(tiny_slaf, batch_size=8)
+        dataloader = build_dataloader(tiny_slaf, batch_size=8)
 
         batch_sizes = []
         for batch in dataloader:
@@ -146,7 +169,7 @@ def test_consistent_batch_sizes(self, tiny_slaf):
 
     def test_cell_id_mapping(self, tiny_slaf):
         """Test that cell IDs are properly mapped"""
-        dataloader = SLAFDataLoader(tiny_slaf, batch_size=5)
+        dataloader = build_dataloader(tiny_slaf, batch_size=5)
 
         for batch in dataloader:
             cell_ids = batch["cell_ids"]
@@ -162,7 +185,7 @@ def test_cell_id_mapping(self, tiny_slaf):
 
     def test_tokenizer_integration(self, tiny_slaf):
         """Test that tokenizer is properly integrated"""
-        dataloader = SLAFDataLoader(tiny_slaf)
+        dataloader = build_dataloader(tiny_slaf)
 
         # Check that tokenizer has expected attributes
         assert hasattr(dataloader.tokenizer, "gene_vocab")
@@ -174,7 +197,7 @@ def test_tokenizer_integration(self, tiny_slaf):
 
     def test_dataloader_cleanup(self, tiny_slaf):
         """Test dataloader cleanup functionality"""
-        dataloader = SLAFDataLoader(tiny_slaf)
+        dataloader = build_dataloader(tiny_slaf)
 
         # Test that cleanup methods exist and don't crash
         # The dataloader doesn't have a stop_streaming method
@@ -185,7 +208,7 @@ def test_dataloader_cleanup(self, tiny_slaf):
 
     def test_dataloader_device(self, tiny_slaf):
         """Test dataloader device handling"""
-        dataloader = SLAFDataLoader(tiny_slaf)
+        dataloader = build_dataloader(tiny_slaf)
 
         # Get a sample batch
         batch = next(iter(dataloader))
@@ -197,7 +220,7 @@ def test_dataloader_device(self, tiny_slaf):
 
     def test_multi_epoch_initialization(self, tiny_slaf):
         """Test SLAFDataLoader initialization with multi-epoch support"""
-        dataloader = SLAFDataLoader(
+        dataloader = build_dataloader(
             tiny_slaf,
             n_epochs=5,  # Test multi-epoch initialization
         )
@@ -208,7 +231,7 @@ def test_multi_epoch_initialization(self, tiny_slaf):
 
     def test_multi_epoch_iteration(self, tiny_slaf):
         """Test dataloader iteration with multiple epochs"""
-        dataloader = SLAFDataLoader(
+        dataloader = build_dataloader(
             tiny_slaf,
             batch_size=4,  # Small batch size for testing
             n_epochs=3,  # Test with 3 epochs
@@ -239,7 +262,7 @@ def test_multi_epoch_iteration(self, tiny_slaf):
 
     def test_multi_epoch_epoch_progression(self, tiny_slaf):
         """Test that epochs progress correctly in multi-epoch mode"""
-        dataloader = SLAFDataLoader(
+        dataloader = build_dataloader(
             tiny_slaf,
             batch_size=2,  # Very small batch size
             n_epochs=4,  # Test with 4 epochs
@@ -272,7 +295,7 @@ def test_multi_epoch_epoch_progression(self, tiny_slaf):
 
     def test_single_epoch_default_behavior(self, tiny_slaf):
         """Test that single epoch (default) behavior is unchanged"""
-        dataloader = SLAFDataLoader(
+        dataloader = build_dataloader(
             tiny_slaf,
             batch_size=8,
             n_epochs=1,  # Single epoch (default)
@@ -296,9 +319,9 @@ def test_single_epoch_default_behavior(self, tiny_slaf):
     def test_multi_epoch_with_different_tokenizers(self, tiny_slaf):
         """Test multi-epoch functionality with different tokenizer types"""
         # Test with Geneformer - limit epochs and batches for speed
-        dataloader_geneformer = SLAFDataLoader(
+        dataloader_geneformer = build_dataloader(
             tiny_slaf,
-            tokenizer_type="geneformer",
+            tokenizer_kind="geneformer",
             batch_size=4,
             n_epochs=2,
         )
@@ -317,9 +340,9 @@ def test_multi_epoch_with_different_tokenizers(self, tiny_slaf):
         )
 
         # Test with scGPT - limit epochs and batches for speed
-        dataloader_scgpt = SLAFDataLoader(
+        dataloader_scgpt = build_dataloader(
             tiny_slaf,
-            tokenizer_type="scgpt",
+            tokenizer_kind="scgpt",
             batch_size=4,
             n_epochs=2,
         )
@@ -339,7 +362,7 @@ def test_multi_epoch_with_different_tokenizers(self, tiny_slaf):
 
     def test_multi_epoch_completion(self, tiny_slaf):
         """Test that dataloader correctly completes all epochs"""
-        dataloader = SLAFDataLoader(
+        dataloader = build_dataloader(
             tiny_slaf,
             batch_size=2,  # Small batch size to complete quickly
             n_epochs=3,  # Test with 3 epochs
@@ -366,7 +389,7 @@ def test_multi_epoch_completion(self, tiny_slaf):
 
     def test_multi_epoch_parameter_passing(self, tiny_slaf):
         """Test that n_epochs parameter is correctly passed through the hierarchy"""
-        dataloader = SLAFDataLoader(
+        dataloader = build_dataloader(
             tiny_slaf,
             n_epochs=7,  # Test with 7 epochs
         )
@@ -379,9 +402,9 @@ def test_multi_epoch_parameter_passing(self, tiny_slaf):
 
     def test_multi_epoch_with_custom_parameters(self, tiny_slaf):
         """Test multi-epoch functionality with custom parameters"""
-        dataloader = SLAFDataLoader(
+        dataloader = build_dataloader(
             tiny_slaf,
-            tokenizer_type="scgpt",
+            tokenizer_kind="scgpt",
             batch_size=6,
             max_genes=512,
             n_epochs=4,
@@ -479,7 +502,7 @@ def test_dataloader_iteration_batch_mode(self, tiny_slaf):
 
     def test_dataloader_tokenized_mode_fragment(self, tiny_slaf):
         """Test dataloader in tokenized mode with fragment loading."""
-        dataloader = SLAFDataLoader(
+        dataloader = build_dataloader(
             tiny_slaf,
             batch_size=32,
             raw_mode=False,
@@ -501,7 +524,7 @@ def test_dataloader_tokenized_mode_fragment(self, tiny_slaf):
 
     def test_dataloader_tokenized_mode_batch(self, tiny_slaf):
         """Test dataloader in tokenized mode with batch loading."""
-        dataloader = SLAFDataLoader(
+        dataloader = build_dataloader(
             tiny_slaf,
             batch_size=32,
             raw_mode=False,
@@ -523,12 +546,9 @@ def test_dataloader_tokenized_mode_batch(self, tiny_slaf):
 
     def test_dataloader_parameters_consistency(self, tiny_slaf):
         """Test that all parameters are properly passed through."""
-        dataloader = SLAFDataLoader(
+        dataloader = build_dataloader(
             tiny_slaf,
             batch_size=64,
-            max_genes=1024,
-            vocab_size=10000,
-            n_expression_bins=5,
             n_epochs=5,
             raw_mode=True,
             verbose=True,
@@ -538,7 +558,7 @@ def test_dataloader_parameters_consistency(self, tiny_slaf):
         )
 
         assert dataloader.batch_size == 64
-        assert dataloader.max_genes == 1024
+        assert dataloader.max_genes == 0
         assert dataloader.n_epochs == 5
         assert dataloader.raw_mode is True
         assert dataloader.verbose is True
@@ -547,7 +567,7 @@ def test_dataloader_parameters_consistency(self, tiny_slaf):
 
     def test_dataloader_length(self, tiny_slaf):
         """Test that dataloader length returns 0 for streaming datasets."""
-        dataloader = SLAFDataLoader(
+        dataloader = build_dataloader(
             tiny_slaf,
             batch_size=32,
             verbose=False,
@@ -557,8 +577,8 @@ def test_dataloader_length(self, tiny_slaf):
 
     def test_mixture_of_scanners_initialization(self, tiny_slaf):
         """Test SLAFDataLoader initialization with Mixture of Scanners (MoS)"""
-        dataloader = SLAFDataLoader(
-            slaf_array=tiny_slaf,
+        dataloader = build_dataloader(
+            tiny_slaf,
             batch_size=32,
             use_mixture_of_scanners=True,
             n_scanners=8,
@@ -576,8 +596,8 @@ def test_mixture_of_scanners_initialization(self, tiny_slaf):
     def test_mixture_of_scanners_parameter_validation(self, tiny_slaf):
         """Test MoS parameter validation in SLAFDataLoader"""
         # Test valid parameters
-        dataloader = SLAFDataLoader(
-            slaf_array=tiny_slaf,
+        dataloader = build_dataloader(
+            tiny_slaf,
             batch_size=32,
             use_mixture_of_scanners=True,
             n_scanners=16,
@@ -588,7 +608,7 @@ def test_mixture_of_scanners_parameter_validation(self, tiny_slaf):
         # Test invalid n_scanners (too low)
         with pytest.raises(ValueError, match="n_scanners must be at least 1"):
             SLAFDataLoader(
-                slaf_array=tiny_slaf,
+                tiny_slaf,
                 batch_size=32,
                 use_mixture_of_scanners=True,
                 n_scanners=0,
@@ -598,7 +618,7 @@ def test_mixture_of_scanners_parameter_validation(self, tiny_slaf):
         # Test invalid n_scanners (too high)
         with pytest.raises(ValueError, match="n_scanners cannot exceed 100"):
             SLAFDataLoader(
-                slaf_array=tiny_slaf,
+                tiny_slaf,
                 batch_size=32,
                 use_mixture_of_scanners=True,
                 n_scanners=101,
@@ -610,7 +630,7 @@ def test_mixture_of_scanners_parameter_validation(self, tiny_slaf):
             ValueError, match="prefetch_batch_size must be at least 1,000"
         ):
             SLAFDataLoader(
-                slaf_array=tiny_slaf,
+                tiny_slaf,
                 batch_size=32,
                 use_mixture_of_scanners=True,
                 n_scanners=16,
@@ -622,7 +642,7 @@ def test_mixture_of_scanners_parameter_validation(self, tiny_slaf):
             ValueError, match="prefetch_batch_size cannot exceed 10,000,000"
         ):
             SLAFDataLoader(
-                slaf_array=tiny_slaf,
+                tiny_slaf,
                 batch_size=32,
                 use_mixture_of_scanners=True,
                 n_scanners=16,
@@ -631,8 +651,8 @@ def test_mixture_of_scanners_parameter_validation(self, tiny_slaf):
 
     def test_mixture_of_scanners_iteration(self, tiny_slaf):
         """Test that MoS dataloader can iterate through batches"""
-        dataloader = SLAFDataLoader(
-            slaf_array=tiny_slaf,
+        dataloader = build_dataloader(
+            tiny_slaf,
             batch_size=8,
             use_mixture_of_scanners=True,
             n_scanners=4,
@@ -654,7 +674,7 @@ def test_mixture_of_scanners_iteration(self, tiny_slaf):
     def test_mixture_of_scanners_with_raw_mode(self, tiny_slaf):
         """Test MoS functionality with raw mode in SLAFDataLoader"""
         dataloader = SLAFDataLoader(
-            slaf_array=tiny_slaf,
+            tiny_slaf,
             batch_size=32,
             raw_mode=True,
             use_mixture_of_scanners=True,
@@ -678,8 +698,8 @@ def test_mixture_of_scanners_with_raw_mode(self, tiny_slaf):
 
     def test_mixture_of_scanners_with_tokenized_mode(self, tiny_slaf):
         """Test MoS functionality with tokenized mode in SLAFDataLoader"""
-        dataloader = SLAFDataLoader(
-            slaf_array=tiny_slaf,
+        dataloader = build_dataloader(
+            tiny_slaf,
             batch_size=32,
             raw_mode=False,
             use_mixture_of_scanners=True,
@@ -705,8 +725,8 @@ def test_mixture_of_scanners_with_tokenized_mode(self, tiny_slaf):
     def test_mixture_of_scanners_backward_compatibility(self, tiny_slaf):
         """Test that MoS is backward compatible (enabled by default) in SLAFDataLoader"""
         # Default behavior (MoS enabled)
-        dataloader_default = SLAFDataLoader(
-            slaf_array=tiny_slaf,
+        dataloader_default = build_dataloader(
+            tiny_slaf,
             batch_size=32,
         )
 
@@ -716,8 +736,8 @@ def test_mixture_of_scanners_backward_compatibility(self, tiny_slaf):
         )
 
         # Explicitly disable MoS
-        dataloader_disabled = SLAFDataLoader(
-            slaf_array=tiny_slaf,
+        dataloader_disabled = build_dataloader(
+            tiny_slaf,
             batch_size=32,
             use_mixture_of_scanners=False,
         )
@@ -730,8 +750,8 @@ def test_mixture_of_scanners_backward_compatibility(self, tiny_slaf):
 
     def test_mixture_of_scanners_parameter_passing(self, tiny_slaf):
         """Test that MoS parameters are correctly passed through the hierarchy"""
-        dataloader = SLAFDataLoader(
-            slaf_array=tiny_slaf,
+        dataloader = build_dataloader(
+            tiny_slaf,
             batch_size=32,
             use_mixture_of_scanners=True,
             n_scanners=12,
@@ -749,9 +769,9 @@ def test_mixture_of_scanners_parameter_passing(self, tiny_slaf):
 
     def test_mixture_of_scanners_with_custom_parameters(self, tiny_slaf):
         """Test MoS functionality with custom parameters in SLAFDataLoader"""
-        dataloader = SLAFDataLoader(
-            slaf_array=tiny_slaf,
-            tokenizer_type="scgpt",
+        dataloader = build_dataloader(
+            tiny_slaf,
+            tokenizer_kind="scgpt",
             batch_size=16,
             use_mixture_of_scanners=True,
             n_scanners=6,
@@ -764,7 +784,7 @@ def test_mixture_of_scanners_with_custom_parameters(self, tiny_slaf):
         assert dataloader.prefetch_batch_size == 1048576
         assert isinstance(dataloader.tokenizer, ScGPTTokenizer)
         assert dataloader.batch_size == 16
-        assert dataloader.max_genes == 2048
+        assert dataloader.max_genes == dataloader.tokenizer.max_genes
 
         # Test iteration
         batch_count = 0
@@ -780,8 +800,8 @@ def test_mixture_of_scanners_with_custom_parameters(self, tiny_slaf):
 
     def test_mixture_of_scanners_multi_epoch(self, tiny_slaf):
         """Test MoS functionality with multi-epoch training"""
-        dataloader = SLAFDataLoader(
-            slaf_array=tiny_slaf,
+        dataloader = build_dataloader(
+            tiny_slaf,
             batch_size=8,
             n_epochs=3,
             use_mixture_of_scanners=True,
@@ -810,8 +830,8 @@ def test_mixture_of_scanners_multi_epoch(self, tiny_slaf):
 
     def test_mixture_of_scanners_fragment_generators_creation(self, tiny_slaf):
         """Test that MoS creates fragment generators correctly"""
-        dataloader = SLAFDataLoader(
-            slaf_array=tiny_slaf,
+        dataloader = build_dataloader(
+            tiny_slaf,
             batch_size=32,
             use_mixture_of_scanners=True,
             n_scanners=4,
@@ -834,8 +854,8 @@ def test_mixture_of_scanners_fragment_generators_creation(self, tiny_slaf):
 
     def test_mixture_of_scanners_random_sampling_behavior(self, tiny_slaf):
         """Test that MoS uses random sampling from fragment generators"""
-        dataloader = SLAFDataLoader(
-            slaf_array=tiny_slaf,
+        dataloader = build_dataloader(
+            tiny_slaf,
             batch_size=8,
             use_mixture_of_scanners=True,
             n_scanners=2,  # Small number for testing
@@ -859,8 +879,8 @@ def test_mixture_of_scanners_random_sampling_behavior(self, tiny_slaf):
 
     def test_mixture_of_scanners_cell_boundary_handling(self, tiny_slaf):
         """Test that MoS handles cell boundaries correctly in SLAFDataLoader"""
-        dataloader = SLAFDataLoader(
-            slaf_array=tiny_slaf,
+        dataloader = build_dataloader(
+            tiny_slaf,
             batch_size=8,
             use_mixture_of_scanners=True,
             n_scanners=4,
@@ -919,7 +939,7 @@ def test_tensor_on_optimal_device(self):
     )
     def test_dataloader_device(self, tiny_slaf):
         """Test that dataloader uses the correct device"""
-        dataloader = SLAFDataLoader(tiny_slaf)
+        dataloader = build_dataloader(tiny_slaf)
 
         # Check that device is set
         if dataloader.device is not None:
diff --git a/tests/test_distributed_worker.py b/tests/test_distributed_worker.py
index 32e2636..89e673a 100644
--- a/tests/test_distributed_worker.py
+++ b/tests/test_distributed_worker.py
@@ -9,10 +9,12 @@
 from unittest.mock import MagicMock, patch
 
 import polars as pl
+import pytest
 
 from slaf.distributed.data_source import DataSource
 from omegaconf import OmegaConf
 
+
 from slaf.distributed.worker import prefetch_worker
 
 
diff --git a/tests/test_tokenizers.py b/tests/test_tokenizers.py
index 7580ca7..42de839 100644
--- a/tests/test_tokenizers.py
+++ b/tests/test_tokenizers.py
@@ -14,16 +14,20 @@
 from slaf.ml.tokenizers import GeneformerTokenizer, ScGPTTokenizer
 
 
+def build_mock_slaf_array():
+    mock_slaf_array = Mock(spec=SLAFArray)
+    mock_var = Mock()
+    mock_var.index = pd.Index(["gene_0", "gene_1", "gene_2"])
+    mock_slaf_array.var = mock_var
+    return mock_slaf_array
+
+
 class TestSLAFTokenizer:
     """Test the new SLAFTokenizer interface."""
 
     def test_tokenizer_initialization(self):
         """Test SLAFTokenizer initialization with different tokenizer types."""
-        # Mock SLAFArray
-        mock_slaf_array = Mock(spec=SLAFArray)
-        mock_var = Mock()
-        mock_var.index = pd.Index(["gene_0", "gene_1", "gene_2"])
-        mock_slaf_array.var = mock_var
+        mock_slaf_array = build_mock_slaf_array()
 
         # Test Geneformer initialization
         tokenizer = GeneformerTokenizer(
@@ -51,11 +55,7 @@ def test_tokenizer_initialization(self):
 
     def test_geneformer_tokenization(self):
         """Test Geneformer tokenization."""
-        # Mock SLAFArray
-        mock_slaf_array = Mock(spec=SLAFArray)
-        mock_var = Mock()
-        mock_var.index = pd.Index(["gene_0", "gene_1", "gene_2"])
-        mock_slaf_array.var = mock_var
+        mock_slaf_array = build_mock_slaf_array()
 
         tokenizer = GeneformerTokenizer(
             slaf_array=mock_slaf_array,
@@ -83,11 +83,7 @@ def test_geneformer_tokenization(self):
 
     def test_scgpt_tokenization(self):
         """Test scGPT tokenization with expressions."""
-        # Mock SLAFArray
-        mock_slaf_array = Mock(spec=SLAFArray)
-        mock_var = Mock()
-        mock_var.index = pd.Index(["gene_0", "gene_1", "gene_2"])
-        mock_slaf_array.var = mock_var
+        mock_slaf_array = build_mock_slaf_array()
 
         tokenizer = ScGPTTokenizer(
             slaf_array=mock_slaf_array,
@@ -128,13 +124,48 @@ def test_scgpt_tokenization(self):
         assert torch.all(values[cls_positions] == tokenizer.special_tokens["PAD"])
         assert torch.all(values[sep_positions] == tokenizer.special_tokens["PAD"])
 
+    def test_scgpt_tokenize_grouped_uses_preencoded_tokens(self):
+        mock_slaf_array = build_mock_slaf_array()
+
+        tokenizer = ScGPTTokenizer(
+            slaf_array=mock_slaf_array,
+            vocab_size=1000,
+            n_expression_bins=10,
+        )
+
+        grouped_df = pd.DataFrame(
+            {
+                "gene_sequence": [[4, 5, 6]],
+                "expr_sequence": [[1001, 1005, 1009]],
+            }
+        )
+        grouped_df = __import__("polars").from_pandas(grouped_df)
+
+        input_ids, attention_mask, values = tokenizer.tokenize_grouped(grouped_df)
+
+        assert input_ids[0, 1] == 4
+        assert values is not None
+        assert values[0, 1] == 1001
+
+    def test_geneformer_tokenize_grouped_uses_preencoded_tokens(self):
+        mock_slaf_array = build_mock_slaf_array()
+
+        tokenizer = GeneformerTokenizer(
+            slaf_array=mock_slaf_array,
+            vocab_size=1000,
+        )
+
+        grouped_df = pd.DataFrame({"gene_sequence": [[4, 5, 6]]})
+        grouped_df = __import__("polars").from_pandas(grouped_df)
+
+        input_ids, attention_mask, values = tokenizer.tokenize_grouped(grouped_df)
+
+        assert input_ids[0, 1] == 4
+        assert values is None
+
     def test_scgpt_tokenization_no_expression(self):
         """Test that scGPT tokenization works without expressions (empty sequences)."""
-        # Mock SLAFArray
-        mock_slaf_array = Mock(spec=SLAFArray)
-        mock_var = Mock()
-        mock_var.index = pd.Index(["gene_0", "gene_1", "gene_2"])
-        mock_slaf_array.var = mock_var
+        mock_slaf_array = build_mock_slaf_array()
 
         tokenizer = ScGPTTokenizer(
             slaf_array=mock_slaf_array,
@@ -155,11 +186,7 @@ def test_scgpt_tokenization_no_expression(self):
 
     def test_tokenization_edge_cases(self):
         """Test edge cases for tokenization."""
-        # Mock SLAFArray
-        mock_slaf_array = Mock(spec=SLAFArray)
-        mock_var = Mock()
-        mock_var.index = pd.Index(["gene_0", "gene_1", "gene_2"])
-        mock_slaf_array.var = mock_var
+        mock_slaf_array = build_mock_slaf_array()
 
         tokenizer = GeneformerTokenizer(
             slaf_array=mock_slaf_array,
@@ -184,10 +211,7 @@ def test_tokenization_edge_cases(self):
     def test_expression_binning(self):
         """Test expression binning functionality."""
         # Mock SLAFArray
-        mock_slaf_array = Mock(spec=SLAFArray)
-        mock_var = Mock()
-        mock_var.index = pd.Index(["gene_0", "gene_1", "gene_2"])
-        mock_slaf_array.var = mock_var
+        mock_slaf_array = build_mock_slaf_array()
 
         tokenizer = ScGPTTokenizer(
             slaf_array=mock_slaf_array,
@@ -214,11 +238,7 @@ def test_expression_binning(self):
 
     def test_gene_id_mapping(self):
         """Test gene ID to token mapping."""
-        # Mock SLAFArray
-        mock_slaf_array = Mock(spec=SLAFArray)
-        mock_var = Mock()
-        mock_var.index = pd.Index(["gene_0", "gene_1", "gene_2"])
-        mock_slaf_array.var = mock_var
+        mock_slaf_array = build_mock_slaf_array()
 
         tokenizer = GeneformerTokenizer(
             slaf_array=mock_slaf_array,
@@ -246,11 +266,7 @@ def test_gene_id_mapping(self):
 
     def test_vocabulary_info(self):
         """Test vocabulary information retrieval."""
-        # Mock SLAFArray
-        mock_slaf_array = Mock(spec=SLAFArray)
-        mock_var = Mock()
-        mock_var.index = pd.Index(["gene_0", "gene_1", "gene_2"])
-        mock_slaf_array.var = mock_var
+        mock_slaf_array = build_mock_slaf_array()
 
         tokenizer = GeneformerTokenizer(
             slaf_array=mock_slaf_array,
@@ -266,11 +282,7 @@ def test_vocabulary_info(self):
 
     def test_token_decoding(self):
         """Test token decoding functionality."""
-        # Mock SLAFArray
-        mock_slaf_array = Mock(spec=SLAFArray)
-        mock_var = Mock()
-        mock_var.index = pd.Index(["gene_0", "gene_1", "gene_2"])
-        mock_slaf_array.var = mock_var
+        mock_slaf_array = build_mock_slaf_array()
 
         tokenizer = ScGPTTokenizer(
             slaf_array=mock_slaf_array,
diff --git a/uv.lock b/uv.lock
index e4a542d..1b9a19a 100644
--- a/uv.lock
+++ b/uv.lock
@@ -227,6 +227,12 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/78/b6/6307fbef88d9b5ee7421e68d78a9f162e0da4900bc5f5793f6d3d0e34fb8/annotated_types-0.7.0-py3-none-any.whl", hash = "sha256:1f02e8b43a8fbbc3f3e0d4f0f4bfc8131bcb4eebe8849b8e5c773f3a1c582a53", size = 13643, upload-time = "2024-05-20T21:33:24.1Z" },
 ]
 
+[[package]]
+name = "antlr4-python3-runtime"
+version = "4.9.3"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/3e/38/7859ff46355f76f8d19459005ca000b6e7012f2f1ca597746cbcd1fbfe5e/antlr4-python3-runtime-4.9.3.tar.gz", hash = "sha256:f224469b4168294902bb1efa80a8bf7855f24c99aef99cbefc1bcd3cce77881b", size = 117034, upload-time = "2021-11-06T17:52:23.524Z" }
+
 [[package]]
 name = "anyio"
 version = "4.13.0"
@@ -1348,28 +1354,28 @@ wheels = [
 
 [[package]]
 name = "google-auth"
-version = "2.49.2"
+version = "2.52.0"
 source = { registry = "https://pypi.org/simple" }
 dependencies = [
     { name = "cryptography" },
     { name = "pyasn1-modules" },
 ]
-sdist = { url = "https://files.pythonhosted.org/packages/c6/fc/e925290a1ad95c975c459e2df070fac2b90954e13a0370ac505dff78cb99/google_auth-2.49.2.tar.gz", hash = "sha256:c1ae38500e73065dcae57355adb6278cf8b5c8e391994ae9cbadbcb9631ab409", size = 333958, upload-time = "2026-04-10T00:41:21.888Z" }
+sdist = { url = "https://files.pythonhosted.org/packages/d4/f8/80d2493cbedece1c623dc3e3cb1883300871af0dcdae254409522985ac23/google_auth-2.52.0.tar.gz", hash = "sha256:01f30e1a9e3638698d89464f5e603ce29d18e1c0e63ec31ac570aba4e164aaf5", size = 335027, upload-time = "2026-05-07T19:45:24.033Z" }
 wheels = [
-    { url = "https://files.pythonhosted.org/packages/73/76/d241a5c927433420507215df6cac1b1fa4ac0ba7a794df42a84326c68da8/google_auth-2.49.2-py3-none-any.whl", hash = "sha256:c2720924dfc82dedb962c9f52cabb2ab16714fd0a6a707e40561d217574ed6d5", size = 240638, upload-time = "2026-04-10T00:41:14.501Z" },
+    { url = "https://files.pythonhosted.org/packages/ee/fc/2cdc74252746f547f81ff3f02d4d4234a3f411b5de5b61af97e633a060b9/google_auth-2.52.0-py3-none-any.whl", hash = "sha256:aee92803ba0ff93a70a3b8a35c7b4797837751cd6380b63ff38372b98f3ed627", size = 245614, upload-time = "2026-05-07T19:45:21.914Z" },
 ]
 
 [[package]]
 name = "google-cloud-core"
-version = "2.5.1"
+version = "2.6.0"
 source = { registry = "https://pypi.org/simple" }
 dependencies = [
     { name = "google-api-core" },
     { name = "google-auth" },
 ]
-sdist = { url = "https://files.pythonhosted.org/packages/dc/24/6ca08b0a03c7b0c620427503ab00353a4ae806b848b93bcea18b6b76fde6/google_cloud_core-2.5.1.tar.gz", hash = "sha256:3dc94bdec9d05a31d9f355045ed0f369fbc0d8c665076c734f065d729800f811", size = 36078, upload-time = "2026-03-30T22:50:08.057Z" }
+sdist = { url = "https://files.pythonhosted.org/packages/a8/dd/1eef226e470369b26824a505c34482c0b493bc35fe8e0c6b003b5feca21a/google_cloud_core-2.6.0.tar.gz", hash = "sha256:e76149739f90fac1fc6757c09f47eaccb3145b54adbd7759b0f7c4b235f46c83", size = 36001, upload-time = "2026-05-07T08:04:04.124Z" }
 wheels = [
-    { url = "https://files.pythonhosted.org/packages/73/d9/5bb050cb32826466aa9b25f79e2ca2879fe66cb76782d4ed798dd7506151/google_cloud_core-2.5.1-py3-none-any.whl", hash = "sha256:ea62cdf502c20e3e14be8a32c05ed02113d7bef454e40ff3fab6fe1ec9f1f4e7", size = 29452, upload-time = "2026-03-30T22:48:31.567Z" },
+    { url = "https://files.pythonhosted.org/packages/84/4a/98da8930ab109c73d9a5d13782a9ebb81ea8c111f6d534a567b71d23e52b/google_cloud_core-2.6.0-py3-none-any.whl", hash = "sha256:6d63ac8e5eca6d9e4319d0a1e2265fadcd7f1049904378caecfa01cf52dd869e", size = 29390, upload-time = "2026-05-07T08:02:34.672Z" },
 ]
 
 [[package]]
@@ -1421,26 +1427,26 @@ wheels = [
 
 [[package]]
 name = "google-resumable-media"
-version = "2.8.2"
+version = "2.9.0"
 source = { registry = "https://pypi.org/simple" }
 dependencies = [
     { name = "google-crc32c" },
 ]
-sdist = { url = "https://files.pythonhosted.org/packages/3f/d1/b1ea14b93b6b78f57fc580125de44e9f593ab88dd2460f1a8a8d18f74754/google_resumable_media-2.8.2.tar.gz", hash = "sha256:f3354a182ebd193ae3f42e3ef95e6c9b10f128320de23ac7637236713b1acd70", size = 2164510, upload-time = "2026-03-30T23:34:25.369Z" }
+sdist = { url = "https://files.pythonhosted.org/packages/00/4b/0b235beccc310d0a48adbc7246b719d173cca6c88c572dfa4b090e39143c/google_resumable_media-2.9.0.tar.gz", hash = "sha256:f7cfb224846a9dd444d125115dfbe8ef02a2b893e78f087762fe716a255a734b", size = 2164534, upload-time = "2026-05-07T08:04:44.236Z" }
 wheels = [
-    { url = "https://files.pythonhosted.org/packages/5e/f8/50bfaf4658431ff9de45c5c3935af7ab01157a4903c603cd0eee6e78e087/google_resumable_media-2.8.2-py3-none-any.whl", hash = "sha256:82b6d8ccd11765268cdd2a2123f417ec806b8eef3000a9a38dfe3033da5fb220", size = 81511, upload-time = "2026-03-30T23:34:09.671Z" },
+    { url = "https://files.pythonhosted.org/packages/07/73/3518e63deb1667c5409a4579e28daf5e84479a87a72c547e0487f7883dcd/google_resumable_media-2.9.0-py3-none-any.whl", hash = "sha256:c8901e88e389af8bed64d9696c74d8bad961865eb2236e13e0bfca9bb0a65ca3", size = 81507, upload-time = "2026-05-07T08:03:23.809Z" },
 ]
 
 [[package]]
 name = "googleapis-common-protos"
-version = "1.74.0"
+version = "1.75.0"
 source = { registry = "https://pypi.org/simple" }
 dependencies = [
     { name = "protobuf" },
 ]
-sdist = { url = "https://files.pythonhosted.org/packages/20/18/a746c8344152d368a5aac738d4c857012f2c5d1fd2eac7e17b647a7861bd/googleapis_common_protos-1.74.0.tar.gz", hash = "sha256:57971e4eeeba6aad1163c1f0fc88543f965bb49129b8bb55b2b7b26ecab084f1", size = 151254, upload-time = "2026-04-02T21:23:26.679Z" }
+sdist = { url = "https://files.pythonhosted.org/packages/b5/c8/f439cffde755cffa462bfbb156278fa6f9d09119719af9814b858fd4f81f/googleapis_common_protos-1.75.0.tar.gz", hash = "sha256:53a062ff3c32552fbd62c11fe23768b78e4ddf0494d5e5fd97d3f4689c75fbbd", size = 151035, upload-time = "2026-05-07T08:04:49.423Z" }
 wheels = [
-    { url = "https://files.pythonhosted.org/packages/b6/b0/be5d3329badb9230b765de6eea66b73abd5944bdeb5afb3562ddcd80ae84/googleapis_common_protos-1.74.0-py3-none-any.whl", hash = "sha256:702216f78610bb510e3f12ac3cafd281b7ac45cc5d86e90ad87e4d301a3426b5", size = 300743, upload-time = "2026-04-02T21:22:49.108Z" },
+    { url = "https://files.pythonhosted.org/packages/e7/c8/e2645aa8ed02fd4c7a2f59d68783b65b1f3cbdfe39a6308e156509d1fee8/googleapis_common_protos-1.75.0-py3-none-any.whl", hash = "sha256:961ed60399c457ceb0ee8f285a84c870aabc9c6a832b9d37bb281b5bebde43ed", size = 300631, upload-time = "2026-05-07T08:03:30.345Z" },
 ]
 
 [[package]]
@@ -3618,6 +3624,19 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/a8/64/3708a90d1ebe202ffdeb7185f878a3c84d15c2b2c31858da2ce0583e2def/nvidia_nvtx-13.0.85-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:cb7780edb6b14107373c835bf8b72e7a178bac7367e23da7acb108f973f157a6", size = 148878, upload-time = "2025-09-04T08:28:53.627Z" },
 ]
 
+[[package]]
+name = "omegaconf"
+version = "2.3.0"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "antlr4-python3-runtime" },
+    { name = "pyyaml" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/09/48/6388f1bb9da707110532cb70ec4d2822858ddfb44f1cdf1233c20a80ea4b/omegaconf-2.3.0.tar.gz", hash = "sha256:d5d4b6d29955cc50ad50c46dc269bcd92c6e00f5f90d23ab5fee7bfca4ba4cc7", size = 3298120, upload-time = "2022-12-08T20:59:22.753Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/e3/94/1843518e420fa3ed6919835845df698c7e27e183cb997394e4a670973a65/omegaconf-2.3.0-py3-none-any.whl", hash = "sha256:7b4df175cdb08ba400f45cae3bdcae7ba8365db4d165fc65fd04b050ab63b46b", size = 79500, upload-time = "2022-12-08T20:59:19.686Z" },
+]
+
 [[package]]
 name = "opt-einsum"
 version = "3.4.0"
@@ -4033,14 +4052,14 @@ wheels = [
 
 [[package]]
 name = "proto-plus"
-version = "1.27.2"
+version = "1.28.0"
 source = { registry = "https://pypi.org/simple" }
 dependencies = [
     { name = "protobuf" },
 ]
-sdist = { url = "https://files.pythonhosted.org/packages/81/0d/94dfe80193e79d55258345901acd2917523d56e8381bc4dee7fd38e3868a/proto_plus-1.27.2.tar.gz", hash = "sha256:b2adde53adadf75737c44d3dcb0104fde65250dfc83ad59168b4aa3e574b6a24", size = 57204, upload-time = "2026-03-26T22:18:57.174Z" }
+sdist = { url = "https://files.pythonhosted.org/packages/c9/56/e647b0c675392d2da368da7b6f158f7368b18542fd6f7d7400a2f39de000/proto_plus-1.28.0.tar.gz", hash = "sha256:38e5696342835b08fc116f30a25665b29531cda9d5d5643e9b81fc312385abd9", size = 57221, upload-time = "2026-05-07T08:04:50.811Z" }
 wheels = [
-    { url = "https://files.pythonhosted.org/packages/84/f3/1fba73eeffafc998a25d59703b63f8be4fe8a5cb12eaff7386a0ba0f7125/proto_plus-1.27.2-py3-none-any.whl", hash = "sha256:6432f75893d3b9e70b9c412f1d2f03f65b11fb164b793d14ae2ca01821d22718", size = 50450, upload-time = "2026-03-26T22:13:42.927Z" },
+    { url = "https://files.pythonhosted.org/packages/7c/20/b122d4626976acb81132036d2ad1bb35a1a8775fceb837ec30964622516a/proto_plus-1.28.0-py3-none-any.whl", hash = "sha256:a630604310899e73c59ec302e5765c058d412b2f090b9c79c8822589f14955b8", size = 50410, upload-time = "2026-05-07T08:03:31.962Z" },
 ]
 
 [[package]]
@@ -5519,11 +5538,13 @@ docs = [
 full = [
     { name = "igraph" },
     { name = "leidenalg" },
+    { name = "omegaconf" },
     { name = "tiledb" },
     { name = "tiledbsoma" },
     { name = "torch" },
 ]
 ml = [
+    { name = "omegaconf" },
     { name = "tiledb" },
     { name = "tiledbsoma" },
     { name = "torch" },
@@ -5581,6 +5602,7 @@ requires-dist = [
     { name = "modal", marker = "extra == 'dev'", specifier = ">=1.2.1" },
     { name = "mypy", marker = "extra == 'dev'", specifier = ">=1.8.0" },
     { name = "numpy", specifier = ">=1.26.0" },
+    { name = "omegaconf", marker = "extra == 'ml'", specifier = ">=2.3.0" },
     { name = "pandas", specifier = ">=2.1.0,<3" },
     { name = "polars", specifier = ">=1.36.0" },
     { name = "psutil", marker = "extra == 'dev'", specifier = ">=6.0.0" },

From afacb09a986554713a5075c4bc8a6110d13fb6ac Mon Sep 17 00:00:00 2001
From: Shahul Alam <alam.shahul@gmail.com>
Date: Mon, 25 May 2026 00:25:59 -0400
Subject: [PATCH 04/10] Improve tokenizer abstraction.

Remove scGPT-specific tokenizer properties.
---
 slaf/ml/datasets.py    | 11 ++++-------
 slaf/ml/distributed.py |  5 -----
 slaf/ml/tokenizers.py  | 11 ++++++-----
 3 files changed, 10 insertions(+), 17 deletions(-)

diff --git a/slaf/ml/datasets.py b/slaf/ml/datasets.py
index 9600f4a..769134f 100644
--- a/slaf/ml/datasets.py
+++ b/slaf/ml/datasets.py
@@ -57,7 +57,7 @@
 from slaf.core.tabular_schema import SLAF_LANCE_COO_SCHEMA
 from slaf.ml.expression_preprocessor import ExpressionPreprocessor
 from slaf.ml.samplers import Shuffle
-from slaf.ml.tokenizers import SLAFTokenizer
+from slaf.ml.tokenizers import ScGPTTokenizer, SLAFTokenizer
 
 # Define union type for both batch types
 PrefetchBatch = Union["TokenizedPrefetchBatch", "RawPrefetchBatch"]
@@ -1026,10 +1026,7 @@ def load_prefetch_batch(self) -> PrefetchBatch:
 
                     shuffle_time = time.time() - shuffle_start
                     window_start = time.time()
-                    window_params: dict[str, Any] = {
-                        "n_expression_bins": self.n_expression_bins,
-                        "use_binned_expressions": self.use_binned_expressions,
-                    }
+                    window_params = {}
                     window_params.update(
                         self.window_kwargs
                     )  # Add any additional kwargs
@@ -1578,8 +1575,8 @@ def __init__(
             tokenizer, "n_expression_bins", 10
         )  # Default value for raw mode
 
-        # Set binning based on tokenizer type
-        use_binned_expressions = use_binned_expressions  # Use parameter value
+        if isinstance(tokenizer, ScGPTTokenizer):
+            tokenizer.use_binned_expressions = use_binned_expressions
 
         self.batch_processor = PrefetchBatchProcessor(
             slaf_array=slaf_array,
diff --git a/slaf/ml/distributed.py b/slaf/ml/distributed.py
index 4ba2a34..30d09ac 100644
--- a/slaf/ml/distributed.py
+++ b/slaf/ml/distributed.py
@@ -296,7 +296,6 @@ def __init__(
         self.seed = seed
 
         window_kwargs = dict(window_kwargs)
-        window_kwargs.setdefault("use_binned_expressions", True)
         if expression_preprocessor is not None:
             window_kwargs["expression_preprocessor"] = expression_preprocessor
         if tokenizer_config is not None and not isinstance(
@@ -315,10 +314,6 @@ def __init__(
             self.tokenizer_type = self.tokenizer.name
             self.max_genes = self.tokenizer.max_genes
             self.special_tokens = self.tokenizer.special_tokens
-            window_kwargs.setdefault(
-                "n_expression_bins",
-                getattr(self.tokenizer, "n_expression_bins", 10),
-            )
         else:
             if tokenizer is not None:
                 raise ValueError("raw_mode=True is incompatible with tokenizer.")
diff --git a/slaf/ml/tokenizers.py b/slaf/ml/tokenizers.py
index 59ee5a1..5f08e1e 100644
--- a/slaf/ml/tokenizers.py
+++ b/slaf/ml/tokenizers.py
@@ -358,6 +358,7 @@ def __init__(
         vocab_size: int = 50000,
         n_expression_bins: int = 10,
         max_genes: int = 1024,
+        use_binned_expressions: bool = True,
     ):
         """
         Initialize ScGPTTokenizer with SLAF array and vocabulary settings.
@@ -373,6 +374,8 @@ def __init__(
                              Range: 1-1000, default: 10.
             max_genes: Maximum gene--expression pairs per cell. Sequence length is
                        ``2 * max_genes + 2`` (CLS, pairs, SEP).
+            use_binned_expressions: Whether `apply` should emit binned expression
+                       values by default. Set to false to emit raw expression values.
 
         Raises:
             ValueError: If vocab_size is invalid.
@@ -402,11 +405,8 @@ def __init__(
         """
 
         self.n_expression_bins = n_expression_bins
-        super().__init__(
-            slaf_array=slaf_array,
-            vocab_size=vocab_size,
-            max_genes=max_genes,
-        )
+        self.use_binned_expressions = use_binned_expressions
+        super().__init__(slaf_array=slaf_array, vocab_size=vocab_size, max_genes=max_genes)
 
     def create_window(self) -> Window:
         return ScGPTWindow()
@@ -421,6 +421,7 @@ def apply(
         kwargs.setdefault("special_token_offset", 4)
         kwargs.setdefault("expr_bin_start", self.expr_bin_start)
         kwargs.setdefault("n_expression_bins", self.n_expression_bins)
+        kwargs.setdefault("use_binned_expressions", self.use_binned_expressions)
         return super().apply(df, schema=schema, max_items=max_items, **kwargs)
 
     def tokenize_grouped(

From 8eb8fa4e36fd0f0e9b527615170bf5921f65f7fe Mon Sep 17 00:00:00 2001
From: Shahul Alam <alam.shahul@gmail.com>
Date: Thu, 14 May 2026 16:11:41 -0400
Subject: [PATCH 05/10] Fix epoch reporting.

---
 slaf/ml/datasets.py | 8 ++++++--
 1 file changed, 6 insertions(+), 2 deletions(-)

diff --git a/slaf/ml/datasets.py b/slaf/ml/datasets.py
index 769134f..5699c0b 100644
--- a/slaf/ml/datasets.py
+++ b/slaf/ml/datasets.py
@@ -276,7 +276,8 @@ class TokenizedPrefetchBatch:
     """
 
     batch_id: int
-    input_ids: torch.Tensor  # Tokenized sequences
+    epoch: int
+    input_ids: torch.Tensor  # Tokenized identity sequences
     attention_mask: torch.Tensor  # Attention masks
     cell_integer_ids: list[int]  # Corresponding cell integer IDs
     values: torch.Tensor | None = None  # scGPT aligned expression/value stream
@@ -291,6 +292,7 @@ class RawPrefetchBatch:
     """Raw prefetch batch containing pre-chunked raw data for fast batch creation."""
 
     batch_id: int
+    epoch: int
     batch_dfs: list[pl.DataFrame]  # List of pre-chunked DataFrames
     cell_integer_ids: list[int]  # List of all cell IDs across all batches
     process_time: float
@@ -1002,6 +1004,7 @@ def load_prefetch_batch(self) -> PrefetchBatch:
                     self.batch_id += 1  # Increment batch_id for raw mode
                     return RawPrefetchBatch(
                         batch_id=self.batch_id - 1,
+                        epoch=self.current_epoch,
                         batch_dfs=shuffled_chunks,  # type: ignore[arg-type]  # List of pre-chunked DataFrames
                         cell_integer_ids=complete_df["cell_integer_id"]  # type: ignore[index]
                         .unique()
@@ -1089,6 +1092,7 @@ def load_prefetch_batch(self) -> PrefetchBatch:
                     cell_ids_ordered = grouped["cell_integer_id"].to_list()  # type: ignore[index]
                     return TokenizedPrefetchBatch(
                         batch_id=self.batch_id - 1,
+                        epoch=self.current_epoch,
                         input_ids=input_ids,
                         attention_mask=attention_mask,
                         values=values,
@@ -1773,7 +1777,7 @@ def __iter__(self) -> Iterator[dict]:
                         break
 
             # Track epoch transitions
-            current_epoch = self.batch_processor.current_epoch
+            current_epoch = data.epoch
             if current_epoch != last_epoch:
                 print_epoch_transition(
                     f"Epoch transition detected: {last_epoch} -> {current_epoch}",

From 6f727e1e28f0615aa3fb1925833270d79f2df7db Mon Sep 17 00:00:00 2001
From: Shahul Alam <alam.shahul@gmail.com>
Date: Sat, 23 May 2026 13:15:53 -0400
Subject: [PATCH 06/10] Improve logging for normalization.

---
 slaf/integrations/scanpy.py | 15 ++++++++-------
 1 file changed, 8 insertions(+), 7 deletions(-)

diff --git a/slaf/integrations/scanpy.py b/slaf/integrations/scanpy.py
index c672edf..5fee28a 100644
--- a/slaf/integrations/scanpy.py
+++ b/slaf/integrations/scanpy.py
@@ -1,6 +1,7 @@
 import numpy as np
 import pandas as pd
 import polars as pl
+from loguru import logger
 
 from slaf.integrations.anndata import LazyAnnData
 
@@ -464,7 +465,7 @@ def filtered_obs() -> pd.DataFrame:
         if inplace:
             # Apply filter to adata (would need proper implementation)
             # For now, just return the original adata
-            print(
+            logger.info(
                 f"Filtered out {np.sum(~cell_mask)} cells, {np.sum(cell_mask)} remaining"
             )
             return None
@@ -599,7 +600,7 @@ def filter_genes(
 
         if inplace:
             # Apply filter to adata (would need proper implementation)
-            print(
+            logger.info(
                 f"Filtered out {np.sum(~gene_mask)} genes, {np.sum(gene_mask)} remaining"
             )
             return None
@@ -706,7 +707,7 @@ def normalize_total(
                 )
 
             except Exception as e:
-                print(
+                logger.warning(
                     f"Fragment processing failed, falling back to global processing: {e}"
                 )
                 # Fall back to global processing
@@ -807,7 +808,7 @@ def normalize_total(
                 "cell_factors": normalization_dict,
             }
 
-            print(f"Applied normalize_total with target_sum={target_sum}")
+            logger.info(f"Applied normalize_total with target_sum={target_sum}")
             return None
         else:
             # Create a copy with the transformation (copy-on-write)
@@ -899,7 +900,7 @@ def log1p(
                 return adata._update_with_log1p_data(result_df, inplace)
 
             except Exception as e:
-                print(
+                logger.warning(
                     f"Fragment processing failed, falling back to global processing: {e}"
                 )
                 # Fall back to global processing
@@ -914,7 +915,7 @@ def log1p(
 
                 adata._transformations["log1p"] = {"type": "log1p", "applied": True}
 
-                print("Applied log1p transformation")
+                logger.info("Applied log1p transformation")
                 return None
             else:
                 # Create a copy with the transformation (copy-on-write)
@@ -1100,7 +1101,7 @@ def highly_variable_genes(
 
         if inplace:
             # Update var metadata (would need implementation)
-            print(f"Identified {hvg_mask.sum()} highly variable genes")
+            logger.info(f"Identified {hvg_mask.sum()} highly variable genes")
             return None
         else:
             return gene_stats_complete

From cf3cca0b60a5aed458759068ab583b0f01a410ea Mon Sep 17 00:00:00 2001
From: Shahul Alam <alam.shahul@gmail.com>
Date: Fri, 29 May 2026 10:12:43 -0400
Subject: [PATCH 07/10] Fortify test.

---
 tests/test_dataloaders.py | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/tests/test_dataloaders.py b/tests/test_dataloaders.py
index 48045d2..43b24f4 100644
--- a/tests/test_dataloaders.py
+++ b/tests/test_dataloaders.py
@@ -69,9 +69,11 @@ def test_dataloader_initialization_custom_params(self, tiny_slaf):
     def test_dataloader_initialization_with_default_loading_mode(self, tiny_slaf):
         """Test SLAFDataLoader default loading mode."""
         dataloader = build_dataloader(tiny_slaf)
+        batch_processor = dataloader._dataset.batch_processor
 
-        assert dataloader._dataset.batch_processor.use_mixture_of_scanners is True
-        assert dataloader._dataset.batch_processor.by_fragment is True
+        assert batch_processor.use_mixture_of_scanners is True
+        assert batch_processor.by_fragment is True
+        assert batch_processor.expression_dataset is not None
 
     def test_geneformer_iteration(self, tiny_slaf):
         """Test dataloader iteration with Geneformer tokenizer"""

From e4034a0cb6f6929f3987df15f85edb9a1a773990 Mon Sep 17 00:00:00 2001
From: Shahul Alam <alam.shahul@gmail.com>
Date: Sat, 23 May 2026 15:46:42 -0400
Subject: [PATCH 08/10] Fix scGPT expression bins to use PAD-relative IDs

This matches the staged diff: it removes the old vocab_size-offset expression bins, makes expression bins live in 1..n_expression_bins with 0 as PAD, and updates tokenizer/aggregator tests plus dataset test compatibility
---
 slaf/ml/aggregators.py         |  15 ++--
 slaf/ml/tokenizers.py          | 123 +++------------------------------
 tests/test_aggregators.py      |  10 +--
 tests/test_pytorch_datasets.py |  17 +++--
 tests/test_tokenizers.py       |  46 +++++++-----
 5 files changed, 66 insertions(+), 145 deletions(-)

diff --git a/slaf/ml/aggregators.py b/slaf/ml/aggregators.py
index f68cb4d..d7ca6cf 100644
--- a/slaf/ml/aggregators.py
+++ b/slaf/ml/aggregators.py
@@ -109,15 +109,16 @@ def apply(
                     )
                 )
                 .with_columns(
-                    pl.when(pl.col("log_value") > 0)
+                    # TODO: Decide whether scGPT expression bins should use log_value
+                    # or the transformed raw value column for normalization.
+                    pl.when(pl.col(vk) > 0)
                     .then(
-                        (
-                            pl.col("log_value")
-                            * n_expression_bins
-                            / pl.col("log_value").max().over(gk)
+                        1
+                        + (
+                            (pl.col(vk) * n_expression_bins / pl.col(vk).max().over(gk))
+                            .floor()
+                            .clip(0, n_expression_bins - 1)
                         )
-                        .floor()
-                        .clip(0, n_expression_bins - 1)
                     )
                     .otherwise(0)
                     .alias("expr_bin")
diff --git a/slaf/ml/tokenizers.py b/slaf/ml/tokenizers.py
index 5f08e1e..0fcdca4 100644
--- a/slaf/ml/tokenizers.py
+++ b/slaf/ml/tokenizers.py
@@ -6,8 +6,8 @@
 import polars as pl
 import torch
 
-from slaf.core.tabular_schema import SLAF_LANCE_COO_SCHEMA, DataSchema
 from slaf.core.slaf import SLAFArray
+from slaf.core.tabular_schema import SLAF_LANCE_COO_SCHEMA, DataSchema
 from slaf.ml.aggregators import GeneformerWindow, ScGPTWindow, Window
 
 TORCH_AVAILABLE = True
@@ -406,7 +406,9 @@ def __init__(
 
         self.n_expression_bins = n_expression_bins
         self.use_binned_expressions = use_binned_expressions
-        super().__init__(slaf_array=slaf_array, vocab_size=vocab_size, max_genes=max_genes)
+        super().__init__(
+            slaf_array=slaf_array, vocab_size=vocab_size, max_genes=max_genes
+        )
 
     def create_window(self) -> Window:
         return ScGPTWindow()
@@ -418,8 +420,6 @@ def apply(
         max_items: int,
         **kwargs: Any,
     ) -> pl.DataFrame:
-        kwargs.setdefault("special_token_offset", 4)
-        kwargs.setdefault("expr_bin_start", self.expr_bin_start)
         kwargs.setdefault("n_expression_bins", self.n_expression_bins)
         kwargs.setdefault("use_binned_expressions", self.use_binned_expressions)
         return super().apply(df, schema=schema, max_items=max_items, **kwargs)
@@ -429,67 +429,11 @@ def tokenize_grouped(
         grouped_df: pl.DataFrame,
         schema: DataSchema = SLAF_LANCE_COO_SCHEMA,
     ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]:
-        gene_sequences = grouped_df[schema.item_list_key].to_list()
-        expr_sequences = (
-            grouped_df[schema.value_list_key].to_list()
-            if schema.value_list_key and schema.value_list_key in grouped_df.columns
-            else None
-        )
-        if expr_sequences is None:
+        if not schema.value_list_key or schema.value_list_key not in grouped_df.columns:
             raise ValueError(
                 "scGPT grouped tokenization requires expression token sequences"
             )
-
-        max_sequence_length = self.max_genes + 2
-        batch_size = len(gene_sequences)
-        gene_token_array = np.full(
-            (batch_size, max_sequence_length),
-            self.special_tokens["PAD"],
-            dtype=np.int64,
-        )
-        value_array = np.full(
-            (batch_size, max_sequence_length),
-            self.special_tokens["PAD"],
-            dtype=np.int64,
-        )
-
-        for i, (genes, exprs) in enumerate(
-            zip(gene_sequences, expr_sequences, strict=False)
-        ):
-            n_pairs = min(len(genes), len(exprs), self.max_genes)
-
-            if n_pairs > 0:
-                gene_ids = np.full(
-                    n_pairs + 2, self.special_tokens["PAD"], dtype=np.int64
-                )
-                value_tokens = np.full(
-                    n_pairs + 2, self.special_tokens["PAD"], dtype=np.int64
-                )
-                gene_ids[0] = self.special_tokens["CLS"]
-                gene_ids[1 : 1 + n_pairs] = np.asarray(genes[:n_pairs], dtype=np.int64)
-                gene_ids[1 + n_pairs] = self.special_tokens["SEP"]
-                value_tokens[1 : 1 + n_pairs] = np.asarray(
-                    exprs[:n_pairs], dtype=np.int64
-                )
-            else:
-                gene_ids = np.array(
-                    [self.special_tokens["CLS"], self.special_tokens["SEP"]],
-                    dtype=np.int64,
-                )
-                value_tokens = np.array(
-                    [self.special_tokens["PAD"], self.special_tokens["PAD"]],
-                    dtype=np.int64,
-                )
-
-            length = min(len(gene_ids), max_sequence_length)
-            gene_token_array[i, :length] = gene_ids[:length]
-            value_array[i, :length] = value_tokens[:length]
-
-        input_ids = torch.from_numpy(gene_token_array)
-        values_tensor = torch.from_numpy(value_array)
-        attention_mask = input_ids != self.special_tokens["PAD"]
-
-        return input_ids, attention_mask, values_tensor
+        return super().tokenize_grouped(grouped_df, schema=schema)
 
     def tokenize(
         self,
@@ -558,9 +502,7 @@ def tokenize(
                 gene_tokens = np.array(genes[:n_pairs], dtype=np.int64) + 4
 
                 if isinstance(exprs[0], int | np.integer):
-                    expr_tokens = (
-                        np.array(exprs[:n_pairs], dtype=np.int64) + self.expr_bin_start
-                    )
+                    expr_tokens = np.array(exprs[:n_pairs], dtype=np.int64)
                 else:
                     expr_tokens = self._expression_to_bin_vectorized(
                         np.array(exprs[:n_pairs], dtype=np.float32)
@@ -623,7 +565,6 @@ def _setup_special_tokens(self):
         super()._setup_special_tokens()
 
         # Expression binning setup for scGPT
-        self.expr_bin_start = self.vocab_size
         self.expr_bin_size = 1.0 / self.n_expression_bins
 
     def _expression_to_bin(self, expression_value: float) -> int:
@@ -635,7 +576,7 @@ def _expression_to_bin(self, expression_value: float) -> int:
         bin_id = min(
             int(expression_value / self.expr_bin_size), self.n_expression_bins - 1
         )
-        return self.expr_bin_start + bin_id
+        return 1 + bin_id
 
     def _expression_to_bin_vectorized(
         self, expression_values: np.ndarray
@@ -655,7 +596,7 @@ def _expression_to_bin_vectorized(
         # Convert to token IDs
         result = np.where(
             expression_values > 0,
-            self.expr_bin_start + bins,
+            1 + bins,
             self.special_tokens["PAD"],
         )
 
@@ -699,8 +640,8 @@ def decode_tokens(self, tokens: list[int]) -> dict[str, Any]:
                 special_tokens.append("PAD")
             elif token == self.special_tokens["MASK"]:
                 special_tokens.append("MASK")
-            elif token >= self.expr_bin_start:  # Expression token
-                bin_id = token - self.expr_bin_start
+            elif 1 <= token <= self.n_expression_bins:  # Expression bin
+                bin_id = token - 1
                 expr_value = bin_id * self.expr_bin_size
                 expressions.append(expr_value)
             else:
@@ -760,50 +701,8 @@ def apply(
         max_items: int,
         **kwargs: Any,
     ) -> pl.DataFrame:
-        kwargs.setdefault("special_token_offset", 4)
         return self.window.apply(df, schema=schema, max_items=max_items, **kwargs)
 
-    def tokenize_grouped(
-        self,
-        grouped_df: pl.DataFrame,
-        schema: DataSchema = SLAF_LANCE_COO_SCHEMA,
-    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]:
-        gene_sequences = grouped_df[schema.item_list_key].to_list()
-        batch_size = len(gene_sequences)
-        token_array = np.full(
-            (batch_size, self.max_genes), self.special_tokens["PAD"], dtype=np.int64
-        )
-
-        for i, genes in enumerate(gene_sequences):
-            gene_tokens = np.asarray(genes, dtype=np.int64)
-            if len(gene_tokens) > 0:
-                tokens = np.concatenate(
-                    [
-                        [self.special_tokens["CLS"]],
-                        gene_tokens,
-                        [self.special_tokens["SEP"]],
-                    ]
-                )
-            else:
-                tokens = np.array(
-                    [self.special_tokens["CLS"], self.special_tokens["SEP"]],
-                    dtype=np.int64,
-                )
-
-            tokens = tokens[: self.max_genes]
-            if len(tokens) < self.max_genes:
-                padding = np.full(
-                    self.max_genes - len(tokens),
-                    self.special_tokens["PAD"],
-                    dtype=np.int64,
-                )
-                tokens = np.concatenate([tokens, padding])
-            token_array[i, :] = tokens
-
-        input_ids = torch.from_numpy(token_array)
-        attention_mask = input_ids != self.special_tokens["PAD"]
-        return input_ids, attention_mask, None
-
     def tokenize(
         self,
         gene_sequences: list[list[int] | list[tuple[int, float]]],
diff --git a/tests/test_aggregators.py b/tests/test_aggregators.py
index 333072e..869a4a0 100644
--- a/tests/test_aggregators.py
+++ b/tests/test_aggregators.py
@@ -79,9 +79,11 @@ def test_apply_ranking(self):
 
         # Check that highest expression gene comes first
         assert gene_seq[0] == 10  # gene with value 5.0
+        assert gene_seq.to_list() == [10, 20, 30]
         # Check that expression bins are calculated (default behavior)
         assert len(expr_seq) == len(gene_seq)
-        assert all(0 <= bin_val < 10 for bin_val in expr_seq)  # Default 10 bins
+        assert all(0 <= bin_val <= 10 for bin_val in expr_seq)  # PAD + 10 bins
+        assert all(1 <= bin_val <= 10 for bin_val in expr_seq)
 
     def test_apply_expression_binning(self):
         """Test that expression binning works correctly"""
@@ -91,7 +93,7 @@ def test_apply_expression_binning(self):
         expr_seq = cell_0_data["expr_sequence"][0]
 
         # Check that expression bins are in the correct range
-        assert all(0 <= bin_val < 5 for bin_val in expr_seq)  # 5 bins
+        assert all(0 <= bin_val <= 5 for bin_val in expr_seq)  # PAD + 5 bins
 
         # Check that zero values get bin 0
         # Values: [5.0, 3.0, 1.0] -> log(1+value): [1.79, 1.39, 0.69]
@@ -122,7 +124,7 @@ def test_apply_custom_binning_parameters(self):
         expr_seq = cell_0_data["expr_sequence"][0]
 
         # Check that expression bins are in the correct range for 3 bins
-        assert all(0 <= bin_val < 3 for bin_val in expr_seq)
+        assert all(0 <= bin_val <= 3 for bin_val in expr_seq)
 
     def test_apply_max_genes_limit(self):
         """Test that max_genes limit is respected"""
@@ -167,7 +169,7 @@ def test_apply_single_cell(self):
         assert len(expr_seq) == 2
 
         # Check that expression bins are calculated (default behavior)
-        assert all(0 <= bin_val < 10 for bin_val in expr_seq)  # Default 10 bins
+        assert all(0 <= bin_val <= 10 for bin_val in expr_seq)  # PAD + 10 bins
 
         # Based on the actual output, the ranking is [10, 30] with values [3.0, 5.0]
         # This suggests the ranking might be by gene_integer_id when expression values are tied
diff --git a/tests/test_pytorch_datasets.py b/tests/test_pytorch_datasets.py
index 6244b96..980ac95 100644
--- a/tests/test_pytorch_datasets.py
+++ b/tests/test_pytorch_datasets.py
@@ -201,6 +201,7 @@ def test_prefetch_batch_dataclass(self):
 
         batch = TokenizedPrefetchBatch(
             batch_id=0,
+            epoch=0,
             input_ids=input_ids,
             attention_mask=attention_mask,
             cell_integer_ids=[100, 101],
@@ -222,6 +223,7 @@ def test_prefetch_batch_geneformer(self):
 
         batch = TokenizedPrefetchBatch(
             batch_id=1,
+            epoch=0,
             input_ids=input_ids,
             attention_mask=attention_mask,
             cell_integer_ids=[200, 201, 202],
@@ -384,13 +386,19 @@ def test_scgpt_values_alignment_contract(self, tiny_slaf):
         batch = next(iter(dataset))
         input_ids = batch["input_ids"]
         values = batch["values"]
+        attention_mask = batch["attention_mask"]
 
         assert values.shape == input_ids.shape
-        cls_positions = input_ids == tokenizer.special_tokens["CLS"]
-        sep_positions = input_ids == tokenizer.special_tokens["SEP"]
         pad_value = tokenizer.special_tokens["PAD"]
-        assert torch.all(values[cls_positions] == pad_value)
-        assert torch.all(values[sep_positions] == pad_value)
+        assert torch.all(input_ids[:, 0] == tokenizer.special_tokens["CLS"])
+        assert torch.all(values[:, 0] == pad_value)
+
+        sep_positions = attention_mask.long().sum(dim=1) - 1
+        row_indices = torch.arange(input_ids.shape[0])
+        assert torch.all(
+            input_ids[row_indices, sep_positions] == tokenizer.special_tokens["SEP"]
+        )
+        assert torch.all(values[row_indices, sep_positions] == pad_value)
 
     def test_device_transfer(self, tiny_slaf):
         """Test device transfer functionality"""
@@ -768,6 +776,7 @@ def test_prefetch_batch_serialization(self):
 
         batch = TokenizedPrefetchBatch(
             batch_id=0,
+            epoch=0,
             input_ids=input_ids,
             attention_mask=attention_mask,
             cell_integer_ids=[100, 101],
diff --git a/tests/test_tokenizers.py b/tests/test_tokenizers.py
index 42de839..6df000a 100644
--- a/tests/test_tokenizers.py
+++ b/tests/test_tokenizers.py
@@ -50,7 +50,6 @@ def test_tokenizer_initialization(self):
         )
 
         assert tokenizer.n_expression_bins == 5
-        assert tokenizer.expr_bin_start == 1000
         assert tokenizer.expr_bin_size == 0.2
 
     def test_geneformer_tokenization(self):
@@ -124,7 +123,7 @@ def test_scgpt_tokenization(self):
         assert torch.all(values[cls_positions] == tokenizer.special_tokens["PAD"])
         assert torch.all(values[sep_positions] == tokenizer.special_tokens["PAD"])
 
-    def test_scgpt_tokenize_grouped_uses_preencoded_tokens(self):
+    def test_scgpt_tokenize_grouped_matches_tokenize(self):
         mock_slaf_array = build_mock_slaf_array()
 
         tokenizer = ScGPTTokenizer(
@@ -133,21 +132,28 @@ def test_scgpt_tokenize_grouped_uses_preencoded_tokens(self):
             n_expression_bins=10,
         )
 
+        gene_sequences = [[0, 1, 2]]
+        expr_sequences = [[1, 5, 9]]
         grouped_df = pd.DataFrame(
             {
-                "gene_sequence": [[4, 5, 6]],
-                "expr_sequence": [[1001, 1005, 1009]],
+                "gene_sequence": gene_sequences,
+                "expr_sequence": expr_sequences,
             }
         )
         grouped_df = __import__("polars").from_pandas(grouped_df)
 
         input_ids, attention_mask, values = tokenizer.tokenize_grouped(grouped_df)
+        expected_input_ids, expected_attention_mask, expected_values = (
+            tokenizer.tokenize(gene_sequences, expr_sequences)
+        )
 
-        assert input_ids[0, 1] == 4
+        assert torch.equal(input_ids, expected_input_ids)
+        assert torch.equal(attention_mask, expected_attention_mask)
         assert values is not None
-        assert values[0, 1] == 1001
+        assert expected_values is not None
+        assert torch.equal(values, expected_values)
 
-    def test_geneformer_tokenize_grouped_uses_preencoded_tokens(self):
+    def test_geneformer_tokenize_grouped_matches_tokenize(self):
         mock_slaf_array = build_mock_slaf_array()
 
         tokenizer = GeneformerTokenizer(
@@ -155,13 +161,19 @@ def test_geneformer_tokenize_grouped_uses_preencoded_tokens(self):
             vocab_size=1000,
         )
 
-        grouped_df = pd.DataFrame({"gene_sequence": [[4, 5, 6]]})
+        gene_sequences = [[0, 1, 2]]
+        grouped_df = pd.DataFrame({"gene_sequence": gene_sequences})
         grouped_df = __import__("polars").from_pandas(grouped_df)
 
         input_ids, attention_mask, values = tokenizer.tokenize_grouped(grouped_df)
+        expected_input_ids, expected_attention_mask, expected_values = (
+            tokenizer.tokenize(gene_sequences)
+        )
 
-        assert input_ids[0, 1] == 4
+        assert torch.equal(input_ids, expected_input_ids)
+        assert torch.equal(attention_mask, expected_attention_mask)
         assert values is None
+        assert expected_values is None
 
     def test_scgpt_tokenization_no_expression(self):
         """Test that scGPT tokenization works without expressions (empty sequences)."""
@@ -222,18 +234,18 @@ def test_expression_binning(self):
         # Test individual expression binning
         assert tokenizer._expression_to_bin(0.0) == 0  # PAD for zero
         assert tokenizer._expression_to_bin(-1.0) == 0  # PAD for negative
-        assert tokenizer._expression_to_bin(0.1) == 1001  # First bin
-        assert tokenizer._expression_to_bin(0.9) == 1009  # Last bin
-        assert tokenizer._expression_to_bin(1.0) == 1009  # Clipped to last bin
+        assert tokenizer._expression_to_bin(0.1) == 2
+        assert tokenizer._expression_to_bin(0.9) == 10
+        assert tokenizer._expression_to_bin(1.0) == 10  # Clipped to last bin
 
         # Test vectorized expression binning
         expr_values = np.array([0.0, 0.1, 0.5, 0.9, -1.0])
         bins = tokenizer._expression_to_bin_vectorized(expr_values)
 
         assert bins[0] == 0  # PAD for 0.0
-        assert bins[1] == 1001  # First bin for 0.1
-        assert bins[2] == 1005  # Fifth bin for 0.5
-        assert bins[3] == 1009  # Last bin for 0.9
+        assert bins[1] == 2
+        assert bins[2] == 6
+        assert bins[3] == 10
         assert bins[4] == 0  # PAD for -1.0
 
     def test_gene_id_mapping(self):
@@ -307,9 +319,7 @@ def test_token_decoding(self):
         assert "PAD" in decoded["special_tokens"]
 
         # Test decoding scGPT tokens
-        tokens = (
-            [1] + gene_tokens + [1001, 1005] + [2, 0]
-        )  # CLS, gene1, gene2, expr1, expr2, SEP, PAD
+        tokens = [1] + gene_tokens + [2, 0]  # CLS, gene1, gene2, SEP, PAD
         decoded = tokenizer.decode_tokens(tokens)
 
         # Check structure

From c337a7d6ffff43b94f12abc8645eec767d669fab Mon Sep 17 00:00:00 2001
From: Shahul Alam <alam.shahul@gmail.com>
Date: Fri, 29 May 2026 11:33:53 -0400
Subject: [PATCH 09/10] Style fix.

---
 slaf/distributed/worker.py       | 4 +++-
 tests/test_distributed_worker.py | 5 +----
 2 files changed, 4 insertions(+), 5 deletions(-)

diff --git a/slaf/distributed/worker.py b/slaf/distributed/worker.py
index d340d04..9524ab2 100644
--- a/slaf/distributed/worker.py
+++ b/slaf/distributed/worker.py
@@ -122,7 +122,9 @@ def prefetch_worker(
                 if "args" in tokenizer_config and tokenizer_config.args is not None
                 else {}
             )
-            tokenizer_instance = tokenizer_class( slaf_array=slaf_array, **tokenizer_kwargs)
+            tokenizer_instance = tokenizer_class(
+                slaf_array=slaf_array, **tokenizer_kwargs
+            )
 
             def tokenize_grouped(
                 grouped_df: pl.DataFrame, schema: DataSchema
diff --git a/tests/test_distributed_worker.py b/tests/test_distributed_worker.py
index 89e673a..4dd55ba 100644
--- a/tests/test_distributed_worker.py
+++ b/tests/test_distributed_worker.py
@@ -9,12 +9,9 @@
 from unittest.mock import MagicMock, patch
 
 import polars as pl
-import pytest
-
-from slaf.distributed.data_source import DataSource
 from omegaconf import OmegaConf
 
-
+from slaf.distributed.data_source import DataSource
 from slaf.distributed.worker import prefetch_worker
 
 

From 7282abbf1510bd8fb1eebe24804c788173eb0014 Mon Sep 17 00:00:00 2001
From: Shahul Alam <alam.shahul@gmail.com>
Date: Fri, 29 May 2026 11:39:35 -0400
Subject: [PATCH 10/10] Add OmegaConf to `test` extra group.

---
 pyproject.toml | 1 +
 uv.lock        | 2 ++
 2 files changed, 3 insertions(+)

diff --git a/pyproject.toml b/pyproject.toml
index 7740fd1..21a51e9 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -104,6 +104,7 @@ test = [
     "psutil>=6.0.0",
     "torch>=2.5.0",
     "tiledbsoma>=1.17.1",
+    "omegaconf>=2.3.0",
 ]
 
 # Full installation includes all optional dependencies
diff --git a/uv.lock b/uv.lock
index 1b9a19a..a25a2a1 100644
--- a/uv.lock
+++ b/uv.lock
@@ -5552,6 +5552,7 @@ ml = [
 test = [
     { name = "coverage" },
     { name = "h5py" },
+    { name = "omegaconf" },
     { name = "psutil" },
     { name = "pytest" },
     { name = "pytest-cov" },
@@ -5603,6 +5604,7 @@ requires-dist = [
     { name = "mypy", marker = "extra == 'dev'", specifier = ">=1.8.0" },
     { name = "numpy", specifier = ">=1.26.0" },
     { name = "omegaconf", marker = "extra == 'ml'", specifier = ">=2.3.0" },
+    { name = "omegaconf", marker = "extra == 'test'", specifier = ">=2.3.0" },
     { name = "pandas", specifier = ">=2.1.0,<3" },
     { name = "polars", specifier = ">=1.36.0" },
     { name = "psutil", marker = "extra == 'dev'", specifier = ">=6.0.0" },