diff --git a/slaf/distributed/worker.py b/slaf/distributed/worker.py index 74b012d..2cd488d 100644 --- a/slaf/distributed/worker.py +++ b/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 = { diff --git a/slaf/ml/datasets.py b/slaf/ml/datasets.py index 857e54b..56034e0 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() @@ -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}", diff --git a/slaf/ml/tokenizers.py b/slaf/ml/tokenizers.py index c71e940..25a1903 100644 --- a/slaf/ml/tokenizers.py +++ b/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. diff --git a/tests/test_pytorch_datasets.py b/tests/test_pytorch_datasets.py index 6244b96..5628280 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], @@ -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], diff --git a/tests/test_tokenizers.py b/tests/test_tokenizers.py index 7580ca7..8e5c786 100644 --- a/tests/test_tokenizers.py +++ b/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