Refactor SLAF tokenizer interfaces

slaf/distributed/worker.py

@@ -124,43 +124,27 @@ def prefetch_worker(
             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
-                gene_sequences = grouped_df[schema.item_list_key].to_list()
+                """Tokenize grouped DataFrame via tokenizer-owned grouping contract."""
                 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
+                if is_scgpt_tokenizer and (
+                    not schema.value_list_key
+                    or schema.value_list_key not 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
+                    raise ValueError(
+                        "scGPT distributed tokenization requires expression/value sequences; "
+                        f"missing grouped column '{schema.value_list_key}'."
                     )
 
-                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."
-                        )
+                input_ids, attention_mask, values = tokenizer_instance.tokenize_grouped(
+                    grouped_df,
+                    schema=schema,
+                )
+
+                if is_scgpt_tokenizer and values is None:
+                    raise ValueError(
+                        "scGPT distributed tokenization requires dual-stream output; "
+                        "tokenizer returned values=None."
+                    )
 
                 # Return as dict (format expected by processor)
                 result = {

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()
@@ -1056,13 +1059,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
@@ -1096,6 +1095,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,
@@ -1780,7 +1780,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}",

slaf/ml/tokenizers.py

@@ -3,9 +3,11 @@
 from typing import Any
 
 import numpy as np
+import polars as pl
 import torch
 
 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
@@ -237,6 +239,22 @@ def create_window(self) -> Window:
         Create a window function based on the tokenizer type.
         """
 
+    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 ``window.apply``."""
+        expr_key = schema.value_list_key
+        return self.tokenize(
+            gene_sequences=grouped_df[schema.item_list_key].to_list(),
+            expr_sequences=(
+                grouped_df[expr_key].to_list()
+                if expr_key and expr_key in grouped_df.columns
+                else None
+            ),
+        )
+
     def get_vocab_info(self) -> dict[str, Any]:
         """
         Get vocabulary information for debugging and analysis.

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],
@@ -768,6 +770,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],

tests/test_tokenizers.py

@@ -7,6 +7,7 @@
 
 import numpy as np
 import pandas as pd
+import polars as pl
 import pytest
 import torch
 
@@ -128,6 +129,53 @@ 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(self):
+        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
+
+        tokenizer = ScGPTTokenizer(
+            slaf_array=mock_slaf_array,
+            vocab_size=1000,
+            n_expression_bins=10,
+        )
+
+        grouped_df = pl.from_pandas(
+            pd.DataFrame(
+                {
+                    "gene_sequence": [[0, 1, 2]],
+                    "expr_sequence": [[0.5, 0.8, 0.2]],
+                }
+            )
+        )
+
+        input_ids, attention_mask, values = tokenizer.tokenize_grouped(grouped_df)
+
+        assert input_ids.shape == (1, 1026)
+        assert attention_mask.shape == (1, 1026)
+        assert values is not None
+        assert input_ids[0, 0] == tokenizer.special_tokens["CLS"]
+
+    def test_geneformer_tokenize_grouped(self):
+        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
+
+        tokenizer = GeneformerTokenizer(
+            slaf_array=mock_slaf_array,
+            vocab_size=1000,
+        )
+
+        grouped_df = pl.from_pandas(pd.DataFrame({"gene_sequence": [[0, 1, 2]]}))
+
+        input_ids, attention_mask, values = tokenizer.tokenize_grouped(grouped_df)
+
+        assert input_ids.shape[0] == 1
+        assert values is None
+        assert input_ids[0, 0] == tokenizer.special_tokens["CLS"]
+
     def test_scgpt_tokenization_no_expression(self):
         """Test that scGPT tokenization works without expressions (empty sequences)."""
         # Mock SLAFArray