From 52da8628659d2063ecbe007231636a0b5c521d08 Mon Sep 17 00:00:00 2001 From: Shahul Alam Date: Mon, 20 Apr 2026 16:43:48 -0400 Subject: [PATCH 01/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 | 172 +++++++++---------- slaf/ml/dataloaders.py | 149 ++++++++--------- slaf/ml/datasets.py | 16 +- slaf/ml/distributed.py | 145 ++++++++-------- slaf/ml/tokenizers.py | 243 ++++++++++++++++++++++++--- tests/test_dataloaders.py | 279 +++++++++++++++++-------------- tests/test_distributed_worker.py | 4 + tests/test_tokenizers.py | 210 ++++++++++++++++------- uv.lock | 52 ++++-- 10 files changed, 792 insertions(+), 485 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 74b012d..b298aab 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 @@ -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 | dict[str, Any], + processor_config: DictConfig | dict[str, Any], queue: Any, # Modal Queue or any queue-like object n_scanners: int = 8, prefetch_batch_count: int = 32, @@ -56,23 +56,47 @@ def prefetch_worker( Returns: Dictionary with worker metrics """ + 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) + + tokenizer_config = OmegaConf.select(processor_config, "tokenizer_config") + shuffle_factory_config = OmegaConf.select(processor_config, "shuffle_factory") + window_factory_config = OmegaConf.select(processor_config, "window_factory") + use_tokenizer_window = bool( + OmegaConf.select(processor_config, "use_tokenizer_window", default=False) + ) + continuity_check = str( + OmegaConf.select( + processor_config, + "continuity_check", + default="sequential", + ) + ) + max_items = int(OmegaConf.select(processor_config, "max_items", default=1024)) + seed_value = int(OmegaConf.select(processor_config, "seed", default=42)) + window_kwargs = ( + OmegaConf.select(processor_config, "window_kwargs", default={}) or {} + ) + # 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 @@ -96,71 +120,39 @@ def prefetch_worker( # 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" + 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", "") - # Recreate SLAFArray in worker from slaf.core.slaf import SLAFArray + from slaf.integrations.anndata import LazyAnnData slaf_array = SLAFArray(slaf_path, load_metadata=False) - - # Create tokenizer instance - tokenizer_instance = tokenizer_class( - slaf_array=slaf_array, **tokenizer_config["kwargs"] + adata = LazyAnnData(slaf_array) + + 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(adata=adata, **tokenizer_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 - 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 - ): - 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 - ) - - 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." - ) + """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 = { @@ -174,56 +166,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 - 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"): + 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) @@ -231,7 +213,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, ) @@ -241,10 +223,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..e5cdb7a 100644 --- a/slaf/ml/dataloaders.py +++ b/slaf/ml/dataloaders.py @@ -2,7 +2,7 @@ from loguru import logger -from slaf.core.slaf import SLAFArray +from slaf.integrations.anndata import LazyAnnData from .expression_preprocessor import ExpressionPreprocessor from .tokenizers import GeneformerTokenizer, ScGPTTokenizer, SLAFTokenizer @@ -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 @@ -268,12 +269,9 @@ class SLAFDataLoader: def __init__( self, - slaf_array: SLAFArray, - tokenizer_type: str = "geneformer", + adata: LazyAnnData, + 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 is missing in tokenized mode 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}") @@ -419,9 +406,9 @@ def __init__( ... print(f"Error: {e}") Error: n_scanners must be at least 1 """ - self.slaf_array = slaf_array + self.adata = adata + self.slaf_array = adata.slaf 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 @@ -462,48 +449,38 @@ 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 + + self._dataset = self._create_dataset(prefetcher_ready_timeout) - # Use IterableDataset - self._dataset = SLAFIterableDataset( - slaf_array=slaf_array, + def _create_dataset(self, prefetcher_ready_timeout: float) -> "SLAFIterableDataset": + return SLAFIterableDataset( + 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 +502,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 +517,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 +564,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 +580,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: @@ -670,7 +660,8 @@ def __del__(self): >>> print("Manual cleanup completed") Manual cleanup completed """ - if hasattr(self, "_dataset"): - # The SLAFIterableDataset doesn't have a stop method, - # so we just let it finish its current epoch. - pass + self.close() + + def close(self): + if hasattr(self, "_dataset") and hasattr(self._dataset, "close"): + self._dataset.close() 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 accec75..e609fcd 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 @@ -95,8 +96,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 | dict[str, Any], + processor_config: DictConfig | dict[str, Any], queue_name: str, n_scanners: int = 8, prefetch_batch_count: int = 32, @@ -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,44 +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, - "max_genes": max_genes, - } - 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 @@ -355,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( @@ -371,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..c301457 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.integrations.anndata import LazyAnnData from slaf.ml.aggregators import GeneformerWindow, ScGPTWindow, Window TORCH_AVAILABLE = True @@ -41,7 +43,7 @@ class SLAFTokenizer(ABC): def __init__( self, - slaf_array: SLAFArray, + adata: LazyAnnData, vocab_size: int = 50000, max_genes: int = 2048, ): @@ -49,9 +51,9 @@ def __init__( Initialize SLAFTokenizer with SLAF array and vocabulary settings. Args: - slaf_array: Initialized SLAFArray instance containing the single-cell data. - Used to build the gene vocabulary and access expression data. - Must be a valid SLAFArray with proper var DataFrame. + adata: LazyAnnData instance containing the single-cell data. + Used to build the gene vocabulary, access expression metadata, + and inspect lazy runtime transformations. vocab_size: Maximum size of gene vocabulary. Genes beyond this limit are excluded from tokenization. Higher values use more memory. max_genes: Max genes per cell for windowing and tokenization (sequence layout @@ -61,9 +63,10 @@ def __init__( Raises: ValueError: If vocab_size or max_genes is invalid. RuntimeError: If SLAF array is not properly initialized. - TypeError: If slaf_array is not a valid SLAFArray instance. + TypeError: If adata is not a valid LazyAnnData instance. """ - self.slaf_array = slaf_array + self.adata = adata + self.slaf_array = adata.slaf self.vocab_size = vocab_size if max_genes < 1: raise ValueError(f"max_genes must be >= 1, got {max_genes}") @@ -75,6 +78,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 +245,73 @@ def create_window(self) -> Window: Create a window function based on the tokenizer type. """ + def _get_runtime_transformations(self) -> dict[str, Any]: + transformations = getattr(self.adata, "_transformations", None) + if not isinstance(transformations, dict): + return {} + return transformations + + def _apply_runtime_transformations( + self, + df: pl.DataFrame, + schema: DataSchema, + ) -> pl.DataFrame: + transformations = self._get_runtime_transformations() + if not transformations: + return df + + transformed_df = df + value_col = schema.value_key + group_col = schema.group_key + + for transform_name, transform_data in transformations.items(): + if transform_name == "normalize_total": + target_sum = float(transform_data.get("target_sum", 1e4)) + transformed_df = transformed_df.with_columns( + ( + pl.col(value_col) + / pl.col(value_col).sum().over(group_col) + * target_sum + ).alias(value_col) + ) + elif transform_name == "log1p": + transformed_df = transformed_df.with_columns( + pl.col(value_col).log1p().alias(value_col) + ) + + return transformed_df + + def apply( + self, + df: pl.DataFrame, + schema: DataSchema, + max_items: int, + **kwargs: Any, + ) -> pl.DataFrame: + """Group per-cell COO rows into tokenizer-ready sequences.""" + transformed_df = self._apply_runtime_transformations(df, schema) + return self.window.apply( + transformed_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. @@ -317,7 +392,7 @@ class ScGPTTokenizer(SLAFTokenizer): def __init__( self, - slaf_array: SLAFArray, + adata: LazyAnnData, vocab_size: int = 50000, n_expression_bins: int = 10, max_genes: int = 1024, @@ -326,9 +401,7 @@ def __init__( Initialize ScGPTTokenizer with SLAF array and vocabulary settings. Args: - slaf_array: Initialized SLAFArray instance containing the single-cell data. - Used to build the gene vocabulary and access expression data. - Must be a valid SLAFArray with proper var DataFrame. + adata: LazyAnnData instance containing the single-cell data. vocab_size: Maximum size of gene vocabulary. Genes beyond this limit are excluded from tokenization. Higher values use more memory. n_expression_bins: Number of expression bins for scGPT tokenization. @@ -340,16 +413,16 @@ def __init__( Raises: ValueError: If vocab_size is invalid. RuntimeError: If SLAF array is not properly initialized. - TypeError: If slaf_array is not a valid SLAFArray instance. + TypeError: If adata is not a valid LazyAnnData instance. Examples: >>> # Basic initialization >>> slaf_array = SLAFArray("path/to/data.slaf") - >>> tokenizer = ScGPTTokenizer(slaf_array) + >>> tokenizer = ScGPTTokenizer(LazyAnnData(slaf_array)) >>> # scGPT with custom settings >>> tokenizer = ScGPTTokenizer( - ... slaf_array=slaf_array, + ... adata=LazyAnnData(slaf_array), ... vocab_size=30000, ... n_expression_bins=20 ... ) @@ -361,17 +434,94 @@ def __init__( ... tokenizer = ScGPTTokenizer(None) ... except TypeError as e: ... print(f"Error: {e}") - Error: slaf_array must be a valid SLAFArray instance + Error: adata must be a valid LazyAnnData instance """ self.n_expression_bins = n_expression_bins - super().__init__( - slaf_array=slaf_array, vocab_size=vocab_size, max_genes=max_genes - ) + super().__init__(adata=adata, 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]]], @@ -621,17 +771,66 @@ class GeneformerTokenizer(SLAFTokenizer): def __init__( self, - slaf_array: SLAFArray, + adata: LazyAnnData, vocab_size: int = 50000, max_genes: int = 2048, ): - super().__init__( - slaf_array=slaf_array, vocab_size=vocab_size, max_genes=max_genes - ) + super().__init__(adata=adata, 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..1880051 100644 --- a/tests/test_dataloaders.py +++ b/tests/test_dataloaders.py @@ -2,21 +2,37 @@ 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(adata, 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(adata, raw_mode=True, **kwargs) + + if tokenizer_kind == "scgpt": + tokenizer = ScGPTTokenizer(adata, **tokenizer_kwargs) + else: + tokenizer = GeneformerTokenizer(adata, **tokenizer_kwargs) + + return SLAFDataLoader( + adata, + tokenizer=tokenizer, + **kwargs, + ) class TestSLAFDataLoader: """Test suite for SLAFDataLoader class with new architecture""" - def test_dataloader_initialization(self, tiny_slaf): + def test_dataloader_initialization(self, tiny_slaf, tiny_lazy_adata): """Test SLAFDataLoader initialization with new architecture""" - dataloader = SLAFDataLoader(tiny_slaf) + dataloader = build_dataloader(tiny_lazy_adata) # Check basic attributes assert dataloader.slaf_array is tiny_slaf @@ -32,11 +48,11 @@ def test_dataloader_initialization(self, tiny_slaf): # Check that we're using the new dataset implementation assert hasattr(dataloader, "_dataset") - def test_dataloader_initialization_custom_params(self, tiny_slaf): + def test_dataloader_initialization_custom_params(self, tiny_lazy_adata): """Test SLAFDataLoader initialization with custom parameters""" - dataloader = SLAFDataLoader( - tiny_slaf, - tokenizer_type="scgpt", + dataloader = build_dataloader( + tiny_lazy_adata, + tokenizer_kind="scgpt", batch_size=16, max_genes=1024, vocab_size=1000, @@ -50,11 +66,19 @@ def test_dataloader_initialization_custom_params(self, tiny_slaf): assert dataloader.tokenizer.vocab_size == 1000 assert dataloader.tokenizer.n_expression_bins == 5 - def test_geneformer_iteration(self, tiny_slaf): + def test_dataloader_initialization_with_default_source(self, tiny_lazy_adata): + """Test SLAFDataLoader defaults to the main expression matrix.""" + dataloader = build_dataloader(tiny_lazy_adata) + + assert dataloader._dataset.batch_processor.source_logical_key is None + 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_lazy_adata): """Test dataloader iteration with Geneformer tokenizer""" - dataloader = SLAFDataLoader( - tiny_slaf, - tokenizer_type="geneformer", + dataloader = build_dataloader( + tiny_lazy_adata, + tokenizer_kind="geneformer", batch_size=5, max_genes=10, ) @@ -88,11 +112,11 @@ def test_geneformer_iteration(self, tiny_slaf): assert batch_count > 0 - def test_scgpt_iteration(self, tiny_slaf): + def test_scgpt_iteration(self, tiny_lazy_adata): """Test dataloader iteration with scGPT tokenizer""" - dataloader = SLAFDataLoader( - tiny_slaf, - tokenizer_type="scgpt", + dataloader = build_dataloader( + tiny_lazy_adata, + tokenizer_kind="scgpt", batch_size=5, max_genes=10, ) @@ -113,7 +137,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 @@ -127,9 +151,9 @@ def test_scgpt_iteration(self, tiny_slaf): assert batch_count > 0 - def test_consistent_batch_sizes(self, tiny_slaf): + def test_consistent_batch_sizes(self, tiny_lazy_adata): """Test that batches have consistent sizes""" - dataloader = SLAFDataLoader(tiny_slaf, batch_size=8) + dataloader = build_dataloader(tiny_lazy_adata, batch_size=8) batch_sizes = [] for batch in dataloader: @@ -144,9 +168,9 @@ def test_consistent_batch_sizes(self, tiny_slaf): for size in batch_sizes[:-1]: assert size == dataloader.batch_size - def test_cell_id_mapping(self, tiny_slaf): + def test_cell_id_mapping(self, tiny_slaf, tiny_lazy_adata): """Test that cell IDs are properly mapped""" - dataloader = SLAFDataLoader(tiny_slaf, batch_size=5) + dataloader = build_dataloader(tiny_lazy_adata, batch_size=5) for batch in dataloader: cell_ids = batch["cell_ids"] @@ -160,9 +184,9 @@ def test_cell_id_mapping(self, tiny_slaf): break # Just test first batch - def test_tokenizer_integration(self, tiny_slaf): + def test_tokenizer_integration(self, tiny_lazy_adata): """Test that tokenizer is properly integrated""" - dataloader = SLAFDataLoader(tiny_slaf) + dataloader = build_dataloader(tiny_lazy_adata) # Check that tokenizer has expected attributes assert hasattr(dataloader.tokenizer, "gene_vocab") @@ -172,9 +196,9 @@ def test_tokenizer_integration(self, tiny_slaf): # Check that special tokens are properly set assert dataloader.special_tokens == dataloader.tokenizer.special_tokens - def test_dataloader_cleanup(self, tiny_slaf): + def test_dataloader_cleanup(self, tiny_lazy_adata): """Test dataloader cleanup functionality""" - dataloader = SLAFDataLoader(tiny_slaf) + dataloader = build_dataloader(tiny_lazy_adata) # Test that cleanup methods exist and don't crash # The dataloader doesn't have a stop_streaming method @@ -183,9 +207,9 @@ def test_dataloader_cleanup(self, tiny_slaf): # Test destructor dataloader.__del__() - def test_dataloader_device(self, tiny_slaf): + def test_dataloader_device(self, tiny_lazy_adata): """Test dataloader device handling""" - dataloader = SLAFDataLoader(tiny_slaf) + dataloader = build_dataloader(tiny_lazy_adata) # Get a sample batch batch = next(iter(dataloader)) @@ -195,10 +219,10 @@ def test_dataloader_device(self, tiny_slaf): assert batch["attention_mask"].device == torch.device("cpu") assert batch["cell_ids"].device == torch.device("cpu") - def test_multi_epoch_initialization(self, tiny_slaf): + def test_multi_epoch_initialization(self, tiny_lazy_adata): """Test SLAFDataLoader initialization with multi-epoch support""" - dataloader = SLAFDataLoader( - tiny_slaf, + dataloader = build_dataloader( + tiny_lazy_adata, n_epochs=5, # Test multi-epoch initialization ) @@ -206,10 +230,10 @@ def test_multi_epoch_initialization(self, tiny_slaf): assert dataloader._dataset.n_epochs == 5 assert dataloader._dataset.batch_processor.n_epochs == 5 - def test_multi_epoch_iteration(self, tiny_slaf): + def test_multi_epoch_iteration(self, tiny_lazy_adata): """Test dataloader iteration with multiple epochs""" - dataloader = SLAFDataLoader( - tiny_slaf, + dataloader = build_dataloader( + tiny_lazy_adata, batch_size=4, # Small batch size for testing n_epochs=3, # Test with 3 epochs ) @@ -237,10 +261,10 @@ def test_multi_epoch_iteration(self, tiny_slaf): assert len(epochs_seen) >= 1, f"Expected at least 1 epoch, got {epochs_seen}" assert total_batches > 0 - def test_multi_epoch_epoch_progression(self, tiny_slaf): + def test_multi_epoch_epoch_progression(self, tiny_lazy_adata): """Test that epochs progress correctly in multi-epoch mode""" - dataloader = SLAFDataLoader( - tiny_slaf, + dataloader = build_dataloader( + tiny_lazy_adata, batch_size=2, # Very small batch size n_epochs=4, # Test with 4 epochs ) @@ -270,10 +294,10 @@ def test_multi_epoch_epoch_progression(self, tiny_slaf): "Epochs should not go backwards" ) - def test_single_epoch_default_behavior(self, tiny_slaf): + def test_single_epoch_default_behavior(self, tiny_lazy_adata): """Test that single epoch (default) behavior is unchanged""" - dataloader = SLAFDataLoader( - tiny_slaf, + dataloader = build_dataloader( + tiny_lazy_adata, batch_size=8, n_epochs=1, # Single epoch (default) ) @@ -293,12 +317,12 @@ def test_single_epoch_default_behavior(self, tiny_slaf): assert epochs_seen == {0}, f"Expected only epoch 0, got {epochs_seen}" assert batch_count > 0 - def test_multi_epoch_with_different_tokenizers(self, tiny_slaf): + def test_multi_epoch_with_different_tokenizers(self, tiny_lazy_adata): """Test multi-epoch functionality with different tokenizer types""" # Test with Geneformer - limit epochs and batches for speed - dataloader_geneformer = SLAFDataLoader( - tiny_slaf, - tokenizer_type="geneformer", + dataloader_geneformer = build_dataloader( + tiny_lazy_adata, + tokenizer_kind="geneformer", batch_size=4, n_epochs=2, ) @@ -317,9 +341,9 @@ def test_multi_epoch_with_different_tokenizers(self, tiny_slaf): ) # Test with scGPT - limit epochs and batches for speed - dataloader_scgpt = SLAFDataLoader( - tiny_slaf, - tokenizer_type="scgpt", + dataloader_scgpt = build_dataloader( + tiny_lazy_adata, + tokenizer_kind="scgpt", batch_size=4, n_epochs=2, ) @@ -337,10 +361,10 @@ def test_multi_epoch_with_different_tokenizers(self, tiny_slaf): f"scGPT: Expected at least 1 epoch, got {epochs_scgpt}" ) - def test_multi_epoch_completion(self, tiny_slaf): + def test_multi_epoch_completion(self, tiny_lazy_adata): """Test that dataloader correctly completes all epochs""" - dataloader = SLAFDataLoader( - tiny_slaf, + dataloader = build_dataloader( + tiny_lazy_adata, batch_size=2, # Small batch size to complete quickly n_epochs=3, # Test with 3 epochs ) @@ -364,10 +388,10 @@ def test_multi_epoch_completion(self, tiny_slaf): # Should see at least one epoch assert len(epochs_seen) >= 1, f"Expected at least 1 epoch, got {epochs_seen}" - def test_multi_epoch_parameter_passing(self, tiny_slaf): + def test_multi_epoch_parameter_passing(self, tiny_lazy_adata): """Test that n_epochs parameter is correctly passed through the hierarchy""" - dataloader = SLAFDataLoader( - tiny_slaf, + dataloader = build_dataloader( + tiny_lazy_adata, n_epochs=7, # Test with 7 epochs ) @@ -377,11 +401,11 @@ def test_multi_epoch_parameter_passing(self, tiny_slaf): assert dataloader._dataset.batch_processor.n_epochs == 7 assert dataloader._dataset.prefetcher.batch_processor.n_epochs == 7 - def test_multi_epoch_with_custom_parameters(self, tiny_slaf): + def test_multi_epoch_with_custom_parameters(self, tiny_lazy_adata): """Test multi-epoch functionality with custom parameters""" - dataloader = SLAFDataLoader( - tiny_slaf, - tokenizer_type="scgpt", + dataloader = build_dataloader( + tiny_lazy_adata, + tokenizer_kind="scgpt", batch_size=6, max_genes=512, n_epochs=4, @@ -411,11 +435,11 @@ def test_multi_epoch_with_custom_parameters(self, tiny_slaf): assert len(epochs_seen) >= 1, f"Expected at least 1 epoch, got {epochs_seen}" assert batch_count > 0 - def test_dataloader_fragment_parameter(self, tiny_slaf): + def test_dataloader_fragment_parameter(self, tiny_lazy_adata): """Test the by_fragment parameter functionality.""" # Test fragment-based loading dataloader_fragment = SLAFDataLoader( - tiny_slaf, + tiny_lazy_adata, batch_size=32, raw_mode=True, verbose=False, @@ -426,7 +450,7 @@ def test_dataloader_fragment_parameter(self, tiny_slaf): # Test batch-based loading dataloader_batch = SLAFDataLoader( - tiny_slaf, + tiny_lazy_adata, batch_size=32, raw_mode=True, verbose=False, @@ -435,10 +459,10 @@ def test_dataloader_fragment_parameter(self, tiny_slaf): assert dataloader_batch.by_fragment is False - def test_dataloader_iteration_fragment_mode(self, tiny_slaf): + def test_dataloader_iteration_fragment_mode(self, tiny_lazy_adata): """Test dataloader iteration in fragment mode.""" dataloader = SLAFDataLoader( - tiny_slaf, + tiny_lazy_adata, batch_size=32, raw_mode=True, verbose=False, @@ -456,10 +480,10 @@ def test_dataloader_iteration_fragment_mode(self, tiny_slaf): assert batch_count > 0 - def test_dataloader_iteration_batch_mode(self, tiny_slaf): + def test_dataloader_iteration_batch_mode(self, tiny_lazy_adata): """Test dataloader iteration in batch mode.""" dataloader = SLAFDataLoader( - tiny_slaf, + tiny_lazy_adata, batch_size=32, raw_mode=True, verbose=False, @@ -477,10 +501,10 @@ def test_dataloader_iteration_batch_mode(self, tiny_slaf): assert batch_count > 0 - def test_dataloader_tokenized_mode_fragment(self, tiny_slaf): + def test_dataloader_tokenized_mode_fragment(self, tiny_lazy_adata): """Test dataloader in tokenized mode with fragment loading.""" - dataloader = SLAFDataLoader( - tiny_slaf, + dataloader = build_dataloader( + tiny_lazy_adata, batch_size=32, raw_mode=False, verbose=False, @@ -499,10 +523,10 @@ def test_dataloader_tokenized_mode_fragment(self, tiny_slaf): assert batch_count > 0 - def test_dataloader_tokenized_mode_batch(self, tiny_slaf): + def test_dataloader_tokenized_mode_batch(self, tiny_lazy_adata): """Test dataloader in tokenized mode with batch loading.""" - dataloader = SLAFDataLoader( - tiny_slaf, + dataloader = build_dataloader( + tiny_lazy_adata, batch_size=32, raw_mode=False, verbose=False, @@ -521,14 +545,11 @@ def test_dataloader_tokenized_mode_batch(self, tiny_slaf): assert batch_count > 0 - def test_dataloader_parameters_consistency(self, tiny_slaf): + def test_dataloader_parameters_consistency(self, tiny_lazy_adata): """Test that all parameters are properly passed through.""" - dataloader = SLAFDataLoader( - tiny_slaf, + dataloader = build_dataloader( + tiny_lazy_adata, batch_size=64, - max_genes=1024, - vocab_size=10000, - n_expression_bins=5, n_epochs=5, raw_mode=True, verbose=True, @@ -538,27 +559,27 @@ 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 assert dataloader.batches_per_chunk == 25 assert dataloader.by_fragment is True - def test_dataloader_length(self, tiny_slaf): + def test_dataloader_length(self, tiny_lazy_adata): """Test that dataloader length returns 0 for streaming datasets.""" - dataloader = SLAFDataLoader( - tiny_slaf, + dataloader = build_dataloader( + tiny_lazy_adata, batch_size=32, verbose=False, ) assert len(dataloader) == 0 # Streaming datasets have unknown length (return 0) - def test_mixture_of_scanners_initialization(self, tiny_slaf): + def test_mixture_of_scanners_initialization(self, tiny_lazy_adata): """Test SLAFDataLoader initialization with Mixture of Scanners (MoS)""" - dataloader = SLAFDataLoader( - slaf_array=tiny_slaf, + dataloader = build_dataloader( + tiny_lazy_adata, batch_size=32, use_mixture_of_scanners=True, n_scanners=8, @@ -573,11 +594,11 @@ def test_mixture_of_scanners_initialization(self, tiny_slaf): assert dataloader._dataset.batch_processor.n_scanners == 8 assert dataloader._dataset.batch_processor.prefetch_batch_size == 1048576 - def test_mixture_of_scanners_parameter_validation(self, tiny_slaf): + def test_mixture_of_scanners_parameter_validation(self, tiny_lazy_adata): """Test MoS parameter validation in SLAFDataLoader""" # Test valid parameters - dataloader = SLAFDataLoader( - slaf_array=tiny_slaf, + dataloader = build_dataloader( + tiny_lazy_adata, batch_size=32, use_mixture_of_scanners=True, n_scanners=16, @@ -588,7 +609,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_lazy_adata, batch_size=32, use_mixture_of_scanners=True, n_scanners=0, @@ -598,7 +619,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_lazy_adata, batch_size=32, use_mixture_of_scanners=True, n_scanners=101, @@ -610,7 +631,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_lazy_adata, batch_size=32, use_mixture_of_scanners=True, n_scanners=16, @@ -622,17 +643,17 @@ 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_lazy_adata, batch_size=32, use_mixture_of_scanners=True, n_scanners=16, prefetch_batch_size=10000001, ) - def test_mixture_of_scanners_iteration(self, tiny_slaf): + def test_mixture_of_scanners_iteration(self, tiny_lazy_adata): """Test that MoS dataloader can iterate through batches""" - dataloader = SLAFDataLoader( - slaf_array=tiny_slaf, + dataloader = build_dataloader( + tiny_lazy_adata, batch_size=8, use_mixture_of_scanners=True, n_scanners=4, @@ -651,10 +672,10 @@ def test_mixture_of_scanners_iteration(self, tiny_slaf): assert batch_count > 0 - def test_mixture_of_scanners_with_raw_mode(self, tiny_slaf): + def test_mixture_of_scanners_with_raw_mode(self, tiny_lazy_adata): """Test MoS functionality with raw mode in SLAFDataLoader""" dataloader = SLAFDataLoader( - slaf_array=tiny_slaf, + tiny_lazy_adata, batch_size=32, raw_mode=True, use_mixture_of_scanners=True, @@ -676,10 +697,10 @@ def test_mixture_of_scanners_with_raw_mode(self, tiny_slaf): assert batch_count > 0 - def test_mixture_of_scanners_with_tokenized_mode(self, tiny_slaf): + def test_mixture_of_scanners_with_tokenized_mode(self, tiny_lazy_adata): """Test MoS functionality with tokenized mode in SLAFDataLoader""" - dataloader = SLAFDataLoader( - slaf_array=tiny_slaf, + dataloader = build_dataloader( + tiny_lazy_adata, batch_size=32, raw_mode=False, use_mixture_of_scanners=True, @@ -702,11 +723,11 @@ def test_mixture_of_scanners_with_tokenized_mode(self, tiny_slaf): assert batch_count > 0 - def test_mixture_of_scanners_backward_compatibility(self, tiny_slaf): + def test_mixture_of_scanners_backward_compatibility(self, tiny_lazy_adata): """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_lazy_adata, batch_size=32, ) @@ -716,8 +737,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_lazy_adata, batch_size=32, use_mixture_of_scanners=False, ) @@ -728,10 +749,10 @@ def test_mixture_of_scanners_backward_compatibility(self, tiny_slaf): is False ) - def test_mixture_of_scanners_parameter_passing(self, tiny_slaf): + def test_mixture_of_scanners_parameter_passing(self, tiny_lazy_adata): """Test that MoS parameters are correctly passed through the hierarchy""" - dataloader = SLAFDataLoader( - slaf_array=tiny_slaf, + dataloader = build_dataloader( + tiny_lazy_adata, batch_size=32, use_mixture_of_scanners=True, n_scanners=12, @@ -747,11 +768,11 @@ def test_mixture_of_scanners_parameter_passing(self, tiny_slaf): assert dataloader._dataset.batch_processor.n_scanners == 12 assert dataloader._dataset.batch_processor.prefetch_batch_size == 2097152 - def test_mixture_of_scanners_with_custom_parameters(self, tiny_slaf): + def test_mixture_of_scanners_with_custom_parameters(self, tiny_lazy_adata): """Test MoS functionality with custom parameters in SLAFDataLoader""" - dataloader = SLAFDataLoader( - slaf_array=tiny_slaf, - tokenizer_type="scgpt", + dataloader = build_dataloader( + tiny_lazy_adata, + tokenizer_kind="scgpt", batch_size=16, use_mixture_of_scanners=True, n_scanners=6, @@ -764,7 +785,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 @@ -778,10 +799,10 @@ def test_mixture_of_scanners_with_custom_parameters(self, tiny_slaf): assert batch_count > 0 - def test_mixture_of_scanners_multi_epoch(self, tiny_slaf): + def test_mixture_of_scanners_multi_epoch(self, tiny_lazy_adata): """Test MoS functionality with multi-epoch training""" - dataloader = SLAFDataLoader( - slaf_array=tiny_slaf, + dataloader = build_dataloader( + tiny_lazy_adata, batch_size=8, n_epochs=3, use_mixture_of_scanners=True, @@ -808,10 +829,10 @@ def test_mixture_of_scanners_multi_epoch(self, tiny_slaf): assert len(epochs_seen) >= 1, f"Expected at least 1 epoch, got {epochs_seen}" assert batch_count > 0 - def test_mixture_of_scanners_fragment_generators_creation(self, tiny_slaf): + def test_mixture_of_scanners_fragment_generators_creation(self, tiny_lazy_adata): """Test that MoS creates fragment generators correctly""" - dataloader = SLAFDataLoader( - slaf_array=tiny_slaf, + dataloader = build_dataloader( + tiny_lazy_adata, batch_size=32, use_mixture_of_scanners=True, n_scanners=4, @@ -832,10 +853,10 @@ def test_mixture_of_scanners_fragment_generators_creation(self, tiny_slaf): dataloader._dataset.batch_processor.fragment_generators ) - def test_mixture_of_scanners_random_sampling_behavior(self, tiny_slaf): + def test_mixture_of_scanners_random_sampling_behavior(self, tiny_lazy_adata): """Test that MoS uses random sampling from fragment generators""" - dataloader = SLAFDataLoader( - slaf_array=tiny_slaf, + dataloader = build_dataloader( + tiny_lazy_adata, batch_size=8, use_mixture_of_scanners=True, n_scanners=2, # Small number for testing @@ -857,10 +878,10 @@ def test_mixture_of_scanners_random_sampling_behavior(self, tiny_slaf): assert "attention_mask" in batch assert "cell_ids" in batch - def test_mixture_of_scanners_cell_boundary_handling(self, tiny_slaf): + def test_mixture_of_scanners_cell_boundary_handling(self, tiny_lazy_adata): """Test that MoS handles cell boundaries correctly in SLAFDataLoader""" - dataloader = SLAFDataLoader( - slaf_array=tiny_slaf, + dataloader = build_dataloader( + tiny_lazy_adata, batch_size=8, use_mixture_of_scanners=True, n_scanners=4, @@ -917,9 +938,9 @@ def test_tensor_on_optimal_device(self): @pytest.mark.skipif( not get_device_info()["torch_available"], reason="PyTorch not available" ) - def test_dataloader_device(self, tiny_slaf): + def test_dataloader_device(self, tiny_lazy_adata): """Test that dataloader uses the correct device""" - dataloader = SLAFDataLoader(tiny_slaf) + dataloader = build_dataloader(tiny_lazy_adata) # 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 c06df3b..c50687a 100644 --- a/tests/test_distributed_worker.py +++ b/tests/test_distributed_worker.py @@ -9,8 +9,12 @@ from unittest.mock import MagicMock, patch import polars as pl +import pytest from slaf.distributed.data_source import DataSource + +pytest.importorskip("omegaconf") + from slaf.distributed.worker import prefetch_worker diff --git a/tests/test_tokenizers.py b/tests/test_tokenizers.py index 7580ca7..7cc95c1 100644 --- a/tests/test_tokenizers.py +++ b/tests/test_tokenizers.py @@ -7,27 +7,37 @@ import numpy as np import pandas as pd +import polars as pl import pytest import torch from slaf.core.slaf import SLAFArray +from slaf.core.tabular_schema import SLAF_LANCE_COO_SCHEMA +from slaf.integrations.anndata import LazyAnnData from slaf.ml.tokenizers import GeneformerTokenizer, ScGPTTokenizer +def build_mock_adata(): + 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_adata = Mock(spec=LazyAnnData) + mock_adata.slaf = mock_slaf_array + mock_adata._transformations = {} + return mock_adata, 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_adata, _ = build_mock_adata() # Test Geneformer initialization tokenizer = GeneformerTokenizer( - slaf_array=mock_slaf_array, + adata=mock_adata, vocab_size=1000, ) @@ -40,7 +50,7 @@ def test_tokenizer_initialization(self): # Test scGPT initialization tokenizer = ScGPTTokenizer( - slaf_array=mock_slaf_array, + adata=mock_adata, vocab_size=1000, n_expression_bins=5, ) @@ -51,14 +61,10 @@ 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_adata, _ = build_mock_adata() tokenizer = GeneformerTokenizer( - slaf_array=mock_slaf_array, + adata=mock_adata, vocab_size=1000, ) @@ -83,14 +89,10 @@ 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_adata, _ = build_mock_adata() tokenizer = ScGPTTokenizer( - slaf_array=mock_slaf_array, + adata=mock_adata, vocab_size=1000, n_expression_bins=10, ) @@ -128,16 +130,51 @@ 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_adata, _ = build_mock_adata() + + tokenizer = ScGPTTokenizer( + adata=mock_adata, + 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_adata, _ = build_mock_adata() + + tokenizer = GeneformerTokenizer( + adata=mock_adata, + 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_adata, _ = build_mock_adata() tokenizer = ScGPTTokenizer( - slaf_array=mock_slaf_array, + adata=mock_adata, vocab_size=1000, ) @@ -155,14 +192,10 @@ 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_adata, _ = build_mock_adata() tokenizer = GeneformerTokenizer( - slaf_array=mock_slaf_array, + adata=mock_adata, vocab_size=1000, ) @@ -184,13 +217,10 @@ 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_adata, _ = build_mock_adata() tokenizer = ScGPTTokenizer( - slaf_array=mock_slaf_array, + adata=mock_adata, vocab_size=1000, n_expression_bins=10, ) @@ -214,14 +244,10 @@ 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_adata, _ = build_mock_adata() tokenizer = GeneformerTokenizer( - slaf_array=mock_slaf_array, + adata=mock_adata, vocab_size=1000, ) @@ -246,14 +272,10 @@ 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_adata, _ = build_mock_adata() tokenizer = GeneformerTokenizer( - slaf_array=mock_slaf_array, + adata=mock_adata, vocab_size=1000, ) @@ -264,16 +286,87 @@ def test_vocabulary_info(self): # The tokenizer creates a fallback vocabulary assert vocab_info["gene_vocab_size"] > 0 + def test_apply_uses_slaf_runtime_transformations(self): + mock_adata, _ = build_mock_adata() + mock_adata._transformations = { + "normalize_total": { + "type": "normalize_total", + "target_sum": 100.0, + }, + "log1p": { + "type": "log1p", + "applied": True, + }, + } + + tokenizer = ScGPTTokenizer( + adata=mock_adata, + vocab_size=1000, + n_expression_bins=10, + ) + + df = pl.DataFrame( + { + "cell_integer_id": [0, 0, 1, 1], + "gene_integer_id": [0, 1, 0, 1], + "value": [1.0, 3.0, 2.0, 2.0], + } + ) + + grouped = tokenizer.apply( + df, + schema=SLAF_LANCE_COO_SCHEMA, + max_items=2, + use_binned_expressions=False, + ) + + expr_sequences = grouped["expr_sequence"].to_list() + expected = [ + [np.log1p(25.0), np.log1p(75.0)], + [np.log1p(50.0), np.log1p(50.0)], + ] + for actual_seq, expected_seq in zip(expr_sequences, expected, strict=False): + assert np.asarray(actual_seq) == pytest.approx(np.asarray(expected_seq)) + + def test_scgpt_window_does_not_double_log1p(self): + mock_adata, _ = build_mock_adata() + mock_adata._transformations = { + "log1p": { + "type": "log1p", + "applied": True, + }, + } + + tokenizer = ScGPTTokenizer( + adata=mock_adata, + vocab_size=1000, + n_expression_bins=10, + ) + + df = pl.DataFrame( + { + "cell_integer_id": [0, 0], + "gene_integer_id": [0, 1], + "value": [np.e - 1.0, np.exp(2.0) - 1.0], + } + ) + + grouped = tokenizer.apply( + df, + schema=SLAF_LANCE_COO_SCHEMA, + max_items=2, + use_binned_expressions=True, + ) + + expr_tokens = grouped["expr_sequence"].to_list()[0] + assert expr_tokens == [1005, 1009] + 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_adata, _ = build_mock_adata() tokenizer = ScGPTTokenizer( - slaf_array=mock_slaf_array, + adata=mock_adata, vocab_size=1000, n_expression_bins=10, ) @@ -315,8 +408,9 @@ class TestSLAFTokenizerWithRealData: def test_tokenizer_with_real_data(self, tiny_slaf): """Test tokenizer with real SLAF data.""" + adata = LazyAnnData(tiny_slaf) tokenizer = GeneformerTokenizer( - slaf_array=tiny_slaf, + adata=adata, vocab_size=1000, ) @@ -335,8 +429,9 @@ def test_tokenizer_with_real_data(self, tiny_slaf): def test_scgpt_with_real_data(self, tiny_slaf): """Test scGPT tokenizer with real SLAF data.""" + adata = LazyAnnData(tiny_slaf) tokenizer = ScGPTTokenizer( - slaf_array=tiny_slaf, + adata=adata, vocab_size=1000, n_expression_bins=10, ) @@ -357,8 +452,9 @@ def test_scgpt_with_real_data(self, tiny_slaf): def test_gene_mapping_with_real_data(self, tiny_slaf): """Test gene ID mapping with real SLAF data.""" + adata = LazyAnnData(tiny_slaf) tokenizer = GeneformerTokenizer( - slaf_array=tiny_slaf, + adata=adata, vocab_size=1000, ) 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 1bed8cf131a315b1c57fb5bc8ef29210d4a95838 Mon Sep 17 00:00:00 2001 From: Shahul Alam Date: Thu, 14 May 2026 16:11:41 -0400 Subject: [PATCH 02/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 9600f4a..17344de 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() @@ -1092,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, @@ -1776,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}", From a4d36c11fcfeb175dd5415dbaebaf4d517669a42 Mon Sep 17 00:00:00 2001 From: Shahul Alam Date: Sat, 23 May 2026 13:15:53 -0400 Subject: [PATCH 03/10] Fix tokenizers to use precomputed factors for normalization. --- slaf/integrations/anndata.py | 71 ++++++----------------------- slaf/integrations/scanpy.py | 88 ++++++++---------------------------- slaf/ml/tokenizers.py | 31 ++++++++++--- tests/test_tokenizers.py | 71 +++++++++++++++++++++++++++++ 4 files changed, 127 insertions(+), 134 deletions(-) diff --git a/slaf/integrations/anndata.py b/slaf/integrations/anndata.py index 0530d16..5227f0a 100644 --- a/slaf/integrations/anndata.py +++ b/slaf/integrations/anndata.py @@ -616,68 +616,25 @@ def _apply_normalize_total( ) -> scipy.sparse.csr_matrix: """Apply normalize_total transformation using vectorized operations""" cell_factors = transform_data.get("cell_factors", {}) - obs_names_local = [] # Always a list - obs_names = None - if hasattr(self, "parent_adata") and self.parent_adata is not None: - try: - obs_names = self.parent_adata.obs_names - except (AttributeError, TypeError): - obs_names = None - if obs_names is None or not isinstance(obs_names, list | np.ndarray | pd.Index): - if ( - matrix is not None - and hasattr(matrix, "shape") - and matrix.shape is not None - ): - obs_names_local = [f"cell_{i}" for i in range(matrix.shape[0])] - else: - obs_names_local = [] - else: - obs_names_local = list(obs_names) - # Now obs_names_local is always a list - # Determine selected_cell_names based on cell_selector - if cell_selector is None or ( - isinstance(cell_selector, slice) and cell_selector == slice(None) - ): - selected_cell_names = obs_names_local + + if isinstance(cell_selector, int | np.integer): + selected_cell_integer_ids = [int(cell_selector)] elif isinstance(cell_selector, slice): - start = cell_selector.start or 0 - stop = cell_selector.stop or len(obs_names_local) - step = cell_selector.step or 1 - # Clamp bounds - start = max(0, min(start, len(obs_names_local))) - stop = max(0, min(stop, len(obs_names_local))) - selected_cell_names = obs_names_local[start:stop:step] + selected_cell_integer_ids = list(range(self.slaf_array.shape[0]))[ + cell_selector + ] + elif isinstance(cell_selector, np.ndarray) and cell_selector.dtype == bool: + selected_cell_integer_ids = ( + np.flatnonzero(cell_selector).astype(int).tolist() + ) elif isinstance(cell_selector, list | np.ndarray): - if len(obs_names_local) > 0: - if ( - isinstance(cell_selector, np.ndarray) - and cell_selector.dtype == bool - ): - selected_cell_names = [ - obs_names_local[i] - for i, keep in enumerate(cell_selector) - if keep and 0 <= i < len(obs_names_local) - ] - else: - selected_cell_names = [ - obs_names_local[i] - for i in cell_selector - if isinstance(i, int | np.integer) - and 0 <= i < len(obs_names_local) - ] - else: - selected_cell_names = [] - elif isinstance(cell_selector, int | np.integer): - if 0 <= cell_selector < len(obs_names_local): - selected_cell_names = [obs_names_local[cell_selector]] - else: - selected_cell_names = [] + selected_cell_integer_ids = [int(i) for i in cell_selector] else: - selected_cell_names = obs_names_local + selected_cell_integer_ids = list(range(matrix.shape[0])) + # Create a vector of factors for all cells at once cell_factors_vector = np.array( - [cell_factors.get(name, 1.0) for name in selected_cell_names] + [cell_factors.get(cell_id, 1.0) for cell_id in selected_cell_integer_ids] ) # Apply vectorized scaling using CSR matrix properties # Create a copy only if we need to modify the data diff --git a/slaf/integrations/scanpy.py b/slaf/integrations/scanpy.py index c672edf..16401d4 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 @@ -729,70 +730,17 @@ def normalize_total( # Work with polars DataFrame internally cell_totals_pl = cell_totals - # Create normalization factors using polars - # Map cell_integer_id to cell names for compatibility with anndata.py - if hasattr(adata.slaf, "obs") and adata.slaf.obs is not None: - # Create mapping from cell_integer_id to cell names - # Use polars DataFrame to get cell names - obs_df = adata.slaf.obs - if "cell_id" in obs_df.columns: - cell_id_to_name = dict( - zip( - obs_df["cell_integer_id"].to_list(), - obs_df["cell_id"].to_list(), - strict=False, - ) - ) - else: - # Fallback: create cell names from integer IDs - cell_id_to_name = {i: f"cell_{i}" for i in range(len(obs_df))} - # Use vectorized polars operations for mapping - # Create mapping DataFrame - cell_map_df = pl.DataFrame( - { - "cell_integer_id": list(cell_id_to_name.keys()), - "cell_id": list(cell_id_to_name.values()), - } - ) - # Join with mapping DataFrame - cell_totals_pl = cell_totals_pl.join( - cell_map_df, on="cell_integer_id", how="left" - ) - # Fill any missing values with default format - cell_totals_pl = cell_totals_pl.with_columns( - [ - pl.col("cell_id").fill_null( - pl.col("cell_integer_id") - .cast(pl.Utf8) - .map_elements(lambda x: f"cell_{x}", return_dtype=pl.Utf8) - ), - (target_sum / pl.col("total_counts")).alias("normalization_factor"), - ] - ) - # Convert to dictionary for compatibility - normalization_dict = dict( - zip( - cell_totals_pl["cell_id"].to_list(), - cell_totals_pl["normalization_factor"].to_list(), - strict=False, - ) - ) - else: - # Fallback: use cell_integer_id as string keys - cell_totals_pl = cell_totals_pl.with_columns( - [ - pl.col("cell_integer_id").cast(pl.Utf8).alias("cell_id"), - (target_sum / pl.col("total_counts")).alias("normalization_factor"), - ] - ) - # Convert to dictionary for compatibility - normalization_dict = dict( - zip( - cell_totals_pl["cell_id"].to_list(), - cell_totals_pl["normalization_factor"].to_list(), - strict=False, - ) + # Store factors keyed by SLAF's internal cell_integer_id. + cell_totals_pl = cell_totals_pl.with_columns( + (target_sum / pl.col("total_counts")).alias("normalization_factor") + ) + normalization_dict = dict( + zip( + cell_totals_pl["cell_integer_id"].to_list(), + cell_totals_pl["normalization_factor"].to_list(), + strict=False, ) + ) if inplace: # Store normalization factors for lazy application @@ -807,7 +755,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 +847,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 +862,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 +1048,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 diff --git a/slaf/ml/tokenizers.py b/slaf/ml/tokenizers.py index c301457..c9ae97c 100644 --- a/slaf/ml/tokenizers.py +++ b/slaf/ml/tokenizers.py @@ -266,13 +266,30 @@ def _apply_runtime_transformations( for transform_name, transform_data in transformations.items(): if transform_name == "normalize_total": - target_sum = float(transform_data.get("target_sum", 1e4)) - transformed_df = transformed_df.with_columns( - ( - pl.col(value_col) - / pl.col(value_col).sum().over(group_col) - * target_sum - ).alias(value_col) + cell_factors = transform_data.get("cell_factors") + if not isinstance(cell_factors, dict) or not cell_factors: + raise ValueError( + "normalize_total runtime transformation requires precomputed " + "cell_factors keyed by cell_integer_id" + ) + + factor_df = pl.DataFrame( + { + group_col: [int(cell_id) for cell_id in cell_factors.keys()], + "_normalization_factor": [ + float(factor) for factor in cell_factors.values() + ], + } + ) + transformed_df = ( + transformed_df.join(factor_df, on=group_col, how="left") + .with_columns( + ( + pl.col(value_col) + * pl.col("_normalization_factor").fill_null(1.0) + ).alias(value_col) + ) + .drop("_normalization_factor") ) elif transform_name == "log1p": transformed_df = transformed_df.with_columns( diff --git a/tests/test_tokenizers.py b/tests/test_tokenizers.py index 7cc95c1..2008d0f 100644 --- a/tests/test_tokenizers.py +++ b/tests/test_tokenizers.py @@ -292,6 +292,7 @@ def test_apply_uses_slaf_runtime_transformations(self): "normalize_total": { "type": "normalize_total", "target_sum": 100.0, + "cell_factors": {0: 25.0, 1: 25.0}, }, "log1p": { "type": "log1p", @@ -328,6 +329,76 @@ def test_apply_uses_slaf_runtime_transformations(self): for actual_seq, expected_seq in zip(expr_sequences, expected, strict=False): assert np.asarray(actual_seq) == pytest.approx(np.asarray(expected_seq)) + def test_apply_requires_precomputed_cell_factors_for_normalize_total(self): + mock_adata, _ = build_mock_adata() + mock_adata._transformations = { + "normalize_total": { + "type": "normalize_total", + "target_sum": 100.0, + }, + } + + tokenizer = ScGPTTokenizer( + adata=mock_adata, + vocab_size=1000, + n_expression_bins=10, + ) + + df = pl.DataFrame( + { + "cell_integer_id": [0, 0, 1, 1], + "gene_integer_id": [0, 1, 0, 1], + "value": [1.0, 3.0, 2.0, 2.0], + } + ) + + with pytest.raises( + ValueError, + match="requires precomputed cell_factors keyed by cell_integer_id", + ): + tokenizer.apply( + df, + schema=SLAF_LANCE_COO_SCHEMA, + max_items=2, + use_binned_expressions=False, + ) + + def test_apply_uses_precomputed_cell_factors_by_integer_id(self): + mock_adata, _ = build_mock_adata() + mock_adata._transformations = { + "normalize_total": { + "type": "normalize_total", + "target_sum": 100.0, + "cell_factors": {0: 10.0, 1: 5.0}, + }, + } + + tokenizer = ScGPTTokenizer( + adata=mock_adata, + vocab_size=1000, + n_expression_bins=10, + ) + + df = pl.DataFrame( + { + "cell_integer_id": [0, 0, 1, 1], + "gene_integer_id": [0, 1, 0, 1], + "value": [1.0, 3.0, 2.0, 2.0], + } + ) + + grouped = tokenizer.apply( + df, + schema=SLAF_LANCE_COO_SCHEMA, + max_items=2, + use_binned_expressions=False, + ) + + expr_sequences = grouped["expr_sequence"].to_list() + expected = [[10.0, 30.0], [10.0, 10.0]] + for actual_seq, expected_seq in zip(expr_sequences, expected, strict=False): + assert np.asarray(actual_seq) == pytest.approx(np.asarray(expected_seq)) + def test_scgpt_window_does_not_double_log1p(self): mock_adata, _ = build_mock_adata() mock_adata._transformations = { From 7e21ca3ec67edafbd2a25bad1e695add8ce70194 Mon Sep 17 00:00:00 2001 From: Shahul Alam Date: Sat, 23 May 2026 13:33:02 -0400 Subject: [PATCH 04/10] Clean up configuration objects with OmegaConf. --- slaf/distributed/worker.py | 37 ++----- slaf/ml/distributed.py | 9 +- tests/test_distributed_worker.py | 173 +++++++++---------------------- 3 files changed, 64 insertions(+), 155 deletions(-) diff --git a/slaf/distributed/worker.py b/slaf/distributed/worker.py index b298aab..98febb5 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: DictConfig | dict[str, Any], - processor_config: DictConfig | 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,29 +56,14 @@ def prefetch_worker( Returns: Dictionary with worker metrics """ - 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) - - tokenizer_config = OmegaConf.select(processor_config, "tokenizer_config") - shuffle_factory_config = OmegaConf.select(processor_config, "shuffle_factory") - window_factory_config = OmegaConf.select(processor_config, "window_factory") - use_tokenizer_window = bool( - OmegaConf.select(processor_config, "use_tokenizer_window", default=False) - ) - continuity_check = str( - OmegaConf.select( - processor_config, - "continuity_check", - default="sequential", - ) - ) - max_items = int(OmegaConf.select(processor_config, "max_items", default=1024)) - seed_value = int(OmegaConf.select(processor_config, "seed", default=42)) - window_kwargs = ( - OmegaConf.select(processor_config, "window_kwargs", default={}) or {} - ) + 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 @@ -312,7 +297,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 e609fcd..6c3eba1 100644 --- a/slaf/ml/distributed.py +++ b/slaf/ml/distributed.py @@ -96,8 +96,8 @@ def create_app( def distributed_prefetch_worker( worker_id: str, partition_indices: list[int], - data_source_config: DictConfig | dict[str, Any], - processor_config: DictConfig | dict[str, Any], + data_source_config: DictConfig, + processor_config: DictConfig, queue_name: str, n_scanners: int = 8, prefetch_batch_count: int = 32, @@ -114,11 +114,6 @@ 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. diff --git a/tests/test_distributed_worker.py b/tests/test_distributed_worker.py index c50687a..deb5ab7 100644 --- a/tests/test_distributed_worker.py +++ b/tests/test_distributed_worker.py @@ -15,6 +15,8 @@ pytest.importorskip("omegaconf") +from omegaconf import OmegaConf + from slaf.distributed.worker import prefetch_worker @@ -68,6 +70,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", @@ -106,18 +135,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( @@ -151,18 +170,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( @@ -203,18 +212,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( @@ -247,18 +246,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( @@ -291,18 +280,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( @@ -335,18 +314,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( @@ -381,20 +350,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( @@ -428,18 +387,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( @@ -473,18 +422,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( @@ -517,18 +456,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 d7ad5abab040e077173bff2e1444a5ed19c323fc Mon Sep 17 00:00:00 2001 From: Shahul Alam Date: Sat, 23 May 2026 13:36:44 -0400 Subject: [PATCH 05/10] Add back lazy_tiny_adata fixture. --- tests/test_dataloaders.py | 13 ++++++++++--- 1 file changed, 10 insertions(+), 3 deletions(-) diff --git a/tests/test_dataloaders.py b/tests/test_dataloaders.py index 1880051..e9bc17e 100644 --- a/tests/test_dataloaders.py +++ b/tests/test_dataloaders.py @@ -2,10 +2,16 @@ import pytest import torch +from slaf.integrations.anndata import LazyAnnData from slaf.ml.dataloaders import SLAFDataLoader, get_device_info, get_optimal_device from slaf.ml.tokenizers import GeneformerTokenizer, ScGPTTokenizer +@pytest.fixture +def tiny_lazy_adata(tiny_slaf): + return LazyAnnData(tiny_slaf) + + def build_dataloader(adata, tokenizer_kind="geneformer", raw_mode=False, **kwargs): tokenizer_kwargs = {} for key in ("max_genes", "vocab_size", "n_expression_bins"): @@ -69,10 +75,11 @@ def test_dataloader_initialization_custom_params(self, tiny_lazy_adata): def test_dataloader_initialization_with_default_source(self, tiny_lazy_adata): """Test SLAFDataLoader defaults to the main expression matrix.""" dataloader = build_dataloader(tiny_lazy_adata) + batch_processor = dataloader._dataset.batch_processor - assert dataloader._dataset.batch_processor.source_logical_key is None - assert dataloader._dataset.batch_processor.use_mixture_of_scanners is True - assert dataloader._dataset.batch_processor.by_fragment is True + assert batch_processor.expression_dataset is not None + assert batch_processor.use_mixture_of_scanners is True + assert batch_processor.by_fragment is True def test_geneformer_iteration(self, tiny_lazy_adata): """Test dataloader iteration with Geneformer tokenizer""" From 589af7447064e094ab1ced7b6f5b8b29472cc711 Mon Sep 17 00:00:00 2001 From: Shahul Alam Date: Sat, 23 May 2026 15:46:42 -0400 Subject: [PATCH 06/10] Refactor transformation + tokenization. --- slaf/ml/aggregators.py | 21 +-- slaf/ml/datasets.py | 2 +- slaf/ml/tokenizers.py | 37 +++--- tests/conftest.py | 6 + tests/test_aggregators.py | 8 +- tests/test_dataloaders.py | 6 - tests/test_pytorch_datasets.py | 233 ++++++++++++++++++--------------- tests/test_tokenizers.py | 31 ++--- 8 files changed, 181 insertions(+), 163 deletions(-) diff --git a/slaf/ml/aggregators.py b/slaf/ml/aggregators.py index f68cb4d..c2ea93b 100644 --- a/slaf/ml/aggregators.py +++ b/slaf/ml/aggregators.py @@ -90,9 +90,6 @@ def apply( preprocessor = kwargs.get("expression_preprocessor") fragment_df = apply_expression_preprocessor(fragment_df, schema, preprocessor) - already_log1p = ( - isinstance(preprocessor, ExpressionPreprocessor) and preprocessor.log1p - ) if use_binned_expressions: grouped = ( @@ -104,20 +101,14 @@ def apply( ) .filter(pl.col("gene_rank") <= max_items) .with_columns( - (pl.col(vk) if already_log1p else pl.col(vk).log1p()).alias( - "log_value" - ) - ) - .with_columns( - pl.when(pl.col("log_value") > 0) + 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/datasets.py b/slaf/ml/datasets.py index 17344de..997487b 100644 --- a/slaf/ml/datasets.py +++ b/slaf/ml/datasets.py @@ -1045,7 +1045,7 @@ def load_prefetch_batch(self) -> PrefetchBatch: if tokenizer is None: raise RuntimeError("Tokenizer is required for tokenized mode") - grouped = tokenizer.apply( + grouped = tokenizer.transform_and_apply( shuffled_df, schema=SLAF_LANCE_COO_SCHEMA, max_items=tokenizer.max_genes, diff --git a/slaf/ml/tokenizers.py b/slaf/ml/tokenizers.py index c9ae97c..bfe654e 100644 --- a/slaf/ml/tokenizers.py +++ b/slaf/ml/tokenizers.py @@ -298,21 +298,26 @@ def _apply_runtime_transformations( return transformed_df - def apply( + def transform_and_apply( self, df: pl.DataFrame, schema: DataSchema, max_items: int, **kwargs: Any, ) -> pl.DataFrame: - """Group per-cell COO rows into tokenizer-ready sequences.""" + """Apply runtime transformations, then group rows into tokenizer-ready sequences.""" transformed_df = self._apply_runtime_transformations(df, schema) - return self.window.apply( - transformed_df, - schema=schema, - max_items=max_items, - **kwargs, - ) + return self.apply(transformed_df, schema=schema, max_items=max_items, **kwargs) + + def apply( + self, + df: pl.DataFrame, + schema: DataSchema, + max_items: int, + **kwargs: Any, + ) -> pl.DataFrame: + """Group already-transformed per-cell COO rows into tokenizer-ready sequences.""" + return self.window.apply(df, schema=schema, max_items=max_items, **kwargs) def tokenize_grouped( self, @@ -468,9 +473,8 @@ def apply( **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) + return self.window.apply(df, schema=schema, max_items=max_items, **kwargs) def tokenize_grouped( self, @@ -606,9 +610,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) @@ -671,7 +673,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: @@ -683,7 +684,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 @@ -703,7 +704,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"], ) @@ -747,8 +748,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: diff --git a/tests/conftest.py b/tests/conftest.py index e96082e..11a5b4d 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -12,6 +12,7 @@ from slaf.core.slaf import SLAFArray from slaf.data import SLAFConverter from slaf.integrations import ensure_h5ad_writable +from slaf.integrations.anndata import LazyAnnData def _wait_for_all_slaf_threads(timeout=0.3): @@ -441,6 +442,11 @@ def tiny_slaf(temp_dir, tiny_adata): return slaf_array +@pytest.fixture +def tiny_lazy_adata(tiny_slaf): + return LazyAnnData(tiny_slaf) + + @pytest.fixture def small_slaf(temp_dir, small_adata): """Create a small SLAFArray for SLAFArray testing.""" diff --git a/tests/test_aggregators.py b/tests/test_aggregators.py index 333072e..59d43d6 100644 --- a/tests/test_aggregators.py +++ b/tests/test_aggregators.py @@ -81,7 +81,7 @@ def test_apply_ranking(self): assert gene_seq[0] == 10 # gene with value 5.0 # 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 def test_apply_expression_binning(self): """Test that expression binning works correctly""" @@ -91,7 +91,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 +122,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 +167,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_dataloaders.py b/tests/test_dataloaders.py index e9bc17e..c8da8a4 100644 --- a/tests/test_dataloaders.py +++ b/tests/test_dataloaders.py @@ -2,16 +2,10 @@ import pytest import torch -from slaf.integrations.anndata import LazyAnnData from slaf.ml.dataloaders import SLAFDataLoader, get_device_info, get_optimal_device from slaf.ml.tokenizers import GeneformerTokenizer, ScGPTTokenizer -@pytest.fixture -def tiny_lazy_adata(tiny_slaf): - return LazyAnnData(tiny_slaf) - - def build_dataloader(adata, tokenizer_kind="geneformer", raw_mode=False, **kwargs): tokenizer_kwargs = {} for key in ("max_genes", "vocab_size", "n_expression_bins"): diff --git a/tests/test_pytorch_datasets.py b/tests/test_pytorch_datasets.py index 6244b96..938718f 100644 --- a/tests/test_pytorch_datasets.py +++ b/tests/test_pytorch_datasets.py @@ -19,9 +19,9 @@ class TestSLAFIterableDataset: """Test suite for SLAFIterableDataset with comprehensive coverage""" - def test_dataset_initialization(self, tiny_slaf): + def test_dataset_initialization(self, tiny_slaf, tiny_lazy_adata): """Test SLAFIterableDataset initialization""" - tokenizer = GeneformerTokenizer(tiny_slaf) + tokenizer = GeneformerTokenizer(tiny_lazy_adata) dataset = SLAFIterableDataset( slaf_array=tiny_slaf, @@ -41,9 +41,9 @@ def test_dataset_initialization(self, tiny_slaf): assert hasattr(dataset, "batch_processor") assert hasattr(dataset, "prefetcher") - def test_dataset_initialization_scgpt(self, tiny_slaf): + def test_dataset_initialization_scgpt(self, tiny_slaf, tiny_lazy_adata): """Test SLAFIterableDataset initialization with scGPT tokenizer""" - tokenizer = ScGPTTokenizer(tiny_slaf) + tokenizer = ScGPTTokenizer(tiny_lazy_adata) dataset = SLAFIterableDataset( slaf_array=tiny_slaf, @@ -57,10 +57,10 @@ def test_dataset_initialization_scgpt(self, tiny_slaf): assert dataset.batch_size == 16 assert dataset.seed == 123 - def test_prefetch_batch_processor_initialization(self, tiny_slaf): + def test_prefetch_batch_processor_initialization(self, tiny_slaf, tiny_lazy_adata): """Test PrefetchBatchProcessor initialization and configuration""" shuffle = RandomShuffle() - tokenizer = ScGPTTokenizer(tiny_slaf) + tokenizer = ScGPTTokenizer(tiny_lazy_adata) processor = PrefetchBatchProcessor( slaf_array=tiny_slaf, @@ -82,10 +82,12 @@ def test_prefetch_batch_processor_initialization(self, tiny_slaf): # With MoS enabled by default, we should have fragment_generators instead of batch_generator assert hasattr(processor, "fragment_generators") - def test_prefetch_batch_processor_fragment_parameter(self, tiny_slaf): + def test_prefetch_batch_processor_fragment_parameter( + self, tiny_slaf, tiny_lazy_adata + ): """Test the by_fragment parameter in PrefetchBatchProcessor.""" shuffle = RandomShuffle() - tokenizer = ScGPTTokenizer(tiny_slaf) + tokenizer = ScGPTTokenizer(tiny_lazy_adata) # Test fragment-based loading (with MoS disabled to test fragment mode) processor_fragment = PrefetchBatchProcessor( @@ -113,10 +115,10 @@ def test_prefetch_batch_processor_fragment_parameter(self, tiny_slaf): assert processor_batch.by_fragment is False - def test_prefetch_batch_processor_reset_fragment(self, tiny_slaf): + def test_prefetch_batch_processor_reset_fragment(self, tiny_slaf, tiny_lazy_adata): """Test processor epoch reset in fragment mode.""" shuffle = RandomShuffle() - tokenizer = ScGPTTokenizer(tiny_slaf) + tokenizer = ScGPTTokenizer(tiny_lazy_adata) processor = PrefetchBatchProcessor( slaf_array=tiny_slaf, @@ -133,10 +135,10 @@ def test_prefetch_batch_processor_reset_fragment(self, tiny_slaf): assert processor.current_epoch == 1 assert processor.batch_id == 0 - def test_prefetch_batch_processor_reset_batch(self, tiny_slaf): + def test_prefetch_batch_processor_reset_batch(self, tiny_slaf, tiny_lazy_adata): """Test processor epoch reset in batch mode.""" shuffle = RandomShuffle() - tokenizer = ScGPTTokenizer(tiny_slaf) + tokenizer = ScGPTTokenizer(tiny_lazy_adata) processor = PrefetchBatchProcessor( slaf_array=tiny_slaf, @@ -153,10 +155,10 @@ def test_prefetch_batch_processor_reset_batch(self, tiny_slaf): assert processor.current_epoch == 1 assert processor.batch_id == 0 - def test_prefetch_batch_processor_load_fragment(self, tiny_slaf): + def test_prefetch_batch_processor_load_fragment(self, tiny_slaf, tiny_lazy_adata): """Test processor batch loading in fragment mode.""" shuffle = RandomShuffle() - tokenizer = ScGPTTokenizer(tiny_slaf) + tokenizer = ScGPTTokenizer(tiny_lazy_adata) processor = PrefetchBatchProcessor( slaf_array=tiny_slaf, @@ -173,10 +175,10 @@ def test_prefetch_batch_processor_load_fragment(self, tiny_slaf): assert hasattr(batch, "attention_mask") assert hasattr(batch, "cell_integer_ids") - def test_prefetch_batch_processor_load_batch_mode(self, tiny_slaf): + def test_prefetch_batch_processor_load_batch_mode(self, tiny_slaf, tiny_lazy_adata): """Test processor batch loading in batch mode.""" shuffle = RandomShuffle() - tokenizer = ScGPTTokenizer(tiny_slaf) + tokenizer = ScGPTTokenizer(tiny_lazy_adata) processor = PrefetchBatchProcessor( slaf_array=tiny_slaf, @@ -201,6 +203,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 +225,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], @@ -236,10 +240,10 @@ def test_prefetch_batch_geneformer(self): assert batch.cell_integer_ids == [200, 201, 202] assert batch.tokenize_time == 0.05 - def test_async_prefetcher_initialization(self, tiny_slaf): + def test_async_prefetcher_initialization(self, tiny_slaf, tiny_lazy_adata): """Test AsyncPrefetcher initialization""" shuffle = RandomShuffle() - tokenizer = ScGPTTokenizer(tiny_slaf) + tokenizer = ScGPTTokenizer(tiny_lazy_adata) processor = PrefetchBatchProcessor( slaf_array=tiny_slaf, @@ -256,10 +260,10 @@ def test_async_prefetcher_initialization(self, tiny_slaf): assert prefetcher.worker_thread is None assert prefetcher.should_stop is False - def test_prefetcher_start_stop(self, tiny_slaf): + def test_prefetcher_start_stop(self, tiny_slaf, tiny_lazy_adata): """Test AsyncPrefetcher start and stop functionality""" shuffle = RandomShuffle() - tokenizer = ScGPTTokenizer(tiny_slaf) + tokenizer = ScGPTTokenizer(tiny_lazy_adata) processor = PrefetchBatchProcessor( slaf_array=tiny_slaf, @@ -279,10 +283,10 @@ def test_prefetcher_start_stop(self, tiny_slaf): prefetcher.stop() assert prefetcher.should_stop is True - def test_prefetcher_queue_operations(self, tiny_slaf): + def test_prefetcher_queue_operations(self, tiny_slaf, tiny_lazy_adata): """Test AsyncPrefetcher queue operations""" shuffle = RandomShuffle() - tokenizer = ScGPTTokenizer(tiny_slaf) + tokenizer = ScGPTTokenizer(tiny_lazy_adata) processor = PrefetchBatchProcessor( slaf_array=tiny_slaf, @@ -303,9 +307,9 @@ def test_prefetcher_queue_operations(self, tiny_slaf): assert "elapsed_time" in stats assert "cells_per_sec" in stats - def test_dataset_iteration_geneformer(self, tiny_slaf): + def test_dataset_iteration_geneformer(self, tiny_slaf, tiny_lazy_adata): """Test dataset iteration with Geneformer tokenizer""" - tokenizer = GeneformerTokenizer(tiny_slaf) + tokenizer = GeneformerTokenizer(tiny_lazy_adata) dataset = SLAFIterableDataset( slaf_array=tiny_slaf, tokenizer=tokenizer, @@ -336,9 +340,9 @@ def test_dataset_iteration_geneformer(self, tiny_slaf): assert batch_count > 0 - def test_dataset_iteration_scgpt(self, tiny_slaf): + def test_dataset_iteration_scgpt(self, tiny_slaf, tiny_lazy_adata): """Test dataset iteration with scGPT tokenizer""" - tokenizer = ScGPTTokenizer(tiny_slaf) + tokenizer = ScGPTTokenizer(tiny_lazy_adata) dataset = SLAFIterableDataset( slaf_array=tiny_slaf, tokenizer=tokenizer, @@ -372,9 +376,9 @@ def test_dataset_iteration_scgpt(self, tiny_slaf): assert batch_count > 0 - def test_scgpt_values_alignment_contract(self, tiny_slaf): + def test_scgpt_values_alignment_contract(self, tiny_slaf, tiny_lazy_adata): """scGPT batches must keep aligned dual streams.""" - tokenizer = ScGPTTokenizer(tiny_slaf, max_genes=16) + tokenizer = ScGPTTokenizer(tiny_lazy_adata, max_genes=16) dataset = SLAFIterableDataset( slaf_array=tiny_slaf, tokenizer=tokenizer, @@ -384,17 +388,23 @@ 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): + def test_device_transfer(self, tiny_slaf, tiny_lazy_adata): """Test device transfer functionality""" - tokenizer = GeneformerTokenizer(tiny_slaf) + tokenizer = GeneformerTokenizer(tiny_lazy_adata) dataset = SLAFIterableDataset( slaf_array=tiny_slaf, tokenizer=tokenizer, @@ -408,9 +418,9 @@ def test_device_transfer(self, tiny_slaf): assert batch["cell_ids"].device.type == "cpu" break # Just test first batch - def test_prefetcher_timeout(self, tiny_slaf): + def test_prefetcher_timeout(self, tiny_slaf, tiny_lazy_adata): """Test prefetcher timeout handling""" - tokenizer = GeneformerTokenizer(tiny_slaf) + tokenizer = GeneformerTokenizer(tiny_lazy_adata) dataset = SLAFIterableDataset( slaf_array=tiny_slaf, tokenizer=tokenizer, @@ -427,10 +437,10 @@ def test_prefetcher_timeout(self, tiny_slaf): assert batch_count > 0 - def test_pytorch_dataloader_integration(self, tiny_slaf): + def test_pytorch_dataloader_integration(self, tiny_slaf, tiny_lazy_adata): """Test integration with PyTorch DataLoader""" - tokenizer = GeneformerTokenizer(tiny_slaf) + tokenizer = GeneformerTokenizer(tiny_lazy_adata) dataset = SLAFIterableDataset( slaf_array=tiny_slaf, tokenizer=tokenizer, @@ -468,10 +478,10 @@ def test_pytorch_dataloader_integration(self, tiny_slaf): assert batch_count > 0 - def test_window_strategy_integration(self, tiny_slaf): + def test_window_strategy_integration(self, tiny_slaf, tiny_lazy_adata): """Test window strategy integration with batch processor""" shuffle = RandomShuffle() - tokenizer = GeneformerTokenizer(tiny_slaf) + tokenizer = GeneformerTokenizer(tiny_lazy_adata) processor = PrefetchBatchProcessor( slaf_array=tiny_slaf, @@ -484,9 +494,9 @@ def test_window_strategy_integration(self, tiny_slaf): assert processor.shuffle is shuffle assert processor.tokenizer.max_genes == 2048 - def test_multi_epoch_initialization(self, tiny_slaf): + def test_multi_epoch_initialization(self, tiny_slaf, tiny_lazy_adata): """Test SLAFIterableDataset initialization with multi-epoch support""" - tokenizer = GeneformerTokenizer(tiny_slaf) + tokenizer = GeneformerTokenizer(tiny_lazy_adata) dataset = SLAFIterableDataset( slaf_array=tiny_slaf, @@ -501,9 +511,9 @@ def test_multi_epoch_initialization(self, tiny_slaf): assert dataset.batch_processor.current_epoch >= 0 assert dataset.batch_processor.current_epoch < 5 - def test_multi_epoch_iteration(self, tiny_slaf): + def test_multi_epoch_iteration(self, tiny_slaf, tiny_lazy_adata): """Test multi-epoch iteration functionality""" - tokenizer = GeneformerTokenizer(tiny_slaf) + tokenizer = GeneformerTokenizer(tiny_lazy_adata) dataset = SLAFIterableDataset( slaf_array=tiny_slaf, @@ -535,9 +545,9 @@ def test_multi_epoch_iteration(self, tiny_slaf): assert len(epochs_seen) >= 1, f"Expected at least 1 epoch, got {epochs_seen}" assert total_batches > 0 - def test_epoch_transition_functionality(self, tiny_slaf): + def test_epoch_transition_functionality(self, tiny_slaf, tiny_lazy_adata): """Test that epoch transitions work correctly""" - tokenizer = GeneformerTokenizer(tiny_slaf) + tokenizer = GeneformerTokenizer(tiny_lazy_adata) dataset = SLAFIterableDataset( slaf_array=tiny_slaf, @@ -571,10 +581,10 @@ def test_epoch_transition_functionality(self, tiny_slaf): "Epochs should not go backwards" ) - def test_batch_processor_epoch_reset(self, tiny_slaf): + def test_batch_processor_epoch_reset(self, tiny_slaf, tiny_lazy_adata): """Test PrefetchBatchProcessor epoch reset functionality""" shuffle = RandomShuffle() - tokenizer = GeneformerTokenizer(tiny_slaf) + tokenizer = GeneformerTokenizer(tiny_lazy_adata) processor = PrefetchBatchProcessor( slaf_array=tiny_slaf, @@ -600,9 +610,9 @@ def test_batch_processor_epoch_reset(self, tiny_slaf): with pytest.raises(ValueError): processor.reset_for_epoch(3) # Invalid epoch (>= n_epochs) - def test_multi_epoch_with_small_dataset(self, tiny_slaf): + def test_multi_epoch_with_small_dataset(self, tiny_slaf, tiny_lazy_adata): """Test multi-epoch functionality with small dataset that gets exhausted""" - tokenizer = GeneformerTokenizer(tiny_slaf) + tokenizer = GeneformerTokenizer(tiny_lazy_adata) dataset = SLAFIterableDataset( slaf_array=tiny_slaf, @@ -632,9 +642,9 @@ def test_multi_epoch_with_small_dataset(self, tiny_slaf): ) assert total_batches > 0 - def test_single_epoch_behavior(self, tiny_slaf): + def test_single_epoch_behavior(self, tiny_slaf, tiny_lazy_adata): """Test that single epoch (default) behavior is unchanged""" - tokenizer = GeneformerTokenizer(tiny_slaf) + tokenizer = GeneformerTokenizer(tiny_lazy_adata) dataset = SLAFIterableDataset( slaf_array=tiny_slaf, @@ -658,9 +668,9 @@ def test_single_epoch_behavior(self, tiny_slaf): assert epochs_seen == {0}, f"Expected only epoch 0, got {epochs_seen}" assert batch_count > 0 - def test_multi_epoch_prefetcher_stats(self, tiny_slaf): + def test_multi_epoch_prefetcher_stats(self, tiny_slaf, tiny_lazy_adata): """Test that AsyncPrefetcher correctly tracks multi-epoch statistics""" - tokenizer = GeneformerTokenizer(tiny_slaf) + tokenizer = GeneformerTokenizer(tiny_lazy_adata) dataset = SLAFIterableDataset( slaf_array=tiny_slaf, @@ -683,9 +693,9 @@ def test_multi_epoch_prefetcher_stats(self, tiny_slaf): assert stats["n_epochs"] == 3 assert stats["current_epoch"] >= 0 - def test_multi_epoch_completion_detection(self, tiny_slaf): + def test_multi_epoch_completion_detection(self, tiny_slaf, tiny_lazy_adata): """Test that the dataset correctly detects when all epochs are completed""" - tokenizer = GeneformerTokenizer(tiny_slaf) + tokenizer = GeneformerTokenizer(tiny_lazy_adata) dataset = SLAFIterableDataset( slaf_array=tiny_slaf, @@ -768,6 +778,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], @@ -782,10 +793,10 @@ def test_prefetch_batch_serialization(self): assert len(batch.cell_integer_ids) == 2 assert batch.tokenize_time == 0.1 - def test_batch_processor_expression_binning(self, tiny_slaf): + def test_batch_processor_expression_binning(self, tiny_slaf, tiny_lazy_adata): """Test batch processor with expression binning""" shuffle = RandomShuffle() - tokenizer = ScGPTTokenizer(tiny_slaf) + tokenizer = ScGPTTokenizer(tiny_lazy_adata) processor = PrefetchBatchProcessor( slaf_array=tiny_slaf, @@ -798,9 +809,9 @@ def test_batch_processor_expression_binning(self, tiny_slaf): assert processor.n_expression_bins == 10 assert processor.use_binned_expressions is True - def test_dataset_fragment_parameter(self, tiny_slaf): + def test_dataset_fragment_parameter(self, tiny_slaf, tiny_lazy_adata): """Test the by_fragment parameter functionality.""" - tokenizer = GeneformerTokenizer(tiny_slaf) + tokenizer = GeneformerTokenizer(tiny_lazy_adata) # Test fragment-based loading dataset_fragment = SLAFIterableDataset( @@ -824,9 +835,9 @@ def test_dataset_fragment_parameter(self, tiny_slaf): assert dataset_batch.by_fragment is False - def test_dataset_iteration_fragment_mode(self, tiny_slaf): + def test_dataset_iteration_fragment_mode(self, tiny_slaf, tiny_lazy_adata): """Test dataset iteration in fragment mode.""" - tokenizer = GeneformerTokenizer(tiny_slaf) + tokenizer = GeneformerTokenizer(tiny_lazy_adata) dataset = SLAFIterableDataset( slaf_array=tiny_slaf, tokenizer=tokenizer, @@ -847,9 +858,9 @@ def test_dataset_iteration_fragment_mode(self, tiny_slaf): assert batch_count > 0 - def test_dataset_iteration_batch_mode(self, tiny_slaf): + def test_dataset_iteration_batch_mode(self, tiny_slaf, tiny_lazy_adata): """Test dataset iteration in batch mode.""" - tokenizer = GeneformerTokenizer(tiny_slaf) + tokenizer = GeneformerTokenizer(tiny_lazy_adata) dataset = SLAFIterableDataset( slaf_array=tiny_slaf, tokenizer=tokenizer, @@ -870,7 +881,7 @@ def test_dataset_iteration_batch_mode(self, tiny_slaf): assert batch_count > 0 - def test_dataset_raw_mode_fragment(self, tiny_slaf): + def test_dataset_raw_mode_fragment(self, tiny_slaf, tiny_lazy_adata): """Test dataset in raw mode with fragment loading.""" dataset = SLAFIterableDataset( slaf_array=tiny_slaf, @@ -892,7 +903,7 @@ def test_dataset_raw_mode_fragment(self, tiny_slaf): assert batch_count > 0 - def test_dataset_raw_mode_batch(self, tiny_slaf): + def test_dataset_raw_mode_batch(self, tiny_slaf, tiny_lazy_adata): """Test dataset in raw mode with batch loading.""" dataset = SLAFIterableDataset( slaf_array=tiny_slaf, @@ -914,9 +925,9 @@ def test_dataset_raw_mode_batch(self, tiny_slaf): assert batch_count > 0 - def test_mixture_of_scanners_initialization(self, tiny_slaf): + def test_mixture_of_scanners_initialization(self, tiny_slaf, tiny_lazy_adata): """Test SLAFIterableDataset initialization with Mixture of Scanners (MoS)""" - tokenizer = GeneformerTokenizer(tiny_slaf) + tokenizer = GeneformerTokenizer(tiny_lazy_adata) dataset = SLAFIterableDataset( slaf_array=tiny_slaf, @@ -932,9 +943,9 @@ def test_mixture_of_scanners_initialization(self, tiny_slaf): assert dataset.batch_processor.n_scanners == 8 assert dataset.batch_processor.prefetch_batch_size == 1048576 - def test_mixture_of_scanners_fragment_generators(self, tiny_slaf): + def test_mixture_of_scanners_fragment_generators(self, tiny_slaf, tiny_lazy_adata): """Test that MoS creates the correct number of fragment generators""" - tokenizer = GeneformerTokenizer(tiny_slaf) + tokenizer = GeneformerTokenizer(tiny_lazy_adata) dataset = SLAFIterableDataset( slaf_array=tiny_slaf, @@ -959,9 +970,9 @@ def test_mixture_of_scanners_fragment_generators(self, tiny_slaf): dataset.batch_processor.fragment_generators ) - def test_mixture_of_scanners_parameter_validation(self, tiny_slaf): + def test_mixture_of_scanners_parameter_validation(self, tiny_slaf, tiny_lazy_adata): """Test MoS parameter validation""" - tokenizer = GeneformerTokenizer(tiny_slaf) + tokenizer = GeneformerTokenizer(tiny_lazy_adata) # Test valid parameters dataset = SLAFIterableDataset( @@ -1022,9 +1033,9 @@ def test_mixture_of_scanners_parameter_validation(self, tiny_slaf): prefetch_batch_size=10000001, ) - def test_mixture_of_scanners_iteration(self, tiny_slaf): + def test_mixture_of_scanners_iteration(self, tiny_slaf, tiny_lazy_adata): """Test that MoS dataset can iterate through batches""" - tokenizer = GeneformerTokenizer(tiny_slaf) + tokenizer = GeneformerTokenizer(tiny_lazy_adata) dataset = SLAFIterableDataset( slaf_array=tiny_slaf, @@ -1047,9 +1058,9 @@ def test_mixture_of_scanners_iteration(self, tiny_slaf): assert batch_count > 0 - def test_mixture_of_scanners_epoch_reset(self, tiny_slaf): + def test_mixture_of_scanners_epoch_reset(self, tiny_slaf, tiny_lazy_adata): """Test that MoS epoch reset works correctly""" - tokenizer = GeneformerTokenizer(tiny_slaf) + tokenizer = GeneformerTokenizer(tiny_lazy_adata) dataset = SLAFIterableDataset( slaf_array=tiny_slaf, @@ -1075,10 +1086,12 @@ def test_mixture_of_scanners_epoch_reset(self, tiny_slaf): dataset.batch_processor.fragment_generators ) - def test_mixture_of_scanners_backward_compatibility(self, tiny_slaf): + def test_mixture_of_scanners_backward_compatibility( + self, tiny_slaf, tiny_lazy_adata + ): """Test that MoS is backward compatible (enabled by default) in PrefetchBatchProcessor""" shuffle = RandomShuffle() - tokenizer = ScGPTTokenizer(tiny_slaf) + tokenizer = ScGPTTokenizer(tiny_lazy_adata) # Default behavior (MoS enabled) processor_default = PrefetchBatchProcessor( @@ -1104,7 +1117,7 @@ def test_mixture_of_scanners_backward_compatibility(self, tiny_slaf): assert not hasattr(processor_disabled, "fragment_generators") assert hasattr(processor_disabled, "batch_generator") - def test_mixture_of_scanners_with_raw_mode(self, tiny_slaf): + def test_mixture_of_scanners_with_raw_mode(self, tiny_slaf, tiny_lazy_adata): """Test MoS functionality with raw mode""" dataset = SLAFIterableDataset( slaf_array=tiny_slaf, @@ -1130,9 +1143,9 @@ def test_mixture_of_scanners_with_raw_mode(self, tiny_slaf): assert batch_count > 0 - def test_mixture_of_scanners_with_fragment_mode(self, tiny_slaf): + def test_mixture_of_scanners_with_fragment_mode(self, tiny_slaf, tiny_lazy_adata): """Test that MoS automatically enables fragment mode""" - tokenizer = GeneformerTokenizer(tiny_slaf) + tokenizer = GeneformerTokenizer(tiny_lazy_adata) dataset = SLAFIterableDataset( slaf_array=tiny_slaf, @@ -1147,9 +1160,9 @@ def test_mixture_of_scanners_with_fragment_mode(self, tiny_slaf): assert dataset.batch_processor.by_fragment is True assert dataset.batch_processor.use_mixture_of_scanners is True - def test_mixture_of_scanners_random_sampling(self, tiny_slaf): + def test_mixture_of_scanners_random_sampling(self, tiny_slaf, tiny_lazy_adata): """Test that MoS uses random sampling from fragment generators""" - tokenizer = GeneformerTokenizer(tiny_slaf) + tokenizer = GeneformerTokenizer(tiny_lazy_adata) dataset = SLAFIterableDataset( slaf_array=tiny_slaf, @@ -1175,9 +1188,11 @@ def test_mixture_of_scanners_random_sampling(self, tiny_slaf): assert "attention_mask" in batch assert "cell_ids" in batch - def test_mixture_of_scanners_generator_exhaustion_handling(self, tiny_slaf): + def test_mixture_of_scanners_generator_exhaustion_handling( + self, tiny_slaf, tiny_lazy_adata + ): """Test that MoS handles generator exhaustion correctly""" - tokenizer = GeneformerTokenizer(tiny_slaf) + tokenizer = GeneformerTokenizer(tiny_lazy_adata) dataset = SLAFIterableDataset( slaf_array=tiny_slaf, @@ -1197,9 +1212,11 @@ def test_mixture_of_scanners_generator_exhaustion_handling(self, tiny_slaf): assert batch_count > 0 - def test_mixture_of_scanners_cell_boundary_handling(self, tiny_slaf): + def test_mixture_of_scanners_cell_boundary_handling( + self, tiny_slaf, tiny_lazy_adata + ): """Test that MoS handles cell boundaries correctly""" - tokenizer = GeneformerTokenizer(tiny_slaf) + tokenizer = GeneformerTokenizer(tiny_lazy_adata) dataset = SLAFIterableDataset( slaf_array=tiny_slaf, @@ -1224,11 +1241,11 @@ def test_mixture_of_scanners_cell_boundary_handling(self, tiny_slaf): assert batch_count > 0 def test_prefetch_batch_processor_mixture_of_scanners_initialization( - self, tiny_slaf + self, tiny_slaf, tiny_lazy_adata ): """Test PrefetchBatchProcessor initialization with Mixture of Scanners (MoS)""" shuffle = RandomShuffle() - tokenizer = ScGPTTokenizer(tiny_slaf) + tokenizer = ScGPTTokenizer(tiny_lazy_adata) processor = PrefetchBatchProcessor( slaf_array=tiny_slaf, @@ -1255,10 +1272,12 @@ def test_prefetch_batch_processor_mixture_of_scanners_initialization( assert len(processor.generator_last_cells) == len(processor.fragment_generators) assert len(processor.generator_active) == len(processor.generator_active) - def test_prefetch_batch_processor_mos_parameter_validation(self, tiny_slaf): + def test_prefetch_batch_processor_mos_parameter_validation( + self, tiny_slaf, tiny_lazy_adata + ): """Test MoS parameter validation in PrefetchBatchProcessor""" shuffle = RandomShuffle() - tokenizer = ScGPTTokenizer(tiny_slaf) + tokenizer = ScGPTTokenizer(tiny_lazy_adata) # Test valid parameters processor = PrefetchBatchProcessor( @@ -1319,10 +1338,10 @@ def test_prefetch_batch_processor_mos_parameter_validation(self, tiny_slaf): prefetch_batch_size=10000001, ) - def test_prefetch_batch_processor_mos_epoch_reset(self, tiny_slaf): + def test_prefetch_batch_processor_mos_epoch_reset(self, tiny_slaf, tiny_lazy_adata): """Test that MoS epoch reset works correctly in PrefetchBatchProcessor""" shuffle = RandomShuffle() - tokenizer = ScGPTTokenizer(tiny_slaf) + tokenizer = ScGPTTokenizer(tiny_lazy_adata) processor = PrefetchBatchProcessor( slaf_array=tiny_slaf, @@ -1344,10 +1363,12 @@ def test_prefetch_batch_processor_mos_epoch_reset(self, tiny_slaf): assert len(processor.generator_last_cells) == len(processor.fragment_generators) assert len(processor.generator_active) == len(processor.fragment_generators) - def test_prefetch_batch_processor_mos_backward_compatibility(self, tiny_slaf): + def test_prefetch_batch_processor_mos_backward_compatibility( + self, tiny_slaf, tiny_lazy_adata + ): """Test that MoS is backward compatible (enabled by default) in PrefetchBatchProcessor""" shuffle = RandomShuffle() - tokenizer = ScGPTTokenizer(tiny_slaf) + tokenizer = ScGPTTokenizer(tiny_lazy_adata) # Default behavior (MoS enabled) processor_default = PrefetchBatchProcessor( @@ -1373,7 +1394,9 @@ def test_prefetch_batch_processor_mos_backward_compatibility(self, tiny_slaf): assert not hasattr(processor_disabled, "fragment_generators") assert hasattr(processor_disabled, "batch_generator") - def test_prefetch_batch_processor_mos_with_raw_mode(self, tiny_slaf): + def test_prefetch_batch_processor_mos_with_raw_mode( + self, tiny_slaf, tiny_lazy_adata + ): """Test MoS functionality with raw mode in PrefetchBatchProcessor""" shuffle = RandomShuffle() tokenizer = None # No tokenizer for raw mode @@ -1396,10 +1419,12 @@ def test_prefetch_batch_processor_mos_with_raw_mode(self, tiny_slaf): assert hasattr(batch, "batch_dfs") assert hasattr(batch, "cell_integer_ids") - def test_prefetch_batch_processor_mos_fragment_mode_automatic(self, tiny_slaf): + def test_prefetch_batch_processor_mos_fragment_mode_automatic( + self, tiny_slaf, tiny_lazy_adata + ): """Test that MoS automatically enables fragment mode in PrefetchBatchProcessor""" shuffle = RandomShuffle() - tokenizer = ScGPTTokenizer(tiny_slaf) + tokenizer = ScGPTTokenizer(tiny_lazy_adata) processor = PrefetchBatchProcessor( slaf_array=tiny_slaf, @@ -1414,10 +1439,12 @@ def test_prefetch_batch_processor_mos_fragment_mode_automatic(self, tiny_slaf): assert processor.by_fragment is True assert processor.use_mixture_of_scanners is True - def test_prefetch_batch_processor_mos_load_prefetch_batch(self, tiny_slaf): + def test_prefetch_batch_processor_mos_load_prefetch_batch( + self, tiny_slaf, tiny_lazy_adata + ): """Test that MoS can load prefetch batches correctly""" shuffle = RandomShuffle() - tokenizer = ScGPTTokenizer(tiny_slaf) + tokenizer = ScGPTTokenizer(tiny_lazy_adata) processor = PrefetchBatchProcessor( slaf_array=tiny_slaf, @@ -1434,10 +1461,12 @@ def test_prefetch_batch_processor_mos_load_prefetch_batch(self, tiny_slaf): assert hasattr(batch, "attention_mask") assert hasattr(batch, "cell_integer_ids") - def test_prefetch_batch_processor_mos_cell_boundary_handling(self, tiny_slaf): + def test_prefetch_batch_processor_mos_cell_boundary_handling( + self, tiny_slaf, tiny_lazy_adata + ): """Test that MoS handles cell boundaries correctly in PrefetchBatchProcessor""" shuffle = RandomShuffle() - tokenizer = ScGPTTokenizer(tiny_slaf) + tokenizer = ScGPTTokenizer(tiny_lazy_adata) processor = PrefetchBatchProcessor( slaf_array=tiny_slaf, diff --git a/tests/test_tokenizers.py b/tests/test_tokenizers.py index 2008d0f..c5f7602 100644 --- a/tests/test_tokenizers.py +++ b/tests/test_tokenizers.py @@ -56,7 +56,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): @@ -142,7 +141,7 @@ def test_scgpt_tokenize_grouped_uses_preencoded_tokens(self): grouped_df = pd.DataFrame( { "gene_sequence": [[4, 5, 6]], - "expr_sequence": [[1001, 1005, 1009]], + "expr_sequence": [[1, 5, 9]], } ) grouped_df = __import__("polars").from_pandas(grouped_df) @@ -151,7 +150,7 @@ def test_scgpt_tokenize_grouped_uses_preencoded_tokens(self): assert input_ids[0, 1] == 4 assert values is not None - assert values[0, 1] == 1001 + assert values[0, 1] == 1 def test_geneformer_tokenize_grouped_uses_preencoded_tokens(self): mock_adata, _ = build_mock_adata() @@ -228,18 +227,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): @@ -314,7 +313,7 @@ def test_apply_uses_slaf_runtime_transformations(self): } ) - grouped = tokenizer.apply( + grouped = tokenizer.transform_and_apply( df, schema=SLAF_LANCE_COO_SCHEMA, max_items=2, @@ -356,7 +355,7 @@ def test_apply_requires_precomputed_cell_factors_for_normalize_total(self): ValueError, match="requires precomputed cell_factors keyed by cell_integer_id", ): - tokenizer.apply( + tokenizer.transform_and_apply( df, schema=SLAF_LANCE_COO_SCHEMA, max_items=2, @@ -387,7 +386,7 @@ def test_apply_uses_precomputed_cell_factors_by_integer_id(self): } ) - grouped = tokenizer.apply( + grouped = tokenizer.transform_and_apply( df, schema=SLAF_LANCE_COO_SCHEMA, max_items=2, @@ -422,7 +421,7 @@ def test_scgpt_window_does_not_double_log1p(self): } ) - grouped = tokenizer.apply( + grouped = tokenizer.transform_and_apply( df, schema=SLAF_LANCE_COO_SCHEMA, max_items=2, @@ -430,7 +429,7 @@ def test_scgpt_window_does_not_double_log1p(self): ) expr_tokens = grouped["expr_sequence"].to_list()[0] - assert expr_tokens == [1005, 1009] + assert expr_tokens == [6, 10] def test_token_decoding(self): """Test token decoding functionality.""" @@ -459,9 +458,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 cdfc0d9d988f92a45f75c6c3ada2dd7fc587b3a0 Mon Sep 17 00:00:00 2001 From: Shahul Alam Date: Mon, 25 May 2026 00:25:59 -0400 Subject: [PATCH 07/10] Improve tokenizer abstraction. Remove scGPT-specific tokenizer properties. --- slaf/ml/aggregators.py | 5 +---- slaf/ml/datasets.py | 15 +++++---------- slaf/ml/distributed.py | 4 ---- slaf/ml/tokenizers.py | 5 +++++ 4 files changed, 11 insertions(+), 18 deletions(-) diff --git a/slaf/ml/aggregators.py b/slaf/ml/aggregators.py index c2ea93b..1aae481 100644 --- a/slaf/ml/aggregators.py +++ b/slaf/ml/aggregators.py @@ -12,10 +12,7 @@ import polars as pl from slaf.core.tabular_schema import DataSchema -from slaf.ml.expression_preprocessor import ( - ExpressionPreprocessor, - apply_expression_preprocessor, -) +from slaf.ml.expression_preprocessor import apply_expression_preprocessor class Window(ABC): diff --git a/slaf/ml/datasets.py b/slaf/ml/datasets.py index 997487b..bc67c5c 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"] @@ -1029,13 +1029,8 @@ 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.update( - self.window_kwargs - ) # Add any additional kwargs + window_params: dict[str, Any] = {} + window_params.update(self.window_kwargs) if self.expression_preprocessor is not None: window_params["expression_preprocessor"] = ( self.expression_preprocessor @@ -1582,8 +1577,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 6c3eba1..1201c83 100644 --- a/slaf/ml/distributed.py +++ b/slaf/ml/distributed.py @@ -310,10 +310,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 bfe654e..20508e9 100644 --- a/slaf/ml/tokenizers.py +++ b/slaf/ml/tokenizers.py @@ -418,6 +418,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. @@ -431,6 +432,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. @@ -460,6 +463,7 @@ def __init__( """ self.n_expression_bins = n_expression_bins + self.use_binned_expressions = use_binned_expressions super().__init__(adata=adata, vocab_size=vocab_size, max_genes=max_genes) def create_window(self) -> Window: @@ -474,6 +478,7 @@ def apply( ) -> pl.DataFrame: kwargs.setdefault("special_token_offset", 4) kwargs.setdefault("n_expression_bins", self.n_expression_bins) + kwargs.setdefault("use_binned_expressions", self.use_binned_expressions) return self.window.apply(df, schema=schema, max_items=max_items, **kwargs) def tokenize_grouped( From e24d8b98e0f6c24d028554609c9af79f8af2fb4a Mon Sep 17 00:00:00 2001 From: Pavan Ramkumar Date: Thu, 28 May 2026 12:25:31 -0700 Subject: [PATCH 08/10] Store expression_preprocessor on SLAFDataLoader --- slaf/ml/dataloaders.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/slaf/ml/dataloaders.py b/slaf/ml/dataloaders.py index e5cdb7a..1fe0178 100644 --- a/slaf/ml/dataloaders.py +++ b/slaf/ml/dataloaders.py @@ -5,7 +5,7 @@ from slaf.integrations.anndata import LazyAnnData from .expression_preprocessor import ExpressionPreprocessor -from .tokenizers import GeneformerTokenizer, ScGPTTokenizer, SLAFTokenizer +from .tokenizers import SLAFTokenizer # Try to import torch, but make it optional try: @@ -423,6 +423,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: @@ -482,7 +483,7 @@ def _create_dataset(self, prefetcher_ready_timeout: float) -> "SLAFIterableDatas 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, + expression_preprocessor=self.expression_preprocessor, ) def __iter__(self): From 8c674541f0c77c47c611adcfed1fe428aac6aeba Mon Sep 17 00:00:00 2001 From: Pavan Ramkumar Date: Thu, 28 May 2026 12:42:08 -0700 Subject: [PATCH 09/10] Revert tokenizer LazyAnnData and in-tokenizer preprocessing. Pass SLAFArray to tokenizers again and apply normalize/log1p upstream via ExpressionPreprocessor in window aggregators. Restore dataloader and worker APIs from post-#57; keep epoch field on prefetch batches for reporting. Co-authored-by: Cursor --- pyproject.toml | 7 +- slaf/distributed/worker.py | 159 ++++++++++------- slaf/integrations/anndata.py | 71 ++++++-- slaf/integrations/scanpy.py | 88 +++++++-- slaf/ml/aggregators.py | 26 ++- slaf/ml/dataloaders.py | 154 ++++++++-------- slaf/ml/datasets.py | 31 ++-- slaf/ml/distributed.py | 136 ++++++++------ slaf/ml/tokenizers.py | 282 ++++------------------------- tests/conftest.py | 6 - tests/test_aggregators.py | 8 +- tests/test_dataloaders.py | 280 +++++++++++++---------------- tests/test_distributed_worker.py | 177 ++++++++++++------ tests/test_pytorch_datasets.py | 230 +++++++++++------------- tests/test_tokenizers.py | 298 +++++++------------------------ uv.lock | 52 ++---- 16 files changed, 890 insertions(+), 1115 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 7740fd1..17babad 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -37,12 +37,7 @@ dependencies = [ [project.optional-dependencies] # Machine learning - for ML modules -ml = [ - "torch>=2.5.0", - "tiledb>=0.34.2", - "tiledbsoma>=1.17.1", - "omegaconf>=2.3.0", -] +ml = ["torch>=2.5.0", "tiledb>=0.34.2", "tiledbsoma>=1.17.1"] # 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 98febb5..74b012d 100644 --- a/slaf/distributed/worker.py +++ b/slaf/distributed/worker.py @@ -5,6 +5,7 @@ import asyncio import importlib +import inspect import pickle import queue as queue_module import random @@ -16,7 +17,6 @@ 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 @@ -26,8 +26,8 @@ def prefetch_worker( worker_id: str, partition_indices: list[int], - data_source_config: DictConfig, - processor_config: DictConfig, + data_source_config: dict[str, Any], + processor_config: dict[str, Any], queue: Any, # Modal Queue or any queue-like object n_scanners: int = 8, prefetch_batch_count: int = 32, @@ -56,32 +56,23 @@ 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 - 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 @@ -105,39 +96,71 @@ def prefetch_worker( # Built before window when use_tokenizer_window is set so tokenizer_instance.window exists. tokenizer_instance = None tokenizer_fn = None - 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 + 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 - from slaf.integrations.anndata import LazyAnnData slaf_array = SLAFArray(slaf_path, load_metadata=False) - adata = LazyAnnData(slaf_array) - - 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 {} + + # Create tokenizer instance + tokenizer_instance = tokenizer_class( + slaf_array=slaf_array, **tokenizer_config["kwargs"] ) - tokenizer_instance = tokenizer_class(adata=adata, **tokenizer_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 tokenizer-owned grouping contract.""" - input_ids, attention_mask, values = tokenizer_instance.tokenize_grouped( - grouped_df, - schema=schema, - ) + """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() + 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 + ): + 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 + ) + + 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 = { @@ -151,46 +174,56 @@ def tokenize_grouped( tokenizer_fn = tokenize_grouped # Shuffle: use factory if provided, otherwise use default - if shuffle_factory_config: + if processor_config.get("shuffle_factory"): # Dynamic import based on config - module path comes from config, not hardcoded - factory_config = shuffle_factory_config + factory_config = processor_config["shuffle_factory"] 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_kwargs = ( - dict(factory_config.kwargs) - if "kwargs" in factory_config and factory_config.kwargs is not None - else {} + shuffle_factory = getattr(shuffle_module, factory_config["function"]) + shuffle = shuffle_factory( + factory_config["type"], **factory_config.get("kwargs", {}) ) - 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 - elif window_factory_config: + 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 = window_factory_config + factory_config = processor_config["window_factory"] 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_factory_kwargs = ( - dict(factory_config.kwargs) - if "kwargs" in factory_config and factory_config.kwargs is not None - else {} + window_factory = getattr(window_module, factory_config["function"]) + window = window_factory( + factory_config["type"], **factory_config.get("kwargs", {}) ) - 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(**OmegaConf.to_container(processor_config.schema, resolve=True)) + schema = DataSchema(**processor_config["schema"]) # Create boundary handler with KV store support # partial_groups_kv is passed as a parameter (created by caller) @@ -198,7 +231,7 @@ def tokenize_grouped( boundary_handler = GroupBoundaryHandler( schema=schema, - continuity_check=continuity_check, + continuity_check=processor_config.get("continuity_check", "sequential"), partial_groups_kv=partial_groups_kv, ) @@ -208,10 +241,10 @@ def tokenize_grouped( shuffle=shuffle, tokenizer=tokenizer_fn, boundary_handler=boundary_handler, - max_items=max_items, - seed=seed_value, - continuity_check=continuity_check, - **dict(window_kwargs), + 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", {}), ) # Queue is passed as a parameter (created by caller) @@ -297,7 +330,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.n_epochs + epochs = processor_config.get("n_epochs", 1) # 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/integrations/anndata.py b/slaf/integrations/anndata.py index 5227f0a..0530d16 100644 --- a/slaf/integrations/anndata.py +++ b/slaf/integrations/anndata.py @@ -616,25 +616,68 @@ def _apply_normalize_total( ) -> scipy.sparse.csr_matrix: """Apply normalize_total transformation using vectorized operations""" cell_factors = transform_data.get("cell_factors", {}) - - if isinstance(cell_selector, int | np.integer): - selected_cell_integer_ids = [int(cell_selector)] + obs_names_local = [] # Always a list + obs_names = None + if hasattr(self, "parent_adata") and self.parent_adata is not None: + try: + obs_names = self.parent_adata.obs_names + except (AttributeError, TypeError): + obs_names = None + if obs_names is None or not isinstance(obs_names, list | np.ndarray | pd.Index): + if ( + matrix is not None + and hasattr(matrix, "shape") + and matrix.shape is not None + ): + obs_names_local = [f"cell_{i}" for i in range(matrix.shape[0])] + else: + obs_names_local = [] + else: + obs_names_local = list(obs_names) + # Now obs_names_local is always a list + # Determine selected_cell_names based on cell_selector + if cell_selector is None or ( + isinstance(cell_selector, slice) and cell_selector == slice(None) + ): + selected_cell_names = obs_names_local elif isinstance(cell_selector, slice): - selected_cell_integer_ids = list(range(self.slaf_array.shape[0]))[ - cell_selector - ] - elif isinstance(cell_selector, np.ndarray) and cell_selector.dtype == bool: - selected_cell_integer_ids = ( - np.flatnonzero(cell_selector).astype(int).tolist() - ) + start = cell_selector.start or 0 + stop = cell_selector.stop or len(obs_names_local) + step = cell_selector.step or 1 + # Clamp bounds + start = max(0, min(start, len(obs_names_local))) + stop = max(0, min(stop, len(obs_names_local))) + selected_cell_names = obs_names_local[start:stop:step] elif isinstance(cell_selector, list | np.ndarray): - selected_cell_integer_ids = [int(i) for i in cell_selector] + if len(obs_names_local) > 0: + if ( + isinstance(cell_selector, np.ndarray) + and cell_selector.dtype == bool + ): + selected_cell_names = [ + obs_names_local[i] + for i, keep in enumerate(cell_selector) + if keep and 0 <= i < len(obs_names_local) + ] + else: + selected_cell_names = [ + obs_names_local[i] + for i in cell_selector + if isinstance(i, int | np.integer) + and 0 <= i < len(obs_names_local) + ] + else: + selected_cell_names = [] + elif isinstance(cell_selector, int | np.integer): + if 0 <= cell_selector < len(obs_names_local): + selected_cell_names = [obs_names_local[cell_selector]] + else: + selected_cell_names = [] else: - selected_cell_integer_ids = list(range(matrix.shape[0])) - + selected_cell_names = obs_names_local # Create a vector of factors for all cells at once cell_factors_vector = np.array( - [cell_factors.get(cell_id, 1.0) for cell_id in selected_cell_integer_ids] + [cell_factors.get(name, 1.0) for name in selected_cell_names] ) # Apply vectorized scaling using CSR matrix properties # Create a copy only if we need to modify the data diff --git a/slaf/integrations/scanpy.py b/slaf/integrations/scanpy.py index 16401d4..c672edf 100644 --- a/slaf/integrations/scanpy.py +++ b/slaf/integrations/scanpy.py @@ -1,7 +1,6 @@ import numpy as np import pandas as pd import polars as pl -from loguru import logger from slaf.integrations.anndata import LazyAnnData @@ -465,7 +464,7 @@ def filtered_obs() -> pd.DataFrame: if inplace: # Apply filter to adata (would need proper implementation) # For now, just return the original adata - logger.info( + print( f"Filtered out {np.sum(~cell_mask)} cells, {np.sum(cell_mask)} remaining" ) return None @@ -600,7 +599,7 @@ def filter_genes( if inplace: # Apply filter to adata (would need proper implementation) - logger.info( + print( f"Filtered out {np.sum(~gene_mask)} genes, {np.sum(gene_mask)} remaining" ) return None @@ -707,7 +706,7 @@ def normalize_total( ) except Exception as e: - logger.warning( + print( f"Fragment processing failed, falling back to global processing: {e}" ) # Fall back to global processing @@ -730,17 +729,70 @@ def normalize_total( # Work with polars DataFrame internally cell_totals_pl = cell_totals - # Store factors keyed by SLAF's internal cell_integer_id. - cell_totals_pl = cell_totals_pl.with_columns( - (target_sum / pl.col("total_counts")).alias("normalization_factor") - ) - normalization_dict = dict( - zip( - cell_totals_pl["cell_integer_id"].to_list(), - cell_totals_pl["normalization_factor"].to_list(), - strict=False, + # Create normalization factors using polars + # Map cell_integer_id to cell names for compatibility with anndata.py + if hasattr(adata.slaf, "obs") and adata.slaf.obs is not None: + # Create mapping from cell_integer_id to cell names + # Use polars DataFrame to get cell names + obs_df = adata.slaf.obs + if "cell_id" in obs_df.columns: + cell_id_to_name = dict( + zip( + obs_df["cell_integer_id"].to_list(), + obs_df["cell_id"].to_list(), + strict=False, + ) + ) + else: + # Fallback: create cell names from integer IDs + cell_id_to_name = {i: f"cell_{i}" for i in range(len(obs_df))} + # Use vectorized polars operations for mapping + # Create mapping DataFrame + cell_map_df = pl.DataFrame( + { + "cell_integer_id": list(cell_id_to_name.keys()), + "cell_id": list(cell_id_to_name.values()), + } + ) + # Join with mapping DataFrame + cell_totals_pl = cell_totals_pl.join( + cell_map_df, on="cell_integer_id", how="left" + ) + # Fill any missing values with default format + cell_totals_pl = cell_totals_pl.with_columns( + [ + pl.col("cell_id").fill_null( + pl.col("cell_integer_id") + .cast(pl.Utf8) + .map_elements(lambda x: f"cell_{x}", return_dtype=pl.Utf8) + ), + (target_sum / pl.col("total_counts")).alias("normalization_factor"), + ] + ) + # Convert to dictionary for compatibility + normalization_dict = dict( + zip( + cell_totals_pl["cell_id"].to_list(), + cell_totals_pl["normalization_factor"].to_list(), + strict=False, + ) + ) + else: + # Fallback: use cell_integer_id as string keys + cell_totals_pl = cell_totals_pl.with_columns( + [ + pl.col("cell_integer_id").cast(pl.Utf8).alias("cell_id"), + (target_sum / pl.col("total_counts")).alias("normalization_factor"), + ] + ) + # Convert to dictionary for compatibility + normalization_dict = dict( + zip( + cell_totals_pl["cell_id"].to_list(), + cell_totals_pl["normalization_factor"].to_list(), + strict=False, + ) ) - ) if inplace: # Store normalization factors for lazy application @@ -755,7 +807,7 @@ def normalize_total( "cell_factors": normalization_dict, } - logger.info(f"Applied normalize_total with target_sum={target_sum}") + print(f"Applied normalize_total with target_sum={target_sum}") return None else: # Create a copy with the transformation (copy-on-write) @@ -847,7 +899,7 @@ def log1p( return adata._update_with_log1p_data(result_df, inplace) except Exception as e: - logger.warning( + print( f"Fragment processing failed, falling back to global processing: {e}" ) # Fall back to global processing @@ -862,7 +914,7 @@ def log1p( adata._transformations["log1p"] = {"type": "log1p", "applied": True} - logger.info("Applied log1p transformation") + print("Applied log1p transformation") return None else: # Create a copy with the transformation (copy-on-write) @@ -1048,7 +1100,7 @@ def highly_variable_genes( if inplace: # Update var metadata (would need implementation) - logger.info(f"Identified {hvg_mask.sum()} highly variable genes") + print(f"Identified {hvg_mask.sum()} highly variable genes") return None else: return gene_stats_complete diff --git a/slaf/ml/aggregators.py b/slaf/ml/aggregators.py index 1aae481..f68cb4d 100644 --- a/slaf/ml/aggregators.py +++ b/slaf/ml/aggregators.py @@ -12,7 +12,10 @@ import polars as pl from slaf.core.tabular_schema import DataSchema -from slaf.ml.expression_preprocessor import apply_expression_preprocessor +from slaf.ml.expression_preprocessor import ( + ExpressionPreprocessor, + apply_expression_preprocessor, +) class Window(ABC): @@ -87,6 +90,9 @@ def apply( preprocessor = kwargs.get("expression_preprocessor") fragment_df = apply_expression_preprocessor(fragment_df, schema, preprocessor) + already_log1p = ( + isinstance(preprocessor, ExpressionPreprocessor) and preprocessor.log1p + ) if use_binned_expressions: grouped = ( @@ -98,14 +104,20 @@ def apply( ) .filter(pl.col("gene_rank") <= max_items) .with_columns( - pl.when(pl.col(vk) > 0) + (pl.col(vk) if already_log1p else pl.col(vk).log1p()).alias( + "log_value" + ) + ) + .with_columns( + pl.when(pl.col("log_value") > 0) .then( - 1 - + ( - (pl.col(vk) * n_expression_bins / pl.col(vk).max().over(gk)) - .floor() - .clip(0, n_expression_bins - 1) + ( + pl.col("log_value") + * n_expression_bins + / pl.col("log_value").max().over(gk) ) + .floor() + .clip(0, n_expression_bins - 1) ) .otherwise(0) .alias("expr_bin") diff --git a/slaf/ml/dataloaders.py b/slaf/ml/dataloaders.py index 1fe0178..a1356c3 100644 --- a/slaf/ml/dataloaders.py +++ b/slaf/ml/dataloaders.py @@ -2,10 +2,10 @@ from loguru import logger -from slaf.integrations.anndata import LazyAnnData +from slaf.core.slaf import SLAFArray from .expression_preprocessor import ExpressionPreprocessor -from .tokenizers import SLAFTokenizer +from .tokenizers import GeneformerTokenizer, ScGPTTokenizer, SLAFTokenizer # Try to import torch, but make it optional try: @@ -188,8 +188,7 @@ class SLAFDataLoader: Examples: >>> # Basic usage with default settings (MoS loading) >>> slaf_array = SLAFArray("path/to/data.slaf") - >>> tokenizer = GeneformerTokenizer(slaf_array) - >>> dataloader = SLAFDataLoader(slaf_array, tokenizer=tokenizer) + >>> dataloader = SLAFDataLoader(slaf_array) >>> for batch in dataloader: ... print(f"Batch shape: {batch['input_ids'].shape}") ... print(f"Cell IDs: {batch['cell_ids']}") @@ -236,14 +235,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=tokenizer, + ... tokenizer_type="scgpt", ... batch_size=64, + ... max_genes=1024 ... ) >>> print(f"Tokenizer type: {dataloader.tokenizer_type}") - Tokenizer type: ScGPTTokenizer + Tokenizer type: scgpt >>> # Training loop example >>> for batch_idx, batch in enumerate(dataloader): @@ -256,12 +255,12 @@ class SLAFDataLoader: >>> print("Training loop completed") Training loop completed - >>> # Error handling for missing tokenizer + >>> # Error handling for invalid tokenizer type >>> try: - ... dataloader = SLAFDataLoader(slaf_array) + ... dataloader = SLAFDataLoader(slaf_array, tokenizer_type="invalid") ... except ValueError as e: ... print(f"Error: {e}") - Error: tokenizer must be provided unless raw_mode=True. + Error: Unsupported tokenizer type: invalid """ device: Optional["torch.device"] # type: ignore @@ -269,9 +268,12 @@ class SLAFDataLoader: def __init__( self, - adata: LazyAnnData, - tokenizer: SLAFTokenizer | None = None, + slaf_array: SLAFArray, + tokenizer_type: str = "geneformer", 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 @@ -293,8 +295,20 @@ def __init__( Must be a valid SLAFArray with proper Lance dataset structure. # Tokenization Configuration - tokenizer: Instantiated tokenizer used for tokenized mode. - Required unless raw_mode=True. + 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. # Training Configuration batch_size: Number of cells per batch. Larger batches use more memory @@ -342,7 +356,7 @@ def __init__( Default: True. Raises: - ValueError: If tokenizer is missing in tokenized mode or parameters are invalid. + ValueError: If tokenizer_type is not supported 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. @@ -357,8 +371,7 @@ def __init__( Examples: >>> # Basic initialization (MoS is now default) >>> slaf_array = SLAFArray("path/to/data.slaf") - >>> tokenizer = GeneformerTokenizer(slaf_array) - >>> dataloader = SLAFDataLoader(slaf_array, tokenizer=tokenizer) + >>> dataloader = SLAFDataLoader(slaf_array) >>> print(f"Batch size: {dataloader.batch_size}") Batch size: 32 >>> print(f"MoS enabled: {dataloader.use_mixture_of_scanners}") @@ -406,9 +419,9 @@ def __init__( ... print(f"Error: {e}") Error: n_scanners must be at least 1 """ - self.adata = adata - self.slaf_array = adata.slaf + 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 @@ -423,7 +436,6 @@ 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: @@ -450,40 +462,50 @@ def __init__( # Initialize tokenizer (only needed for non-raw mode) if not self.raw_mode: - 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 + 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 self.special_tokens = self.tokenizer.special_tokens else: - if tokenizer is not None: - raise ValueError("raw_mode=True is incompatible with tokenizer.") + # For raw mode, we don't need a tokenizer self.tokenizer = None self.special_tokens = None - self.tokenizer_type = "raw" - self.max_genes = 0 - - self._dataset = self._create_dataset(prefetcher_ready_timeout) - def _create_dataset(self, prefetcher_ready_timeout: float) -> "SLAFIterableDataset": - return SLAFIterableDataset( - slaf_array=self.slaf_array, + # Use IterableDataset + self._dataset = SLAFIterableDataset( + slaf_array=slaf_array, tokenizer=self.tokenizer, - batch_size=self.batch_size, + batch_size=batch_size, seed=42, # TODO: make configurable - 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 + 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 prefetcher_ready_timeout=prefetcher_ready_timeout, # Pass prefetcher_ready_timeout - expression_preprocessor=self.expression_preprocessor, + expression_preprocessor=expression_preprocessor, ) def __iter__(self): @@ -503,13 +525,12 @@ def __iter__(self): Yields: dict: Batch dictionary containing: - **Tokenized mode** (raw_mode=False): - - input_ids: Pre-tokenized gene identity data (torch.Tensor) - - values: Aligned expression/value stream (torch.Tensor) + - input_ids: Pre-tokenized gene expression data (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 sparse cell × gene data (polars.DataFrame) - - cell_ids: Integer IDs of cells in the batch (list[int]) + - x: Raw cell × gene data as Polars DataFrame + - cell_ids: List of cell integer IDs in the batch - **Multi-epoch** (when n_epochs > 1): - epoch: Current epoch number (int) @@ -518,24 +539,19 @@ def __iter__(self): The training loop should handle device transfer as needed. Raises: - ValueError: If tokenizer configuration is invalid. + ValueError: If the tokenizer type is not supported. RuntimeError: If batch processing fails. Examples: >>> # Basic iteration (tokenized mode) >>> slaf_array = SLAFArray("path/to/data.slaf") - >>> tokenizer = GeneformerTokenizer(slaf_array) - >>> dataloader = SLAFDataLoader( - ... slaf_array, - ... tokenizer=tokenizer, - ... batch_size=16, - ... ) + >>> dataloader = SLAFDataLoader(slaf_array, 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', 'values', 'attention_mask', 'cell_ids'] + Batch keys: ['input_ids', 'attention_mask', 'cell_ids'] Input shape: (16, 2048) Cell IDs: tensor([0, 1, 2, ..., 13, 14, 15]) @@ -565,7 +581,6 @@ 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"] @@ -581,15 +596,9 @@ def __iter__(self): Processed batch 1 Processed batch 2 - >>> # Different tokenizer instances - >>> dataloader_geneformer = SLAFDataLoader( - ... slaf_array, - ... tokenizer=GeneformerTokenizer(slaf_array), - ... ) - >>> dataloader_scgpt = SLAFDataLoader( - ... slaf_array, - ... tokenizer=ScGPTTokenizer(slaf_array), - ... ) + >>> # Different tokenizer types + >>> dataloader_geneformer = SLAFDataLoader(slaf_array, tokenizer_type="geneformer") + >>> dataloader_scgpt = SLAFDataLoader(slaf_array, tokenizer_type="scgpt") >>> >>> # Compare batch shapes >>> for batch in dataloader_geneformer: @@ -661,8 +670,7 @@ def __del__(self): >>> print("Manual cleanup completed") Manual cleanup completed """ - self.close() - - def close(self): - if hasattr(self, "_dataset") and hasattr(self._dataset, "close"): - self._dataset.close() + if hasattr(self, "_dataset"): + # The SLAFIterableDataset doesn't have a stop method, + # so we just let it finish its current epoch. + pass diff --git a/slaf/ml/datasets.py b/slaf/ml/datasets.py index bc67c5c..49eb36d 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 ScGPTTokenizer, SLAFTokenizer +from slaf.ml.tokenizers import SLAFTokenizer # Define union type for both batch types PrefetchBatch = Union["TokenizedPrefetchBatch", "RawPrefetchBatch"] @@ -1029,8 +1029,13 @@ def load_prefetch_batch(self) -> PrefetchBatch: shuffle_time = time.time() - shuffle_start window_start = time.time() - window_params: dict[str, Any] = {} - window_params.update(self.window_kwargs) + window_params: dict[str, Any] = { + "n_expression_bins": self.n_expression_bins, + "use_binned_expressions": self.use_binned_expressions, + } + window_params.update( + self.window_kwargs + ) # Add any additional kwargs if self.expression_preprocessor is not None: window_params["expression_preprocessor"] = ( self.expression_preprocessor @@ -1040,10 +1045,10 @@ def load_prefetch_batch(self) -> PrefetchBatch: if tokenizer is None: raise RuntimeError("Tokenizer is required for tokenized mode") - grouped = tokenizer.transform_and_apply( + grouped = tokenizer.window.apply( shuffled_df, - schema=SLAF_LANCE_COO_SCHEMA, - max_items=tokenizer.max_genes, + SLAF_LANCE_COO_SCHEMA, + tokenizer.max_genes, **window_params, ) window_time = time.time() - window_start @@ -1054,9 +1059,13 @@ 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_grouped( - grouped, - schema=SLAF_LANCE_COO_SCHEMA, + 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 + ), ) tokenize_time = time.time() - tokenize_start @@ -1577,8 +1586,8 @@ def __init__( tokenizer, "n_expression_bins", 10 ) # Default value for raw mode - if isinstance(tokenizer, ScGPTTokenizer): - tokenizer.use_binned_expressions = use_binned_expressions + # Set binning based on tokenizer type + use_binned_expressions = use_binned_expressions # Use parameter value self.batch_processor = PrefetchBatchProcessor( slaf_array=slaf_array, diff --git a/slaf/ml/distributed.py b/slaf/ml/distributed.py index 1201c83..accec75 100644 --- a/slaf/ml/distributed.py +++ b/slaf/ml/distributed.py @@ -10,7 +10,6 @@ import modal from loguru import logger -from omegaconf import DictConfig, OmegaConf from slaf.core.slaf import SLAFArray from slaf.core.tabular_schema import DataSchema @@ -22,7 +21,7 @@ # Import SLAF-specific components (for type hints and adapters) from slaf.ml.expression_preprocessor import ExpressionPreprocessor -from slaf.ml.tokenizers import SLAFTokenizer +from slaf.ml.tokenizers import GeneformerTokenizer, ScGPTTokenizer # Configure Modal image for SLAF workers. # Cache bust must run *before* ``uv_pip_install`` so the install layer is not @@ -96,8 +95,8 @@ def create_app( def distributed_prefetch_worker( worker_id: str, partition_indices: list[int], - data_source_config: DictConfig, - processor_config: DictConfig, + data_source_config: dict[str, Any], + processor_config: dict[str, Any], queue_name: str, n_scanners: int = 8, prefetch_batch_count: int = 32, @@ -127,7 +126,7 @@ def distributed_prefetch_worker( queue = modal.Queue.from_name(queue_name, create_if_missing=True) partial_groups_kv = None if ( - bool(processor_config.enable_cross_worker_boundary_merging) + processor_config.get("enable_cross_worker_boundary_merging", False) and partial_groups_kv_name ): if modal_queue_environment is not None: @@ -210,13 +209,12 @@ class DistributedSLAFDataLoader: (raw_mode=True, wait for queue size >= 50, then iterate). """ - tokenizer: SLAFTokenizer | None + tokenizer: GeneformerTokenizer | ScGPTTokenizer | None def __init__( self, slaf_array: SLAFArray, - tokenizer: SLAFTokenizer | None = None, - tokenizer_config: DictConfig | dict[str, Any] | None = None, + tokenizer_type: str = "geneformer", n_workers: int = 64, n_scanners: int = 16, cpu: float = 8, @@ -224,6 +222,9 @@ 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, @@ -240,10 +241,8 @@ def __init__( Args: slaf_array: SLAFArray instance containing the data - 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. + tokenizer_type [PRODUCER]: Tokenization strategy ("geneformer", "scgpt", or "raw") + If "raw", raw_mode is automatically enabled 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=...). @@ -254,6 +253,9 @@ 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). @@ -277,6 +279,7 @@ 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 @@ -290,32 +293,44 @@ 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 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 + if tokenizer_type == "geneformer": + tokenizer_cls = GeneformerTokenizer + else: + tokenizer_cls = ScGPTTokenizer + tokenizer_factory_kwargs = { + "vocab_size": vocab_size, + "max_genes": max_genes, + } + if tokenizer_type == "scgpt": + tokenizer_factory_kwargs["n_expression_bins"] = n_expression_bins + self.tokenizer = tokenizer_cls( + slaf_array=slaf_array, + **tokenizer_factory_kwargs, + ) self.special_tokens = self.tokenizer.special_tokens else: - 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 + tokenizer_factory_kwargs = None + tokenizer_cls = None self.tokenizer = None self.special_tokens = None # Create data source @@ -340,12 +355,10 @@ def __init__( partial_groups_kv_name = f"{queue_name}-partial-groups" # Prepare configs for workers - data_source_config = OmegaConf.create( - { - "type": "lance", - "path": lance_path, - } - ) + data_source_config = { + "type": "lance", + "path": lance_path, + } # Data schema for SLAF (maps generic schema to SLAF column names) schema = DataSchema( @@ -358,37 +371,42 @@ def __init__( value_list_key="expr_sequence", ) - 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, - } - ) + 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, + } # Create the KV dict to ensure it exists before workers try to access it - if bool(processor_config.enable_cross_worker_boundary_merging): + if processor_config.get("enable_cross_worker_boundary_merging", True): _modal_dict_from_name( partial_groups_kv_name, create_if_missing=True, environment_name=modal_queue_environment, ) - processor_config.tokenizer_config = tokenizer_config + 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 # Spawn workers # NOTE: The app must be deployed before spawning workers: diff --git a/slaf/ml/tokenizers.py b/slaf/ml/tokenizers.py index 20508e9..c71e940 100644 --- a/slaf/ml/tokenizers.py +++ b/slaf/ml/tokenizers.py @@ -3,11 +3,9 @@ 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.integrations.anndata import LazyAnnData +from slaf.core.slaf import SLAFArray from slaf.ml.aggregators import GeneformerWindow, ScGPTWindow, Window TORCH_AVAILABLE = True @@ -43,7 +41,7 @@ class SLAFTokenizer(ABC): def __init__( self, - adata: LazyAnnData, + slaf_array: SLAFArray, vocab_size: int = 50000, max_genes: int = 2048, ): @@ -51,9 +49,9 @@ def __init__( Initialize SLAFTokenizer with SLAF array and vocabulary settings. Args: - adata: LazyAnnData instance containing the single-cell data. - Used to build the gene vocabulary, access expression metadata, - and inspect lazy runtime transformations. + slaf_array: Initialized SLAFArray instance containing the single-cell data. + Used to build the gene vocabulary and access expression data. + Must be a valid SLAFArray with proper var DataFrame. vocab_size: Maximum size of gene vocabulary. Genes beyond this limit are excluded from tokenization. Higher values use more memory. max_genes: Max genes per cell for windowing and tokenization (sequence layout @@ -63,10 +61,9 @@ def __init__( Raises: ValueError: If vocab_size or max_genes is invalid. RuntimeError: If SLAF array is not properly initialized. - TypeError: If adata is not a valid LazyAnnData instance. + TypeError: If slaf_array is not a valid SLAFArray instance. """ - self.adata = adata - self.slaf_array = adata.slaf + self.slaf_array = slaf_array self.vocab_size = vocab_size if max_genes < 1: raise ValueError(f"max_genes must be >= 1, got {max_genes}") @@ -78,11 +75,6 @@ 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: @@ -245,95 +237,6 @@ def create_window(self) -> Window: Create a window function based on the tokenizer type. """ - def _get_runtime_transformations(self) -> dict[str, Any]: - transformations = getattr(self.adata, "_transformations", None) - if not isinstance(transformations, dict): - return {} - return transformations - - def _apply_runtime_transformations( - self, - df: pl.DataFrame, - schema: DataSchema, - ) -> pl.DataFrame: - transformations = self._get_runtime_transformations() - if not transformations: - return df - - transformed_df = df - value_col = schema.value_key - group_col = schema.group_key - - for transform_name, transform_data in transformations.items(): - if transform_name == "normalize_total": - cell_factors = transform_data.get("cell_factors") - if not isinstance(cell_factors, dict) or not cell_factors: - raise ValueError( - "normalize_total runtime transformation requires precomputed " - "cell_factors keyed by cell_integer_id" - ) - - factor_df = pl.DataFrame( - { - group_col: [int(cell_id) for cell_id in cell_factors.keys()], - "_normalization_factor": [ - float(factor) for factor in cell_factors.values() - ], - } - ) - transformed_df = ( - transformed_df.join(factor_df, on=group_col, how="left") - .with_columns( - ( - pl.col(value_col) - * pl.col("_normalization_factor").fill_null(1.0) - ).alias(value_col) - ) - .drop("_normalization_factor") - ) - elif transform_name == "log1p": - transformed_df = transformed_df.with_columns( - pl.col(value_col).log1p().alias(value_col) - ) - - return transformed_df - - def transform_and_apply( - self, - df: pl.DataFrame, - schema: DataSchema, - max_items: int, - **kwargs: Any, - ) -> pl.DataFrame: - """Apply runtime transformations, then group rows into tokenizer-ready sequences.""" - transformed_df = self._apply_runtime_transformations(df, schema) - return self.apply(transformed_df, schema=schema, max_items=max_items, **kwargs) - - def apply( - self, - df: pl.DataFrame, - schema: DataSchema, - max_items: int, - **kwargs: Any, - ) -> pl.DataFrame: - """Group already-transformed 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. @@ -414,17 +317,18 @@ class ScGPTTokenizer(SLAFTokenizer): def __init__( self, - adata: LazyAnnData, + slaf_array: SLAFArray, 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. Args: - adata: LazyAnnData instance containing the single-cell data. + slaf_array: Initialized SLAFArray instance containing the single-cell data. + Used to build the gene vocabulary and access expression data. + Must be a valid SLAFArray with proper var DataFrame. vocab_size: Maximum size of gene vocabulary. Genes beyond this limit are excluded from tokenization. Higher values use more memory. n_expression_bins: Number of expression bins for scGPT tokenization. @@ -432,22 +336,20 @@ 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. RuntimeError: If SLAF array is not properly initialized. - TypeError: If adata is not a valid LazyAnnData instance. + TypeError: If slaf_array is not a valid SLAFArray instance. Examples: >>> # Basic initialization >>> slaf_array = SLAFArray("path/to/data.slaf") - >>> tokenizer = ScGPTTokenizer(LazyAnnData(slaf_array)) + >>> tokenizer = ScGPTTokenizer(slaf_array) >>> # scGPT with custom settings >>> tokenizer = ScGPTTokenizer( - ... adata=LazyAnnData(slaf_array), + ... slaf_array=slaf_array, ... vocab_size=30000, ... n_expression_bins=20 ... ) @@ -459,95 +361,17 @@ def __init__( ... tokenizer = ScGPTTokenizer(None) ... except TypeError as e: ... print(f"Error: {e}") - Error: adata must be a valid LazyAnnData instance + Error: slaf_array must be a valid SLAFArray instance """ self.n_expression_bins = n_expression_bins - self.use_binned_expressions = use_binned_expressions - super().__init__(adata=adata, 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() - def apply( - self, - df: pl.DataFrame, - schema: DataSchema, - max_items: int, - **kwargs: Any, - ) -> pl.DataFrame: - kwargs.setdefault("special_token_offset", 4) - kwargs.setdefault("n_expression_bins", self.n_expression_bins) - kwargs.setdefault("use_binned_expressions", self.use_binned_expressions) - 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() - 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]]], @@ -615,7 +439,9 @@ 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) + expr_tokens = ( + np.array(exprs[:n_pairs], dtype=np.int64) + self.expr_bin_start + ) else: expr_tokens = self._expression_to_bin_vectorized( np.array(exprs[:n_pairs], dtype=np.float32) @@ -678,6 +504,7 @@ 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: @@ -689,7 +516,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 1 + bin_id + return self.expr_bin_start + bin_id def _expression_to_bin_vectorized( self, expression_values: np.ndarray @@ -709,7 +536,7 @@ def _expression_to_bin_vectorized( # Convert to token IDs result = np.where( expression_values > 0, - 1 + bins, + self.expr_bin_start + bins, self.special_tokens["PAD"], ) @@ -753,8 +580,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 1 <= token <= self.n_expression_bins: # Expression bin - bin_id = token - 1 + elif token >= self.expr_bin_start: # Expression token + bin_id = token - self.expr_bin_start expr_value = bin_id * self.expr_bin_size expressions.append(expr_value) else: @@ -794,66 +621,17 @@ class GeneformerTokenizer(SLAFTokenizer): def __init__( self, - adata: LazyAnnData, + slaf_array: SLAFArray, vocab_size: int = 50000, max_genes: int = 2048, ): - super().__init__(adata=adata, 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 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/conftest.py b/tests/conftest.py index 11a5b4d..e96082e 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -12,7 +12,6 @@ from slaf.core.slaf import SLAFArray from slaf.data import SLAFConverter from slaf.integrations import ensure_h5ad_writable -from slaf.integrations.anndata import LazyAnnData def _wait_for_all_slaf_threads(timeout=0.3): @@ -442,11 +441,6 @@ def tiny_slaf(temp_dir, tiny_adata): return slaf_array -@pytest.fixture -def tiny_lazy_adata(tiny_slaf): - return LazyAnnData(tiny_slaf) - - @pytest.fixture def small_slaf(temp_dir, small_adata): """Create a small SLAFArray for SLAFArray testing.""" diff --git a/tests/test_aggregators.py b/tests/test_aggregators.py index 59d43d6..333072e 100644 --- a/tests/test_aggregators.py +++ b/tests/test_aggregators.py @@ -81,7 +81,7 @@ def test_apply_ranking(self): assert gene_seq[0] == 10 # gene with value 5.0 # 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) # PAD + 10 bins + assert all(0 <= bin_val < 10 for bin_val in expr_seq) # Default 10 bins def test_apply_expression_binning(self): """Test that expression binning works correctly""" @@ -91,7 +91,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) # PAD + 5 bins + assert all(0 <= bin_val < 5 for bin_val in expr_seq) # 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 +122,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 +167,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) # PAD + 10 bins + assert all(0 <= bin_val < 10 for bin_val in expr_seq) # Default 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_dataloaders.py b/tests/test_dataloaders.py index c8da8a4..78695af 100644 --- a/tests/test_dataloaders.py +++ b/tests/test_dataloaders.py @@ -2,37 +2,21 @@ import pytest import torch -from slaf.ml.dataloaders import SLAFDataLoader, get_device_info, get_optimal_device -from slaf.ml.tokenizers import GeneformerTokenizer, ScGPTTokenizer - - -def build_dataloader(adata, 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(adata, raw_mode=True, **kwargs) - - if tokenizer_kind == "scgpt": - tokenizer = ScGPTTokenizer(adata, **tokenizer_kwargs) - else: - tokenizer = GeneformerTokenizer(adata, **tokenizer_kwargs) - - return SLAFDataLoader( - adata, - tokenizer=tokenizer, - **kwargs, - ) +from slaf.ml.dataloaders import ( + GeneformerTokenizer, + ScGPTTokenizer, + SLAFDataLoader, + get_device_info, + get_optimal_device, +) class TestSLAFDataLoader: """Test suite for SLAFDataLoader class with new architecture""" - def test_dataloader_initialization(self, tiny_slaf, tiny_lazy_adata): + def test_dataloader_initialization(self, tiny_slaf): """Test SLAFDataLoader initialization with new architecture""" - dataloader = build_dataloader(tiny_lazy_adata) + dataloader = SLAFDataLoader(tiny_slaf) # Check basic attributes assert dataloader.slaf_array is tiny_slaf @@ -48,11 +32,11 @@ def test_dataloader_initialization(self, tiny_slaf, tiny_lazy_adata): # Check that we're using the new dataset implementation assert hasattr(dataloader, "_dataset") - def test_dataloader_initialization_custom_params(self, tiny_lazy_adata): + def test_dataloader_initialization_custom_params(self, tiny_slaf): """Test SLAFDataLoader initialization with custom parameters""" - dataloader = build_dataloader( - tiny_lazy_adata, - tokenizer_kind="scgpt", + dataloader = SLAFDataLoader( + tiny_slaf, + tokenizer_type="scgpt", batch_size=16, max_genes=1024, vocab_size=1000, @@ -66,20 +50,11 @@ def test_dataloader_initialization_custom_params(self, tiny_lazy_adata): assert dataloader.tokenizer.vocab_size == 1000 assert dataloader.tokenizer.n_expression_bins == 5 - def test_dataloader_initialization_with_default_source(self, tiny_lazy_adata): - """Test SLAFDataLoader defaults to the main expression matrix.""" - dataloader = build_dataloader(tiny_lazy_adata) - batch_processor = dataloader._dataset.batch_processor - - assert batch_processor.expression_dataset is not None - assert batch_processor.use_mixture_of_scanners is True - assert batch_processor.by_fragment is True - - def test_geneformer_iteration(self, tiny_lazy_adata): + def test_geneformer_iteration(self, tiny_slaf): """Test dataloader iteration with Geneformer tokenizer""" - dataloader = build_dataloader( - tiny_lazy_adata, - tokenizer_kind="geneformer", + dataloader = SLAFDataLoader( + tiny_slaf, + tokenizer_type="geneformer", batch_size=5, max_genes=10, ) @@ -113,11 +88,11 @@ def test_geneformer_iteration(self, tiny_lazy_adata): assert batch_count > 0 - def test_scgpt_iteration(self, tiny_lazy_adata): + def test_scgpt_iteration(self, tiny_slaf): """Test dataloader iteration with scGPT tokenizer""" - dataloader = build_dataloader( - tiny_lazy_adata, - tokenizer_kind="scgpt", + dataloader = SLAFDataLoader( + tiny_slaf, + tokenizer_type="scgpt", batch_size=5, max_genes=10, ) @@ -138,7 +113,7 @@ def test_scgpt_iteration(self, tiny_lazy_adata): 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: max_genes + 2 + assert input_ids.shape[1] == 12 # scGPT dual stream: max_genes + 2 assert values.shape == input_ids.shape # Check data types @@ -152,9 +127,9 @@ def test_scgpt_iteration(self, tiny_lazy_adata): assert batch_count > 0 - def test_consistent_batch_sizes(self, tiny_lazy_adata): + def test_consistent_batch_sizes(self, tiny_slaf): """Test that batches have consistent sizes""" - dataloader = build_dataloader(tiny_lazy_adata, batch_size=8) + dataloader = SLAFDataLoader(tiny_slaf, batch_size=8) batch_sizes = [] for batch in dataloader: @@ -169,9 +144,9 @@ def test_consistent_batch_sizes(self, tiny_lazy_adata): for size in batch_sizes[:-1]: assert size == dataloader.batch_size - def test_cell_id_mapping(self, tiny_slaf, tiny_lazy_adata): + def test_cell_id_mapping(self, tiny_slaf): """Test that cell IDs are properly mapped""" - dataloader = build_dataloader(tiny_lazy_adata, batch_size=5) + dataloader = SLAFDataLoader(tiny_slaf, batch_size=5) for batch in dataloader: cell_ids = batch["cell_ids"] @@ -185,9 +160,9 @@ def test_cell_id_mapping(self, tiny_slaf, tiny_lazy_adata): break # Just test first batch - def test_tokenizer_integration(self, tiny_lazy_adata): + def test_tokenizer_integration(self, tiny_slaf): """Test that tokenizer is properly integrated""" - dataloader = build_dataloader(tiny_lazy_adata) + dataloader = SLAFDataLoader(tiny_slaf) # Check that tokenizer has expected attributes assert hasattr(dataloader.tokenizer, "gene_vocab") @@ -197,9 +172,9 @@ def test_tokenizer_integration(self, tiny_lazy_adata): # Check that special tokens are properly set assert dataloader.special_tokens == dataloader.tokenizer.special_tokens - def test_dataloader_cleanup(self, tiny_lazy_adata): + def test_dataloader_cleanup(self, tiny_slaf): """Test dataloader cleanup functionality""" - dataloader = build_dataloader(tiny_lazy_adata) + dataloader = SLAFDataLoader(tiny_slaf) # Test that cleanup methods exist and don't crash # The dataloader doesn't have a stop_streaming method @@ -208,9 +183,9 @@ def test_dataloader_cleanup(self, tiny_lazy_adata): # Test destructor dataloader.__del__() - def test_dataloader_device(self, tiny_lazy_adata): + def test_dataloader_device(self, tiny_slaf): """Test dataloader device handling""" - dataloader = build_dataloader(tiny_lazy_adata) + dataloader = SLAFDataLoader(tiny_slaf) # Get a sample batch batch = next(iter(dataloader)) @@ -220,10 +195,10 @@ def test_dataloader_device(self, tiny_lazy_adata): assert batch["attention_mask"].device == torch.device("cpu") assert batch["cell_ids"].device == torch.device("cpu") - def test_multi_epoch_initialization(self, tiny_lazy_adata): + def test_multi_epoch_initialization(self, tiny_slaf): """Test SLAFDataLoader initialization with multi-epoch support""" - dataloader = build_dataloader( - tiny_lazy_adata, + dataloader = SLAFDataLoader( + tiny_slaf, n_epochs=5, # Test multi-epoch initialization ) @@ -231,10 +206,10 @@ def test_multi_epoch_initialization(self, tiny_lazy_adata): assert dataloader._dataset.n_epochs == 5 assert dataloader._dataset.batch_processor.n_epochs == 5 - def test_multi_epoch_iteration(self, tiny_lazy_adata): + def test_multi_epoch_iteration(self, tiny_slaf): """Test dataloader iteration with multiple epochs""" - dataloader = build_dataloader( - tiny_lazy_adata, + dataloader = SLAFDataLoader( + tiny_slaf, batch_size=4, # Small batch size for testing n_epochs=3, # Test with 3 epochs ) @@ -262,10 +237,10 @@ def test_multi_epoch_iteration(self, tiny_lazy_adata): assert len(epochs_seen) >= 1, f"Expected at least 1 epoch, got {epochs_seen}" assert total_batches > 0 - def test_multi_epoch_epoch_progression(self, tiny_lazy_adata): + def test_multi_epoch_epoch_progression(self, tiny_slaf): """Test that epochs progress correctly in multi-epoch mode""" - dataloader = build_dataloader( - tiny_lazy_adata, + dataloader = SLAFDataLoader( + tiny_slaf, batch_size=2, # Very small batch size n_epochs=4, # Test with 4 epochs ) @@ -295,10 +270,10 @@ def test_multi_epoch_epoch_progression(self, tiny_lazy_adata): "Epochs should not go backwards" ) - def test_single_epoch_default_behavior(self, tiny_lazy_adata): + def test_single_epoch_default_behavior(self, tiny_slaf): """Test that single epoch (default) behavior is unchanged""" - dataloader = build_dataloader( - tiny_lazy_adata, + dataloader = SLAFDataLoader( + tiny_slaf, batch_size=8, n_epochs=1, # Single epoch (default) ) @@ -318,12 +293,12 @@ def test_single_epoch_default_behavior(self, tiny_lazy_adata): assert epochs_seen == {0}, f"Expected only epoch 0, got {epochs_seen}" assert batch_count > 0 - def test_multi_epoch_with_different_tokenizers(self, tiny_lazy_adata): + 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 = build_dataloader( - tiny_lazy_adata, - tokenizer_kind="geneformer", + dataloader_geneformer = SLAFDataLoader( + tiny_slaf, + tokenizer_type="geneformer", batch_size=4, n_epochs=2, ) @@ -342,9 +317,9 @@ def test_multi_epoch_with_different_tokenizers(self, tiny_lazy_adata): ) # Test with scGPT - limit epochs and batches for speed - dataloader_scgpt = build_dataloader( - tiny_lazy_adata, - tokenizer_kind="scgpt", + dataloader_scgpt = SLAFDataLoader( + tiny_slaf, + tokenizer_type="scgpt", batch_size=4, n_epochs=2, ) @@ -362,10 +337,10 @@ def test_multi_epoch_with_different_tokenizers(self, tiny_lazy_adata): f"scGPT: Expected at least 1 epoch, got {epochs_scgpt}" ) - def test_multi_epoch_completion(self, tiny_lazy_adata): + def test_multi_epoch_completion(self, tiny_slaf): """Test that dataloader correctly completes all epochs""" - dataloader = build_dataloader( - tiny_lazy_adata, + dataloader = SLAFDataLoader( + tiny_slaf, batch_size=2, # Small batch size to complete quickly n_epochs=3, # Test with 3 epochs ) @@ -389,10 +364,10 @@ def test_multi_epoch_completion(self, tiny_lazy_adata): # Should see at least one epoch assert len(epochs_seen) >= 1, f"Expected at least 1 epoch, got {epochs_seen}" - def test_multi_epoch_parameter_passing(self, tiny_lazy_adata): + def test_multi_epoch_parameter_passing(self, tiny_slaf): """Test that n_epochs parameter is correctly passed through the hierarchy""" - dataloader = build_dataloader( - tiny_lazy_adata, + dataloader = SLAFDataLoader( + tiny_slaf, n_epochs=7, # Test with 7 epochs ) @@ -402,11 +377,11 @@ def test_multi_epoch_parameter_passing(self, tiny_lazy_adata): assert dataloader._dataset.batch_processor.n_epochs == 7 assert dataloader._dataset.prefetcher.batch_processor.n_epochs == 7 - def test_multi_epoch_with_custom_parameters(self, tiny_lazy_adata): + def test_multi_epoch_with_custom_parameters(self, tiny_slaf): """Test multi-epoch functionality with custom parameters""" - dataloader = build_dataloader( - tiny_lazy_adata, - tokenizer_kind="scgpt", + dataloader = SLAFDataLoader( + tiny_slaf, + tokenizer_type="scgpt", batch_size=6, max_genes=512, n_epochs=4, @@ -436,11 +411,11 @@ def test_multi_epoch_with_custom_parameters(self, tiny_lazy_adata): assert len(epochs_seen) >= 1, f"Expected at least 1 epoch, got {epochs_seen}" assert batch_count > 0 - def test_dataloader_fragment_parameter(self, tiny_lazy_adata): + def test_dataloader_fragment_parameter(self, tiny_slaf): """Test the by_fragment parameter functionality.""" # Test fragment-based loading dataloader_fragment = SLAFDataLoader( - tiny_lazy_adata, + tiny_slaf, batch_size=32, raw_mode=True, verbose=False, @@ -451,7 +426,7 @@ def test_dataloader_fragment_parameter(self, tiny_lazy_adata): # Test batch-based loading dataloader_batch = SLAFDataLoader( - tiny_lazy_adata, + tiny_slaf, batch_size=32, raw_mode=True, verbose=False, @@ -460,10 +435,10 @@ def test_dataloader_fragment_parameter(self, tiny_lazy_adata): assert dataloader_batch.by_fragment is False - def test_dataloader_iteration_fragment_mode(self, tiny_lazy_adata): + def test_dataloader_iteration_fragment_mode(self, tiny_slaf): """Test dataloader iteration in fragment mode.""" dataloader = SLAFDataLoader( - tiny_lazy_adata, + tiny_slaf, batch_size=32, raw_mode=True, verbose=False, @@ -481,10 +456,10 @@ def test_dataloader_iteration_fragment_mode(self, tiny_lazy_adata): assert batch_count > 0 - def test_dataloader_iteration_batch_mode(self, tiny_lazy_adata): + def test_dataloader_iteration_batch_mode(self, tiny_slaf): """Test dataloader iteration in batch mode.""" dataloader = SLAFDataLoader( - tiny_lazy_adata, + tiny_slaf, batch_size=32, raw_mode=True, verbose=False, @@ -502,10 +477,10 @@ def test_dataloader_iteration_batch_mode(self, tiny_lazy_adata): assert batch_count > 0 - def test_dataloader_tokenized_mode_fragment(self, tiny_lazy_adata): + def test_dataloader_tokenized_mode_fragment(self, tiny_slaf): """Test dataloader in tokenized mode with fragment loading.""" - dataloader = build_dataloader( - tiny_lazy_adata, + dataloader = SLAFDataLoader( + tiny_slaf, batch_size=32, raw_mode=False, verbose=False, @@ -524,10 +499,10 @@ def test_dataloader_tokenized_mode_fragment(self, tiny_lazy_adata): assert batch_count > 0 - def test_dataloader_tokenized_mode_batch(self, tiny_lazy_adata): + def test_dataloader_tokenized_mode_batch(self, tiny_slaf): """Test dataloader in tokenized mode with batch loading.""" - dataloader = build_dataloader( - tiny_lazy_adata, + dataloader = SLAFDataLoader( + tiny_slaf, batch_size=32, raw_mode=False, verbose=False, @@ -546,11 +521,14 @@ def test_dataloader_tokenized_mode_batch(self, tiny_lazy_adata): assert batch_count > 0 - def test_dataloader_parameters_consistency(self, tiny_lazy_adata): + def test_dataloader_parameters_consistency(self, tiny_slaf): """Test that all parameters are properly passed through.""" - dataloader = build_dataloader( - tiny_lazy_adata, + dataloader = SLAFDataLoader( + tiny_slaf, batch_size=64, + max_genes=1024, + vocab_size=10000, + n_expression_bins=5, n_epochs=5, raw_mode=True, verbose=True, @@ -560,27 +538,27 @@ def test_dataloader_parameters_consistency(self, tiny_lazy_adata): ) assert dataloader.batch_size == 64 - assert dataloader.max_genes == 0 + assert dataloader.max_genes == 1024 assert dataloader.n_epochs == 5 assert dataloader.raw_mode is True assert dataloader.verbose is True assert dataloader.batches_per_chunk == 25 assert dataloader.by_fragment is True - def test_dataloader_length(self, tiny_lazy_adata): + def test_dataloader_length(self, tiny_slaf): """Test that dataloader length returns 0 for streaming datasets.""" - dataloader = build_dataloader( - tiny_lazy_adata, + dataloader = SLAFDataLoader( + tiny_slaf, batch_size=32, verbose=False, ) assert len(dataloader) == 0 # Streaming datasets have unknown length (return 0) - def test_mixture_of_scanners_initialization(self, tiny_lazy_adata): + def test_mixture_of_scanners_initialization(self, tiny_slaf): """Test SLAFDataLoader initialization with Mixture of Scanners (MoS)""" - dataloader = build_dataloader( - tiny_lazy_adata, + dataloader = SLAFDataLoader( + slaf_array=tiny_slaf, batch_size=32, use_mixture_of_scanners=True, n_scanners=8, @@ -595,11 +573,11 @@ def test_mixture_of_scanners_initialization(self, tiny_lazy_adata): assert dataloader._dataset.batch_processor.n_scanners == 8 assert dataloader._dataset.batch_processor.prefetch_batch_size == 1048576 - def test_mixture_of_scanners_parameter_validation(self, tiny_lazy_adata): + def test_mixture_of_scanners_parameter_validation(self, tiny_slaf): """Test MoS parameter validation in SLAFDataLoader""" # Test valid parameters - dataloader = build_dataloader( - tiny_lazy_adata, + dataloader = SLAFDataLoader( + slaf_array=tiny_slaf, batch_size=32, use_mixture_of_scanners=True, n_scanners=16, @@ -610,7 +588,7 @@ def test_mixture_of_scanners_parameter_validation(self, tiny_lazy_adata): # Test invalid n_scanners (too low) with pytest.raises(ValueError, match="n_scanners must be at least 1"): SLAFDataLoader( - tiny_lazy_adata, + slaf_array=tiny_slaf, batch_size=32, use_mixture_of_scanners=True, n_scanners=0, @@ -620,7 +598,7 @@ def test_mixture_of_scanners_parameter_validation(self, tiny_lazy_adata): # Test invalid n_scanners (too high) with pytest.raises(ValueError, match="n_scanners cannot exceed 100"): SLAFDataLoader( - tiny_lazy_adata, + slaf_array=tiny_slaf, batch_size=32, use_mixture_of_scanners=True, n_scanners=101, @@ -632,7 +610,7 @@ def test_mixture_of_scanners_parameter_validation(self, tiny_lazy_adata): ValueError, match="prefetch_batch_size must be at least 1,000" ): SLAFDataLoader( - tiny_lazy_adata, + slaf_array=tiny_slaf, batch_size=32, use_mixture_of_scanners=True, n_scanners=16, @@ -644,17 +622,17 @@ def test_mixture_of_scanners_parameter_validation(self, tiny_lazy_adata): ValueError, match="prefetch_batch_size cannot exceed 10,000,000" ): SLAFDataLoader( - tiny_lazy_adata, + slaf_array=tiny_slaf, batch_size=32, use_mixture_of_scanners=True, n_scanners=16, prefetch_batch_size=10000001, ) - def test_mixture_of_scanners_iteration(self, tiny_lazy_adata): + def test_mixture_of_scanners_iteration(self, tiny_slaf): """Test that MoS dataloader can iterate through batches""" - dataloader = build_dataloader( - tiny_lazy_adata, + dataloader = SLAFDataLoader( + slaf_array=tiny_slaf, batch_size=8, use_mixture_of_scanners=True, n_scanners=4, @@ -673,10 +651,10 @@ def test_mixture_of_scanners_iteration(self, tiny_lazy_adata): assert batch_count > 0 - def test_mixture_of_scanners_with_raw_mode(self, tiny_lazy_adata): + def test_mixture_of_scanners_with_raw_mode(self, tiny_slaf): """Test MoS functionality with raw mode in SLAFDataLoader""" dataloader = SLAFDataLoader( - tiny_lazy_adata, + slaf_array=tiny_slaf, batch_size=32, raw_mode=True, use_mixture_of_scanners=True, @@ -698,10 +676,10 @@ def test_mixture_of_scanners_with_raw_mode(self, tiny_lazy_adata): assert batch_count > 0 - def test_mixture_of_scanners_with_tokenized_mode(self, tiny_lazy_adata): + def test_mixture_of_scanners_with_tokenized_mode(self, tiny_slaf): """Test MoS functionality with tokenized mode in SLAFDataLoader""" - dataloader = build_dataloader( - tiny_lazy_adata, + dataloader = SLAFDataLoader( + slaf_array=tiny_slaf, batch_size=32, raw_mode=False, use_mixture_of_scanners=True, @@ -724,11 +702,11 @@ def test_mixture_of_scanners_with_tokenized_mode(self, tiny_lazy_adata): assert batch_count > 0 - def test_mixture_of_scanners_backward_compatibility(self, tiny_lazy_adata): + 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 = build_dataloader( - tiny_lazy_adata, + dataloader_default = SLAFDataLoader( + slaf_array=tiny_slaf, batch_size=32, ) @@ -738,8 +716,8 @@ def test_mixture_of_scanners_backward_compatibility(self, tiny_lazy_adata): ) # Explicitly disable MoS - dataloader_disabled = build_dataloader( - tiny_lazy_adata, + dataloader_disabled = SLAFDataLoader( + slaf_array=tiny_slaf, batch_size=32, use_mixture_of_scanners=False, ) @@ -750,10 +728,10 @@ def test_mixture_of_scanners_backward_compatibility(self, tiny_lazy_adata): is False ) - def test_mixture_of_scanners_parameter_passing(self, tiny_lazy_adata): + def test_mixture_of_scanners_parameter_passing(self, tiny_slaf): """Test that MoS parameters are correctly passed through the hierarchy""" - dataloader = build_dataloader( - tiny_lazy_adata, + dataloader = SLAFDataLoader( + slaf_array=tiny_slaf, batch_size=32, use_mixture_of_scanners=True, n_scanners=12, @@ -769,11 +747,11 @@ def test_mixture_of_scanners_parameter_passing(self, tiny_lazy_adata): assert dataloader._dataset.batch_processor.n_scanners == 12 assert dataloader._dataset.batch_processor.prefetch_batch_size == 2097152 - def test_mixture_of_scanners_with_custom_parameters(self, tiny_lazy_adata): + def test_mixture_of_scanners_with_custom_parameters(self, tiny_slaf): """Test MoS functionality with custom parameters in SLAFDataLoader""" - dataloader = build_dataloader( - tiny_lazy_adata, - tokenizer_kind="scgpt", + dataloader = SLAFDataLoader( + slaf_array=tiny_slaf, + tokenizer_type="scgpt", batch_size=16, use_mixture_of_scanners=True, n_scanners=6, @@ -786,7 +764,7 @@ def test_mixture_of_scanners_with_custom_parameters(self, tiny_lazy_adata): assert dataloader.prefetch_batch_size == 1048576 assert isinstance(dataloader.tokenizer, ScGPTTokenizer) assert dataloader.batch_size == 16 - assert dataloader.max_genes == dataloader.tokenizer.max_genes + assert dataloader.max_genes == 2048 # Test iteration batch_count = 0 @@ -800,10 +778,10 @@ def test_mixture_of_scanners_with_custom_parameters(self, tiny_lazy_adata): assert batch_count > 0 - def test_mixture_of_scanners_multi_epoch(self, tiny_lazy_adata): + def test_mixture_of_scanners_multi_epoch(self, tiny_slaf): """Test MoS functionality with multi-epoch training""" - dataloader = build_dataloader( - tiny_lazy_adata, + dataloader = SLAFDataLoader( + slaf_array=tiny_slaf, batch_size=8, n_epochs=3, use_mixture_of_scanners=True, @@ -830,10 +808,10 @@ def test_mixture_of_scanners_multi_epoch(self, tiny_lazy_adata): assert len(epochs_seen) >= 1, f"Expected at least 1 epoch, got {epochs_seen}" assert batch_count > 0 - def test_mixture_of_scanners_fragment_generators_creation(self, tiny_lazy_adata): + def test_mixture_of_scanners_fragment_generators_creation(self, tiny_slaf): """Test that MoS creates fragment generators correctly""" - dataloader = build_dataloader( - tiny_lazy_adata, + dataloader = SLAFDataLoader( + slaf_array=tiny_slaf, batch_size=32, use_mixture_of_scanners=True, n_scanners=4, @@ -854,10 +832,10 @@ def test_mixture_of_scanners_fragment_generators_creation(self, tiny_lazy_adata) dataloader._dataset.batch_processor.fragment_generators ) - def test_mixture_of_scanners_random_sampling_behavior(self, tiny_lazy_adata): + def test_mixture_of_scanners_random_sampling_behavior(self, tiny_slaf): """Test that MoS uses random sampling from fragment generators""" - dataloader = build_dataloader( - tiny_lazy_adata, + dataloader = SLAFDataLoader( + slaf_array=tiny_slaf, batch_size=8, use_mixture_of_scanners=True, n_scanners=2, # Small number for testing @@ -879,10 +857,10 @@ def test_mixture_of_scanners_random_sampling_behavior(self, tiny_lazy_adata): assert "attention_mask" in batch assert "cell_ids" in batch - def test_mixture_of_scanners_cell_boundary_handling(self, tiny_lazy_adata): + def test_mixture_of_scanners_cell_boundary_handling(self, tiny_slaf): """Test that MoS handles cell boundaries correctly in SLAFDataLoader""" - dataloader = build_dataloader( - tiny_lazy_adata, + dataloader = SLAFDataLoader( + slaf_array=tiny_slaf, batch_size=8, use_mixture_of_scanners=True, n_scanners=4, @@ -939,9 +917,9 @@ def test_tensor_on_optimal_device(self): @pytest.mark.skipif( not get_device_info()["torch_available"], reason="PyTorch not available" ) - def test_dataloader_device(self, tiny_lazy_adata): + def test_dataloader_device(self, tiny_slaf): """Test that dataloader uses the correct device""" - dataloader = build_dataloader(tiny_lazy_adata) + dataloader = SLAFDataLoader(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 deb5ab7..c06df3b 100644 --- a/tests/test_distributed_worker.py +++ b/tests/test_distributed_worker.py @@ -9,14 +9,8 @@ from unittest.mock import MagicMock, patch import polars as pl -import pytest from slaf.distributed.data_source import DataSource - -pytest.importorskip("omegaconf") - -from omegaconf import OmegaConf - from slaf.distributed.worker import prefetch_worker @@ -70,33 +64,6 @@ 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", @@ -135,8 +102,18 @@ def test_worker_initialization(self, mock_lance_ds): queue = MockQueue() # Configs - data_source_config = make_data_source_config() - processor_config = make_processor_config() + 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, + } # Call worker result = prefetch_worker( @@ -170,8 +147,18 @@ def test_worker_single_partition(self, mock_lance_ds): queue = MockQueue() # Configs - data_source_config = make_data_source_config() - processor_config = make_processor_config() + 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, + } # Call worker result = prefetch_worker( @@ -212,8 +199,18 @@ def create_reader(partition_index, batch_size): queue = MockQueue() # Configs - data_source_config = make_data_source_config() - processor_config = make_processor_config() + 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, + } # Call worker result = prefetch_worker( @@ -246,8 +243,18 @@ def test_worker_prefetch_batch_count(self, mock_lance_ds): queue = MockQueue() # Configs - data_source_config = make_data_source_config() - processor_config = make_processor_config() + 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, + } # Call worker with prefetch_batch_count result = prefetch_worker( @@ -280,8 +287,18 @@ def test_worker_prefetch_batch_size(self, mock_lance_ds): queue = MockQueue() # Configs - data_source_config = make_data_source_config() - processor_config = make_processor_config() + 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, + } # Call worker with prefetch_batch_size result = prefetch_worker( @@ -314,8 +331,18 @@ def test_worker_max_batches(self, mock_lance_ds): queue = MockQueue() # Configs - data_source_config = make_data_source_config() - processor_config = make_processor_config() + 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, + } # Call worker with max_batches result = prefetch_worker( @@ -350,10 +377,20 @@ def test_worker_cross_worker_boundary_merging(self, mock_lance_ds): kv_store = MockKVStore() # Configs with cross-worker merging enabled - data_source_config = make_data_source_config() - processor_config = make_processor_config( - enable_cross_worker_boundary_merging=True - ) + 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", + } # Call worker with KV store result = prefetch_worker( @@ -387,8 +424,18 @@ def test_worker_partition_exhaustion(self, mock_lance_ds): queue = MockQueue() # Configs - data_source_config = make_data_source_config() - processor_config = make_processor_config() + 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, + } # Call worker result = prefetch_worker( @@ -422,8 +469,18 @@ def create_reader(partition_index, batch_size): queue = MockQueue() # Configs - data_source_config = make_data_source_config() - processor_config = make_processor_config() + 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, + } # Call worker - should handle error gracefully result = prefetch_worker( @@ -456,8 +513,18 @@ def test_worker_metrics(self, mock_lance_ds): queue = MockQueue() # Configs - data_source_config = make_data_source_config() - processor_config = make_processor_config() + 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, + } # Call worker result = prefetch_worker( diff --git a/tests/test_pytorch_datasets.py b/tests/test_pytorch_datasets.py index 938718f..5628280 100644 --- a/tests/test_pytorch_datasets.py +++ b/tests/test_pytorch_datasets.py @@ -19,9 +19,9 @@ class TestSLAFIterableDataset: """Test suite for SLAFIterableDataset with comprehensive coverage""" - def test_dataset_initialization(self, tiny_slaf, tiny_lazy_adata): + def test_dataset_initialization(self, tiny_slaf): """Test SLAFIterableDataset initialization""" - tokenizer = GeneformerTokenizer(tiny_lazy_adata) + tokenizer = GeneformerTokenizer(tiny_slaf) dataset = SLAFIterableDataset( slaf_array=tiny_slaf, @@ -41,9 +41,9 @@ def test_dataset_initialization(self, tiny_slaf, tiny_lazy_adata): assert hasattr(dataset, "batch_processor") assert hasattr(dataset, "prefetcher") - def test_dataset_initialization_scgpt(self, tiny_slaf, tiny_lazy_adata): + def test_dataset_initialization_scgpt(self, tiny_slaf): """Test SLAFIterableDataset initialization with scGPT tokenizer""" - tokenizer = ScGPTTokenizer(tiny_lazy_adata) + tokenizer = ScGPTTokenizer(tiny_slaf) dataset = SLAFIterableDataset( slaf_array=tiny_slaf, @@ -57,10 +57,10 @@ def test_dataset_initialization_scgpt(self, tiny_slaf, tiny_lazy_adata): assert dataset.batch_size == 16 assert dataset.seed == 123 - def test_prefetch_batch_processor_initialization(self, tiny_slaf, tiny_lazy_adata): + def test_prefetch_batch_processor_initialization(self, tiny_slaf): """Test PrefetchBatchProcessor initialization and configuration""" shuffle = RandomShuffle() - tokenizer = ScGPTTokenizer(tiny_lazy_adata) + tokenizer = ScGPTTokenizer(tiny_slaf) processor = PrefetchBatchProcessor( slaf_array=tiny_slaf, @@ -82,12 +82,10 @@ def test_prefetch_batch_processor_initialization(self, tiny_slaf, tiny_lazy_adat # With MoS enabled by default, we should have fragment_generators instead of batch_generator assert hasattr(processor, "fragment_generators") - def test_prefetch_batch_processor_fragment_parameter( - self, tiny_slaf, tiny_lazy_adata - ): + def test_prefetch_batch_processor_fragment_parameter(self, tiny_slaf): """Test the by_fragment parameter in PrefetchBatchProcessor.""" shuffle = RandomShuffle() - tokenizer = ScGPTTokenizer(tiny_lazy_adata) + tokenizer = ScGPTTokenizer(tiny_slaf) # Test fragment-based loading (with MoS disabled to test fragment mode) processor_fragment = PrefetchBatchProcessor( @@ -115,10 +113,10 @@ def test_prefetch_batch_processor_fragment_parameter( assert processor_batch.by_fragment is False - def test_prefetch_batch_processor_reset_fragment(self, tiny_slaf, tiny_lazy_adata): + def test_prefetch_batch_processor_reset_fragment(self, tiny_slaf): """Test processor epoch reset in fragment mode.""" shuffle = RandomShuffle() - tokenizer = ScGPTTokenizer(tiny_lazy_adata) + tokenizer = ScGPTTokenizer(tiny_slaf) processor = PrefetchBatchProcessor( slaf_array=tiny_slaf, @@ -135,10 +133,10 @@ def test_prefetch_batch_processor_reset_fragment(self, tiny_slaf, tiny_lazy_adat assert processor.current_epoch == 1 assert processor.batch_id == 0 - def test_prefetch_batch_processor_reset_batch(self, tiny_slaf, tiny_lazy_adata): + def test_prefetch_batch_processor_reset_batch(self, tiny_slaf): """Test processor epoch reset in batch mode.""" shuffle = RandomShuffle() - tokenizer = ScGPTTokenizer(tiny_lazy_adata) + tokenizer = ScGPTTokenizer(tiny_slaf) processor = PrefetchBatchProcessor( slaf_array=tiny_slaf, @@ -155,10 +153,10 @@ def test_prefetch_batch_processor_reset_batch(self, tiny_slaf, tiny_lazy_adata): assert processor.current_epoch == 1 assert processor.batch_id == 0 - def test_prefetch_batch_processor_load_fragment(self, tiny_slaf, tiny_lazy_adata): + def test_prefetch_batch_processor_load_fragment(self, tiny_slaf): """Test processor batch loading in fragment mode.""" shuffle = RandomShuffle() - tokenizer = ScGPTTokenizer(tiny_lazy_adata) + tokenizer = ScGPTTokenizer(tiny_slaf) processor = PrefetchBatchProcessor( slaf_array=tiny_slaf, @@ -175,10 +173,10 @@ def test_prefetch_batch_processor_load_fragment(self, tiny_slaf, tiny_lazy_adata assert hasattr(batch, "attention_mask") assert hasattr(batch, "cell_integer_ids") - def test_prefetch_batch_processor_load_batch_mode(self, tiny_slaf, tiny_lazy_adata): + def test_prefetch_batch_processor_load_batch_mode(self, tiny_slaf): """Test processor batch loading in batch mode.""" shuffle = RandomShuffle() - tokenizer = ScGPTTokenizer(tiny_lazy_adata) + tokenizer = ScGPTTokenizer(tiny_slaf) processor = PrefetchBatchProcessor( slaf_array=tiny_slaf, @@ -240,10 +238,10 @@ def test_prefetch_batch_geneformer(self): assert batch.cell_integer_ids == [200, 201, 202] assert batch.tokenize_time == 0.05 - def test_async_prefetcher_initialization(self, tiny_slaf, tiny_lazy_adata): + def test_async_prefetcher_initialization(self, tiny_slaf): """Test AsyncPrefetcher initialization""" shuffle = RandomShuffle() - tokenizer = ScGPTTokenizer(tiny_lazy_adata) + tokenizer = ScGPTTokenizer(tiny_slaf) processor = PrefetchBatchProcessor( slaf_array=tiny_slaf, @@ -260,10 +258,10 @@ def test_async_prefetcher_initialization(self, tiny_slaf, tiny_lazy_adata): assert prefetcher.worker_thread is None assert prefetcher.should_stop is False - def test_prefetcher_start_stop(self, tiny_slaf, tiny_lazy_adata): + def test_prefetcher_start_stop(self, tiny_slaf): """Test AsyncPrefetcher start and stop functionality""" shuffle = RandomShuffle() - tokenizer = ScGPTTokenizer(tiny_lazy_adata) + tokenizer = ScGPTTokenizer(tiny_slaf) processor = PrefetchBatchProcessor( slaf_array=tiny_slaf, @@ -283,10 +281,10 @@ def test_prefetcher_start_stop(self, tiny_slaf, tiny_lazy_adata): prefetcher.stop() assert prefetcher.should_stop is True - def test_prefetcher_queue_operations(self, tiny_slaf, tiny_lazy_adata): + def test_prefetcher_queue_operations(self, tiny_slaf): """Test AsyncPrefetcher queue operations""" shuffle = RandomShuffle() - tokenizer = ScGPTTokenizer(tiny_lazy_adata) + tokenizer = ScGPTTokenizer(tiny_slaf) processor = PrefetchBatchProcessor( slaf_array=tiny_slaf, @@ -307,9 +305,9 @@ def test_prefetcher_queue_operations(self, tiny_slaf, tiny_lazy_adata): assert "elapsed_time" in stats assert "cells_per_sec" in stats - def test_dataset_iteration_geneformer(self, tiny_slaf, tiny_lazy_adata): + def test_dataset_iteration_geneformer(self, tiny_slaf): """Test dataset iteration with Geneformer tokenizer""" - tokenizer = GeneformerTokenizer(tiny_lazy_adata) + tokenizer = GeneformerTokenizer(tiny_slaf) dataset = SLAFIterableDataset( slaf_array=tiny_slaf, tokenizer=tokenizer, @@ -340,9 +338,9 @@ def test_dataset_iteration_geneformer(self, tiny_slaf, tiny_lazy_adata): assert batch_count > 0 - def test_dataset_iteration_scgpt(self, tiny_slaf, tiny_lazy_adata): + def test_dataset_iteration_scgpt(self, tiny_slaf): """Test dataset iteration with scGPT tokenizer""" - tokenizer = ScGPTTokenizer(tiny_lazy_adata) + tokenizer = ScGPTTokenizer(tiny_slaf) dataset = SLAFIterableDataset( slaf_array=tiny_slaf, tokenizer=tokenizer, @@ -376,9 +374,9 @@ def test_dataset_iteration_scgpt(self, tiny_slaf, tiny_lazy_adata): assert batch_count > 0 - def test_scgpt_values_alignment_contract(self, tiny_slaf, tiny_lazy_adata): + def test_scgpt_values_alignment_contract(self, tiny_slaf): """scGPT batches must keep aligned dual streams.""" - tokenizer = ScGPTTokenizer(tiny_lazy_adata, max_genes=16) + tokenizer = ScGPTTokenizer(tiny_slaf, max_genes=16) dataset = SLAFIterableDataset( slaf_array=tiny_slaf, tokenizer=tokenizer, @@ -388,23 +386,17 @@ def test_scgpt_values_alignment_contract(self, tiny_slaf, tiny_lazy_adata): 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(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) + assert torch.all(values[cls_positions] == pad_value) + assert torch.all(values[sep_positions] == pad_value) - def test_device_transfer(self, tiny_slaf, tiny_lazy_adata): + def test_device_transfer(self, tiny_slaf): """Test device transfer functionality""" - tokenizer = GeneformerTokenizer(tiny_lazy_adata) + tokenizer = GeneformerTokenizer(tiny_slaf) dataset = SLAFIterableDataset( slaf_array=tiny_slaf, tokenizer=tokenizer, @@ -418,9 +410,9 @@ def test_device_transfer(self, tiny_slaf, tiny_lazy_adata): assert batch["cell_ids"].device.type == "cpu" break # Just test first batch - def test_prefetcher_timeout(self, tiny_slaf, tiny_lazy_adata): + def test_prefetcher_timeout(self, tiny_slaf): """Test prefetcher timeout handling""" - tokenizer = GeneformerTokenizer(tiny_lazy_adata) + tokenizer = GeneformerTokenizer(tiny_slaf) dataset = SLAFIterableDataset( slaf_array=tiny_slaf, tokenizer=tokenizer, @@ -437,10 +429,10 @@ def test_prefetcher_timeout(self, tiny_slaf, tiny_lazy_adata): assert batch_count > 0 - def test_pytorch_dataloader_integration(self, tiny_slaf, tiny_lazy_adata): + def test_pytorch_dataloader_integration(self, tiny_slaf): """Test integration with PyTorch DataLoader""" - tokenizer = GeneformerTokenizer(tiny_lazy_adata) + tokenizer = GeneformerTokenizer(tiny_slaf) dataset = SLAFIterableDataset( slaf_array=tiny_slaf, tokenizer=tokenizer, @@ -478,10 +470,10 @@ def test_pytorch_dataloader_integration(self, tiny_slaf, tiny_lazy_adata): assert batch_count > 0 - def test_window_strategy_integration(self, tiny_slaf, tiny_lazy_adata): + def test_window_strategy_integration(self, tiny_slaf): """Test window strategy integration with batch processor""" shuffle = RandomShuffle() - tokenizer = GeneformerTokenizer(tiny_lazy_adata) + tokenizer = GeneformerTokenizer(tiny_slaf) processor = PrefetchBatchProcessor( slaf_array=tiny_slaf, @@ -494,9 +486,9 @@ def test_window_strategy_integration(self, tiny_slaf, tiny_lazy_adata): assert processor.shuffle is shuffle assert processor.tokenizer.max_genes == 2048 - def test_multi_epoch_initialization(self, tiny_slaf, tiny_lazy_adata): + def test_multi_epoch_initialization(self, tiny_slaf): """Test SLAFIterableDataset initialization with multi-epoch support""" - tokenizer = GeneformerTokenizer(tiny_lazy_adata) + tokenizer = GeneformerTokenizer(tiny_slaf) dataset = SLAFIterableDataset( slaf_array=tiny_slaf, @@ -511,9 +503,9 @@ def test_multi_epoch_initialization(self, tiny_slaf, tiny_lazy_adata): assert dataset.batch_processor.current_epoch >= 0 assert dataset.batch_processor.current_epoch < 5 - def test_multi_epoch_iteration(self, tiny_slaf, tiny_lazy_adata): + def test_multi_epoch_iteration(self, tiny_slaf): """Test multi-epoch iteration functionality""" - tokenizer = GeneformerTokenizer(tiny_lazy_adata) + tokenizer = GeneformerTokenizer(tiny_slaf) dataset = SLAFIterableDataset( slaf_array=tiny_slaf, @@ -545,9 +537,9 @@ def test_multi_epoch_iteration(self, tiny_slaf, tiny_lazy_adata): assert len(epochs_seen) >= 1, f"Expected at least 1 epoch, got {epochs_seen}" assert total_batches > 0 - def test_epoch_transition_functionality(self, tiny_slaf, tiny_lazy_adata): + def test_epoch_transition_functionality(self, tiny_slaf): """Test that epoch transitions work correctly""" - tokenizer = GeneformerTokenizer(tiny_lazy_adata) + tokenizer = GeneformerTokenizer(tiny_slaf) dataset = SLAFIterableDataset( slaf_array=tiny_slaf, @@ -581,10 +573,10 @@ def test_epoch_transition_functionality(self, tiny_slaf, tiny_lazy_adata): "Epochs should not go backwards" ) - def test_batch_processor_epoch_reset(self, tiny_slaf, tiny_lazy_adata): + def test_batch_processor_epoch_reset(self, tiny_slaf): """Test PrefetchBatchProcessor epoch reset functionality""" shuffle = RandomShuffle() - tokenizer = GeneformerTokenizer(tiny_lazy_adata) + tokenizer = GeneformerTokenizer(tiny_slaf) processor = PrefetchBatchProcessor( slaf_array=tiny_slaf, @@ -610,9 +602,9 @@ def test_batch_processor_epoch_reset(self, tiny_slaf, tiny_lazy_adata): with pytest.raises(ValueError): processor.reset_for_epoch(3) # Invalid epoch (>= n_epochs) - def test_multi_epoch_with_small_dataset(self, tiny_slaf, tiny_lazy_adata): + def test_multi_epoch_with_small_dataset(self, tiny_slaf): """Test multi-epoch functionality with small dataset that gets exhausted""" - tokenizer = GeneformerTokenizer(tiny_lazy_adata) + tokenizer = GeneformerTokenizer(tiny_slaf) dataset = SLAFIterableDataset( slaf_array=tiny_slaf, @@ -642,9 +634,9 @@ def test_multi_epoch_with_small_dataset(self, tiny_slaf, tiny_lazy_adata): ) assert total_batches > 0 - def test_single_epoch_behavior(self, tiny_slaf, tiny_lazy_adata): + def test_single_epoch_behavior(self, tiny_slaf): """Test that single epoch (default) behavior is unchanged""" - tokenizer = GeneformerTokenizer(tiny_lazy_adata) + tokenizer = GeneformerTokenizer(tiny_slaf) dataset = SLAFIterableDataset( slaf_array=tiny_slaf, @@ -668,9 +660,9 @@ def test_single_epoch_behavior(self, tiny_slaf, tiny_lazy_adata): assert epochs_seen == {0}, f"Expected only epoch 0, got {epochs_seen}" assert batch_count > 0 - def test_multi_epoch_prefetcher_stats(self, tiny_slaf, tiny_lazy_adata): + def test_multi_epoch_prefetcher_stats(self, tiny_slaf): """Test that AsyncPrefetcher correctly tracks multi-epoch statistics""" - tokenizer = GeneformerTokenizer(tiny_lazy_adata) + tokenizer = GeneformerTokenizer(tiny_slaf) dataset = SLAFIterableDataset( slaf_array=tiny_slaf, @@ -693,9 +685,9 @@ def test_multi_epoch_prefetcher_stats(self, tiny_slaf, tiny_lazy_adata): assert stats["n_epochs"] == 3 assert stats["current_epoch"] >= 0 - def test_multi_epoch_completion_detection(self, tiny_slaf, tiny_lazy_adata): + def test_multi_epoch_completion_detection(self, tiny_slaf): """Test that the dataset correctly detects when all epochs are completed""" - tokenizer = GeneformerTokenizer(tiny_lazy_adata) + tokenizer = GeneformerTokenizer(tiny_slaf) dataset = SLAFIterableDataset( slaf_array=tiny_slaf, @@ -793,10 +785,10 @@ def test_prefetch_batch_serialization(self): assert len(batch.cell_integer_ids) == 2 assert batch.tokenize_time == 0.1 - def test_batch_processor_expression_binning(self, tiny_slaf, tiny_lazy_adata): + def test_batch_processor_expression_binning(self, tiny_slaf): """Test batch processor with expression binning""" shuffle = RandomShuffle() - tokenizer = ScGPTTokenizer(tiny_lazy_adata) + tokenizer = ScGPTTokenizer(tiny_slaf) processor = PrefetchBatchProcessor( slaf_array=tiny_slaf, @@ -809,9 +801,9 @@ def test_batch_processor_expression_binning(self, tiny_slaf, tiny_lazy_adata): assert processor.n_expression_bins == 10 assert processor.use_binned_expressions is True - def test_dataset_fragment_parameter(self, tiny_slaf, tiny_lazy_adata): + def test_dataset_fragment_parameter(self, tiny_slaf): """Test the by_fragment parameter functionality.""" - tokenizer = GeneformerTokenizer(tiny_lazy_adata) + tokenizer = GeneformerTokenizer(tiny_slaf) # Test fragment-based loading dataset_fragment = SLAFIterableDataset( @@ -835,9 +827,9 @@ def test_dataset_fragment_parameter(self, tiny_slaf, tiny_lazy_adata): assert dataset_batch.by_fragment is False - def test_dataset_iteration_fragment_mode(self, tiny_slaf, tiny_lazy_adata): + def test_dataset_iteration_fragment_mode(self, tiny_slaf): """Test dataset iteration in fragment mode.""" - tokenizer = GeneformerTokenizer(tiny_lazy_adata) + tokenizer = GeneformerTokenizer(tiny_slaf) dataset = SLAFIterableDataset( slaf_array=tiny_slaf, tokenizer=tokenizer, @@ -858,9 +850,9 @@ def test_dataset_iteration_fragment_mode(self, tiny_slaf, tiny_lazy_adata): assert batch_count > 0 - def test_dataset_iteration_batch_mode(self, tiny_slaf, tiny_lazy_adata): + def test_dataset_iteration_batch_mode(self, tiny_slaf): """Test dataset iteration in batch mode.""" - tokenizer = GeneformerTokenizer(tiny_lazy_adata) + tokenizer = GeneformerTokenizer(tiny_slaf) dataset = SLAFIterableDataset( slaf_array=tiny_slaf, tokenizer=tokenizer, @@ -881,7 +873,7 @@ def test_dataset_iteration_batch_mode(self, tiny_slaf, tiny_lazy_adata): assert batch_count > 0 - def test_dataset_raw_mode_fragment(self, tiny_slaf, tiny_lazy_adata): + def test_dataset_raw_mode_fragment(self, tiny_slaf): """Test dataset in raw mode with fragment loading.""" dataset = SLAFIterableDataset( slaf_array=tiny_slaf, @@ -903,7 +895,7 @@ def test_dataset_raw_mode_fragment(self, tiny_slaf, tiny_lazy_adata): assert batch_count > 0 - def test_dataset_raw_mode_batch(self, tiny_slaf, tiny_lazy_adata): + def test_dataset_raw_mode_batch(self, tiny_slaf): """Test dataset in raw mode with batch loading.""" dataset = SLAFIterableDataset( slaf_array=tiny_slaf, @@ -925,9 +917,9 @@ def test_dataset_raw_mode_batch(self, tiny_slaf, tiny_lazy_adata): assert batch_count > 0 - def test_mixture_of_scanners_initialization(self, tiny_slaf, tiny_lazy_adata): + def test_mixture_of_scanners_initialization(self, tiny_slaf): """Test SLAFIterableDataset initialization with Mixture of Scanners (MoS)""" - tokenizer = GeneformerTokenizer(tiny_lazy_adata) + tokenizer = GeneformerTokenizer(tiny_slaf) dataset = SLAFIterableDataset( slaf_array=tiny_slaf, @@ -943,9 +935,9 @@ def test_mixture_of_scanners_initialization(self, tiny_slaf, tiny_lazy_adata): assert dataset.batch_processor.n_scanners == 8 assert dataset.batch_processor.prefetch_batch_size == 1048576 - def test_mixture_of_scanners_fragment_generators(self, tiny_slaf, tiny_lazy_adata): + def test_mixture_of_scanners_fragment_generators(self, tiny_slaf): """Test that MoS creates the correct number of fragment generators""" - tokenizer = GeneformerTokenizer(tiny_lazy_adata) + tokenizer = GeneformerTokenizer(tiny_slaf) dataset = SLAFIterableDataset( slaf_array=tiny_slaf, @@ -970,9 +962,9 @@ def test_mixture_of_scanners_fragment_generators(self, tiny_slaf, tiny_lazy_adat dataset.batch_processor.fragment_generators ) - def test_mixture_of_scanners_parameter_validation(self, tiny_slaf, tiny_lazy_adata): + def test_mixture_of_scanners_parameter_validation(self, tiny_slaf): """Test MoS parameter validation""" - tokenizer = GeneformerTokenizer(tiny_lazy_adata) + tokenizer = GeneformerTokenizer(tiny_slaf) # Test valid parameters dataset = SLAFIterableDataset( @@ -1033,9 +1025,9 @@ def test_mixture_of_scanners_parameter_validation(self, tiny_slaf, tiny_lazy_ada prefetch_batch_size=10000001, ) - def test_mixture_of_scanners_iteration(self, tiny_slaf, tiny_lazy_adata): + def test_mixture_of_scanners_iteration(self, tiny_slaf): """Test that MoS dataset can iterate through batches""" - tokenizer = GeneformerTokenizer(tiny_lazy_adata) + tokenizer = GeneformerTokenizer(tiny_slaf) dataset = SLAFIterableDataset( slaf_array=tiny_slaf, @@ -1058,9 +1050,9 @@ def test_mixture_of_scanners_iteration(self, tiny_slaf, tiny_lazy_adata): assert batch_count > 0 - def test_mixture_of_scanners_epoch_reset(self, tiny_slaf, tiny_lazy_adata): + def test_mixture_of_scanners_epoch_reset(self, tiny_slaf): """Test that MoS epoch reset works correctly""" - tokenizer = GeneformerTokenizer(tiny_lazy_adata) + tokenizer = GeneformerTokenizer(tiny_slaf) dataset = SLAFIterableDataset( slaf_array=tiny_slaf, @@ -1086,12 +1078,10 @@ def test_mixture_of_scanners_epoch_reset(self, tiny_slaf, tiny_lazy_adata): dataset.batch_processor.fragment_generators ) - def test_mixture_of_scanners_backward_compatibility( - self, tiny_slaf, tiny_lazy_adata - ): + def test_mixture_of_scanners_backward_compatibility(self, tiny_slaf): """Test that MoS is backward compatible (enabled by default) in PrefetchBatchProcessor""" shuffle = RandomShuffle() - tokenizer = ScGPTTokenizer(tiny_lazy_adata) + tokenizer = ScGPTTokenizer(tiny_slaf) # Default behavior (MoS enabled) processor_default = PrefetchBatchProcessor( @@ -1117,7 +1107,7 @@ def test_mixture_of_scanners_backward_compatibility( assert not hasattr(processor_disabled, "fragment_generators") assert hasattr(processor_disabled, "batch_generator") - def test_mixture_of_scanners_with_raw_mode(self, tiny_slaf, tiny_lazy_adata): + def test_mixture_of_scanners_with_raw_mode(self, tiny_slaf): """Test MoS functionality with raw mode""" dataset = SLAFIterableDataset( slaf_array=tiny_slaf, @@ -1143,9 +1133,9 @@ def test_mixture_of_scanners_with_raw_mode(self, tiny_slaf, tiny_lazy_adata): assert batch_count > 0 - def test_mixture_of_scanners_with_fragment_mode(self, tiny_slaf, tiny_lazy_adata): + def test_mixture_of_scanners_with_fragment_mode(self, tiny_slaf): """Test that MoS automatically enables fragment mode""" - tokenizer = GeneformerTokenizer(tiny_lazy_adata) + tokenizer = GeneformerTokenizer(tiny_slaf) dataset = SLAFIterableDataset( slaf_array=tiny_slaf, @@ -1160,9 +1150,9 @@ def test_mixture_of_scanners_with_fragment_mode(self, tiny_slaf, tiny_lazy_adata assert dataset.batch_processor.by_fragment is True assert dataset.batch_processor.use_mixture_of_scanners is True - def test_mixture_of_scanners_random_sampling(self, tiny_slaf, tiny_lazy_adata): + def test_mixture_of_scanners_random_sampling(self, tiny_slaf): """Test that MoS uses random sampling from fragment generators""" - tokenizer = GeneformerTokenizer(tiny_lazy_adata) + tokenizer = GeneformerTokenizer(tiny_slaf) dataset = SLAFIterableDataset( slaf_array=tiny_slaf, @@ -1188,11 +1178,9 @@ def test_mixture_of_scanners_random_sampling(self, tiny_slaf, tiny_lazy_adata): assert "attention_mask" in batch assert "cell_ids" in batch - def test_mixture_of_scanners_generator_exhaustion_handling( - self, tiny_slaf, tiny_lazy_adata - ): + def test_mixture_of_scanners_generator_exhaustion_handling(self, tiny_slaf): """Test that MoS handles generator exhaustion correctly""" - tokenizer = GeneformerTokenizer(tiny_lazy_adata) + tokenizer = GeneformerTokenizer(tiny_slaf) dataset = SLAFIterableDataset( slaf_array=tiny_slaf, @@ -1212,11 +1200,9 @@ def test_mixture_of_scanners_generator_exhaustion_handling( assert batch_count > 0 - def test_mixture_of_scanners_cell_boundary_handling( - self, tiny_slaf, tiny_lazy_adata - ): + def test_mixture_of_scanners_cell_boundary_handling(self, tiny_slaf): """Test that MoS handles cell boundaries correctly""" - tokenizer = GeneformerTokenizer(tiny_lazy_adata) + tokenizer = GeneformerTokenizer(tiny_slaf) dataset = SLAFIterableDataset( slaf_array=tiny_slaf, @@ -1241,11 +1227,11 @@ def test_mixture_of_scanners_cell_boundary_handling( assert batch_count > 0 def test_prefetch_batch_processor_mixture_of_scanners_initialization( - self, tiny_slaf, tiny_lazy_adata + self, tiny_slaf ): """Test PrefetchBatchProcessor initialization with Mixture of Scanners (MoS)""" shuffle = RandomShuffle() - tokenizer = ScGPTTokenizer(tiny_lazy_adata) + tokenizer = ScGPTTokenizer(tiny_slaf) processor = PrefetchBatchProcessor( slaf_array=tiny_slaf, @@ -1272,12 +1258,10 @@ def test_prefetch_batch_processor_mixture_of_scanners_initialization( assert len(processor.generator_last_cells) == len(processor.fragment_generators) assert len(processor.generator_active) == len(processor.generator_active) - def test_prefetch_batch_processor_mos_parameter_validation( - self, tiny_slaf, tiny_lazy_adata - ): + def test_prefetch_batch_processor_mos_parameter_validation(self, tiny_slaf): """Test MoS parameter validation in PrefetchBatchProcessor""" shuffle = RandomShuffle() - tokenizer = ScGPTTokenizer(tiny_lazy_adata) + tokenizer = ScGPTTokenizer(tiny_slaf) # Test valid parameters processor = PrefetchBatchProcessor( @@ -1338,10 +1322,10 @@ def test_prefetch_batch_processor_mos_parameter_validation( prefetch_batch_size=10000001, ) - def test_prefetch_batch_processor_mos_epoch_reset(self, tiny_slaf, tiny_lazy_adata): + def test_prefetch_batch_processor_mos_epoch_reset(self, tiny_slaf): """Test that MoS epoch reset works correctly in PrefetchBatchProcessor""" shuffle = RandomShuffle() - tokenizer = ScGPTTokenizer(tiny_lazy_adata) + tokenizer = ScGPTTokenizer(tiny_slaf) processor = PrefetchBatchProcessor( slaf_array=tiny_slaf, @@ -1363,12 +1347,10 @@ def test_prefetch_batch_processor_mos_epoch_reset(self, tiny_slaf, tiny_lazy_ada assert len(processor.generator_last_cells) == len(processor.fragment_generators) assert len(processor.generator_active) == len(processor.fragment_generators) - def test_prefetch_batch_processor_mos_backward_compatibility( - self, tiny_slaf, tiny_lazy_adata - ): + def test_prefetch_batch_processor_mos_backward_compatibility(self, tiny_slaf): """Test that MoS is backward compatible (enabled by default) in PrefetchBatchProcessor""" shuffle = RandomShuffle() - tokenizer = ScGPTTokenizer(tiny_lazy_adata) + tokenizer = ScGPTTokenizer(tiny_slaf) # Default behavior (MoS enabled) processor_default = PrefetchBatchProcessor( @@ -1394,9 +1376,7 @@ def test_prefetch_batch_processor_mos_backward_compatibility( assert not hasattr(processor_disabled, "fragment_generators") assert hasattr(processor_disabled, "batch_generator") - def test_prefetch_batch_processor_mos_with_raw_mode( - self, tiny_slaf, tiny_lazy_adata - ): + def test_prefetch_batch_processor_mos_with_raw_mode(self, tiny_slaf): """Test MoS functionality with raw mode in PrefetchBatchProcessor""" shuffle = RandomShuffle() tokenizer = None # No tokenizer for raw mode @@ -1419,12 +1399,10 @@ def test_prefetch_batch_processor_mos_with_raw_mode( assert hasattr(batch, "batch_dfs") assert hasattr(batch, "cell_integer_ids") - def test_prefetch_batch_processor_mos_fragment_mode_automatic( - self, tiny_slaf, tiny_lazy_adata - ): + def test_prefetch_batch_processor_mos_fragment_mode_automatic(self, tiny_slaf): """Test that MoS automatically enables fragment mode in PrefetchBatchProcessor""" shuffle = RandomShuffle() - tokenizer = ScGPTTokenizer(tiny_lazy_adata) + tokenizer = ScGPTTokenizer(tiny_slaf) processor = PrefetchBatchProcessor( slaf_array=tiny_slaf, @@ -1439,12 +1417,10 @@ def test_prefetch_batch_processor_mos_fragment_mode_automatic( assert processor.by_fragment is True assert processor.use_mixture_of_scanners is True - def test_prefetch_batch_processor_mos_load_prefetch_batch( - self, tiny_slaf, tiny_lazy_adata - ): + def test_prefetch_batch_processor_mos_load_prefetch_batch(self, tiny_slaf): """Test that MoS can load prefetch batches correctly""" shuffle = RandomShuffle() - tokenizer = ScGPTTokenizer(tiny_lazy_adata) + tokenizer = ScGPTTokenizer(tiny_slaf) processor = PrefetchBatchProcessor( slaf_array=tiny_slaf, @@ -1461,12 +1437,10 @@ def test_prefetch_batch_processor_mos_load_prefetch_batch( assert hasattr(batch, "attention_mask") assert hasattr(batch, "cell_integer_ids") - def test_prefetch_batch_processor_mos_cell_boundary_handling( - self, tiny_slaf, tiny_lazy_adata - ): + def test_prefetch_batch_processor_mos_cell_boundary_handling(self, tiny_slaf): """Test that MoS handles cell boundaries correctly in PrefetchBatchProcessor""" shuffle = RandomShuffle() - tokenizer = ScGPTTokenizer(tiny_lazy_adata) + tokenizer = ScGPTTokenizer(tiny_slaf) processor = PrefetchBatchProcessor( slaf_array=tiny_slaf, diff --git a/tests/test_tokenizers.py b/tests/test_tokenizers.py index c5f7602..7580ca7 100644 --- a/tests/test_tokenizers.py +++ b/tests/test_tokenizers.py @@ -7,37 +7,27 @@ import numpy as np import pandas as pd -import polars as pl import pytest import torch from slaf.core.slaf import SLAFArray -from slaf.core.tabular_schema import SLAF_LANCE_COO_SCHEMA -from slaf.integrations.anndata import LazyAnnData from slaf.ml.tokenizers import GeneformerTokenizer, ScGPTTokenizer -def build_mock_adata(): - 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_adata = Mock(spec=LazyAnnData) - mock_adata.slaf = mock_slaf_array - mock_adata._transformations = {} - return mock_adata, mock_slaf_array - - class TestSLAFTokenizer: """Test the new SLAFTokenizer interface.""" def test_tokenizer_initialization(self): """Test SLAFTokenizer initialization with different tokenizer types.""" - mock_adata, _ = build_mock_adata() + # 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 # Test Geneformer initialization tokenizer = GeneformerTokenizer( - adata=mock_adata, + slaf_array=mock_slaf_array, vocab_size=1000, ) @@ -50,20 +40,25 @@ def test_tokenizer_initialization(self): # Test scGPT initialization tokenizer = ScGPTTokenizer( - adata=mock_adata, + slaf_array=mock_slaf_array, vocab_size=1000, n_expression_bins=5, ) assert tokenizer.n_expression_bins == 5 + assert tokenizer.expr_bin_start == 1000 assert tokenizer.expr_bin_size == 0.2 def test_geneformer_tokenization(self): """Test Geneformer tokenization.""" - mock_adata, _ = build_mock_adata() + # 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 tokenizer = GeneformerTokenizer( - adata=mock_adata, + slaf_array=mock_slaf_array, vocab_size=1000, ) @@ -88,10 +83,14 @@ def test_geneformer_tokenization(self): def test_scgpt_tokenization(self): """Test scGPT tokenization with expressions.""" - mock_adata, _ = build_mock_adata() + # 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 tokenizer = ScGPTTokenizer( - adata=mock_adata, + slaf_array=mock_slaf_array, vocab_size=1000, n_expression_bins=10, ) @@ -129,51 +128,16 @@ 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_adata, _ = build_mock_adata() - - tokenizer = ScGPTTokenizer( - adata=mock_adata, - vocab_size=1000, - n_expression_bins=10, - ) - - grouped_df = pd.DataFrame( - { - "gene_sequence": [[4, 5, 6]], - "expr_sequence": [[1, 5, 9]], - } - ) - 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] == 1 - - def test_geneformer_tokenize_grouped_uses_preencoded_tokens(self): - mock_adata, _ = build_mock_adata() - - tokenizer = GeneformerTokenizer( - adata=mock_adata, - 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_adata, _ = build_mock_adata() + # 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 tokenizer = ScGPTTokenizer( - adata=mock_adata, + slaf_array=mock_slaf_array, vocab_size=1000, ) @@ -191,10 +155,14 @@ def test_scgpt_tokenization_no_expression(self): def test_tokenization_edge_cases(self): """Test edge cases for tokenization.""" - mock_adata, _ = build_mock_adata() + # 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 tokenizer = GeneformerTokenizer( - adata=mock_adata, + slaf_array=mock_slaf_array, vocab_size=1000, ) @@ -216,10 +184,13 @@ def test_tokenization_edge_cases(self): def test_expression_binning(self): """Test expression binning functionality.""" # Mock SLAFArray - mock_adata, _ = build_mock_adata() + 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( - adata=mock_adata, + slaf_array=mock_slaf_array, vocab_size=1000, n_expression_bins=10, ) @@ -227,26 +198,30 @@ 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) == 2 - assert tokenizer._expression_to_bin(0.9) == 10 - assert tokenizer._expression_to_bin(1.0) == 10 # Clipped to last bin + 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 # 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] == 2 - assert bins[2] == 6 - assert bins[3] == 10 + 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[4] == 0 # PAD for -1.0 def test_gene_id_mapping(self): """Test gene ID to token mapping.""" - mock_adata, _ = build_mock_adata() + # 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 tokenizer = GeneformerTokenizer( - adata=mock_adata, + slaf_array=mock_slaf_array, vocab_size=1000, ) @@ -271,10 +246,14 @@ def test_gene_id_mapping(self): def test_vocabulary_info(self): """Test vocabulary information retrieval.""" - mock_adata, _ = build_mock_adata() + # 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 tokenizer = GeneformerTokenizer( - adata=mock_adata, + slaf_array=mock_slaf_array, vocab_size=1000, ) @@ -285,158 +264,16 @@ def test_vocabulary_info(self): # The tokenizer creates a fallback vocabulary assert vocab_info["gene_vocab_size"] > 0 - def test_apply_uses_slaf_runtime_transformations(self): - mock_adata, _ = build_mock_adata() - mock_adata._transformations = { - "normalize_total": { - "type": "normalize_total", - "target_sum": 100.0, - "cell_factors": {0: 25.0, 1: 25.0}, - }, - "log1p": { - "type": "log1p", - "applied": True, - }, - } - - tokenizer = ScGPTTokenizer( - adata=mock_adata, - vocab_size=1000, - n_expression_bins=10, - ) - - df = pl.DataFrame( - { - "cell_integer_id": [0, 0, 1, 1], - "gene_integer_id": [0, 1, 0, 1], - "value": [1.0, 3.0, 2.0, 2.0], - } - ) - - grouped = tokenizer.transform_and_apply( - df, - schema=SLAF_LANCE_COO_SCHEMA, - max_items=2, - use_binned_expressions=False, - ) - - expr_sequences = grouped["expr_sequence"].to_list() - expected = [ - [np.log1p(25.0), np.log1p(75.0)], - [np.log1p(50.0), np.log1p(50.0)], - ] - for actual_seq, expected_seq in zip(expr_sequences, expected, strict=False): - assert np.asarray(actual_seq) == pytest.approx(np.asarray(expected_seq)) - - def test_apply_requires_precomputed_cell_factors_for_normalize_total(self): - mock_adata, _ = build_mock_adata() - mock_adata._transformations = { - "normalize_total": { - "type": "normalize_total", - "target_sum": 100.0, - }, - } - - tokenizer = ScGPTTokenizer( - adata=mock_adata, - vocab_size=1000, - n_expression_bins=10, - ) - - df = pl.DataFrame( - { - "cell_integer_id": [0, 0, 1, 1], - "gene_integer_id": [0, 1, 0, 1], - "value": [1.0, 3.0, 2.0, 2.0], - } - ) - - with pytest.raises( - ValueError, - match="requires precomputed cell_factors keyed by cell_integer_id", - ): - tokenizer.transform_and_apply( - df, - schema=SLAF_LANCE_COO_SCHEMA, - max_items=2, - use_binned_expressions=False, - ) - - def test_apply_uses_precomputed_cell_factors_by_integer_id(self): - mock_adata, _ = build_mock_adata() - mock_adata._transformations = { - "normalize_total": { - "type": "normalize_total", - "target_sum": 100.0, - "cell_factors": {0: 10.0, 1: 5.0}, - }, - } - - tokenizer = ScGPTTokenizer( - adata=mock_adata, - vocab_size=1000, - n_expression_bins=10, - ) - - df = pl.DataFrame( - { - "cell_integer_id": [0, 0, 1, 1], - "gene_integer_id": [0, 1, 0, 1], - "value": [1.0, 3.0, 2.0, 2.0], - } - ) - - grouped = tokenizer.transform_and_apply( - df, - schema=SLAF_LANCE_COO_SCHEMA, - max_items=2, - use_binned_expressions=False, - ) - - expr_sequences = grouped["expr_sequence"].to_list() - expected = [[10.0, 30.0], [10.0, 10.0]] - for actual_seq, expected_seq in zip(expr_sequences, expected, strict=False): - assert np.asarray(actual_seq) == pytest.approx(np.asarray(expected_seq)) - - def test_scgpt_window_does_not_double_log1p(self): - mock_adata, _ = build_mock_adata() - mock_adata._transformations = { - "log1p": { - "type": "log1p", - "applied": True, - }, - } - - tokenizer = ScGPTTokenizer( - adata=mock_adata, - vocab_size=1000, - n_expression_bins=10, - ) - - df = pl.DataFrame( - { - "cell_integer_id": [0, 0], - "gene_integer_id": [0, 1], - "value": [np.e - 1.0, np.exp(2.0) - 1.0], - } - ) - - grouped = tokenizer.transform_and_apply( - df, - schema=SLAF_LANCE_COO_SCHEMA, - max_items=2, - use_binned_expressions=True, - ) - - expr_tokens = grouped["expr_sequence"].to_list()[0] - assert expr_tokens == [6, 10] - def test_token_decoding(self): """Test token decoding functionality.""" - mock_adata, _ = build_mock_adata() + # 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 tokenizer = ScGPTTokenizer( - adata=mock_adata, + slaf_array=mock_slaf_array, vocab_size=1000, n_expression_bins=10, ) @@ -458,7 +295,9 @@ def test_token_decoding(self): assert "PAD" in decoded["special_tokens"] # Test decoding scGPT tokens - tokens = [1] + gene_tokens + [2, 0] # CLS, gene1, gene2, SEP, PAD + tokens = ( + [1] + gene_tokens + [1001, 1005] + [2, 0] + ) # CLS, gene1, gene2, expr1, expr2, SEP, PAD decoded = tokenizer.decode_tokens(tokens) # Check structure @@ -476,9 +315,8 @@ class TestSLAFTokenizerWithRealData: def test_tokenizer_with_real_data(self, tiny_slaf): """Test tokenizer with real SLAF data.""" - adata = LazyAnnData(tiny_slaf) tokenizer = GeneformerTokenizer( - adata=adata, + slaf_array=tiny_slaf, vocab_size=1000, ) @@ -497,9 +335,8 @@ def test_tokenizer_with_real_data(self, tiny_slaf): def test_scgpt_with_real_data(self, tiny_slaf): """Test scGPT tokenizer with real SLAF data.""" - adata = LazyAnnData(tiny_slaf) tokenizer = ScGPTTokenizer( - adata=adata, + slaf_array=tiny_slaf, vocab_size=1000, n_expression_bins=10, ) @@ -520,9 +357,8 @@ def test_scgpt_with_real_data(self, tiny_slaf): def test_gene_mapping_with_real_data(self, tiny_slaf): """Test gene ID mapping with real SLAF data.""" - adata = LazyAnnData(tiny_slaf) tokenizer = GeneformerTokenizer( - adata=adata, + slaf_array=tiny_slaf, vocab_size=1000, ) diff --git a/uv.lock b/uv.lock index 1b9a19a..e4a542d 100644 --- a/uv.lock +++ b/uv.lock @@ -227,12 +227,6 @@ 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" @@ -1354,28 +1348,28 @@ wheels = [ [[package]] name = "google-auth" -version = "2.52.0" +version = "2.49.2" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "cryptography" }, { name = "pyasn1-modules" }, ] -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" } +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" } wheels = [ - { 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" }, + { 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" }, ] [[package]] name = "google-cloud-core" -version = "2.6.0" +version = "2.5.1" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "google-api-core" }, { name = "google-auth" }, ] -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" } +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" } wheels = [ - { 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" }, + { 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" }, ] [[package]] @@ -1427,26 +1421,26 @@ wheels = [ [[package]] name = "google-resumable-media" -version = "2.9.0" +version = "2.8.2" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "google-crc32c" }, ] -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" } +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" } wheels = [ - { 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" }, + { 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" }, ] [[package]] name = "googleapis-common-protos" -version = "1.75.0" +version = "1.74.0" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "protobuf" }, ] -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" } +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" } wheels = [ - { 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" }, + { 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" }, ] [[package]] @@ -3624,19 +3618,6 @@ 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" @@ -4052,14 +4033,14 @@ wheels = [ [[package]] name = "proto-plus" -version = "1.28.0" +version = "1.27.2" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "protobuf" }, ] -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" } +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" } wheels = [ - { 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" }, + { 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" }, ] [[package]] @@ -5538,13 +5519,11 @@ docs = [ full = [ { name = "igraph" }, { name = "leidenalg" }, - { name = "omegaconf" }, { name = "tiledb" }, { name = "tiledbsoma" }, { name = "torch" }, ] ml = [ - { name = "omegaconf" }, { name = "tiledb" }, { name = "tiledbsoma" }, { name = "torch" }, @@ -5602,7 +5581,6 @@ 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 350fe1da844be6a31dffd13d3c56d82a3e66316a Mon Sep 17 00:00:00 2001 From: Pavan Ramkumar Date: Thu, 28 May 2026 12:56:31 -0700 Subject: [PATCH 10/10] Add SLAFTokenizer.tokenize_grouped for grouped DataFrame input. Centralize schema-based extraction in the tokenizer so datasets and the distributed worker delegate to tokenizer_instance.tokenize_grouped. Co-authored-by: Cursor --- slaf/distributed/worker.py | 50 +++++++++++++------------------------- slaf/ml/datasets.py | 10 +++----- slaf/ml/tokenizers.py | 18 ++++++++++++++ tests/test_tokenizers.py | 48 ++++++++++++++++++++++++++++++++++++ 4 files changed, 86 insertions(+), 40 deletions(-) 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 49eb36d..56034e0 100644 --- a/slaf/ml/datasets.py +++ b/slaf/ml/datasets.py @@ -1059,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 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_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