Feature/tokenizer update

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"]
@@ -99,6 +104,7 @@ test = [
     "psutil>=6.0.0",
     "torch>=2.5.0",
     "tiledbsoma>=1.17.1",
+    "omegaconf>=2.3.0",
 ]
 
 # Full installation includes all optional dependencies

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,
+    processor_config: DictConfig,
     queue: Any,  # Modal Queue or any queue-like object
     n_scanners: int = 8,
     prefetch_batch_count: int = 32,
@@ -56,23 +56,32 @@ 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
@@ -92,75 +101,39 @@ def prefetch_worker(
     # - Out of the box: uses Window/Shuffle from slaf.distributed if no factory config
     # - No hardcoded dependencies on slaf.ml in slaf.distributed code
 
-    # Tokenizer is passed as a factory function name (will be created in ml/distributed.py).
-    # Built before window when use_tokenizer_window is set so tokenizer_instance.window exists.
     tokenizer_instance = None
     tokenizer_fn = None
-    if processor_config.get("tokenizer_factory"):
-        # Dynamic import and factory call
-        tokenizer_config = processor_config["tokenizer_factory"]
-        tokenizer_module = importlib.import_module(tokenizer_config["module"])
-        tokenizer_class = getattr(tokenizer_module, tokenizer_config["class"])
-
-        # SLAFTokenizer needs a slaf_array, so we need to recreate it from the data source path
-        # Extract slaf_path from data_source_config (assumes Lance path is under slaf_path/expression.lance)
-        if data_source_config["type"] == "lance":
-            lance_path = data_source_config["path"]
-            # Assume lance_path is like "path/to/slaf/expression.lance"
+    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
 
             slaf_array = SLAFArray(slaf_path, load_metadata=False)
 
-            # Create tokenizer instance
+            target = tokenizer_config.type
+            module_name, class_name = target.rsplit(".", 1)
+            tokenizer_module = importlib.import_module(module_name)
+            tokenizer_class = getattr(tokenizer_module, class_name)
+            tokenizer_kwargs = (
+                dict(tokenizer_config.args)
+                if "args" in tokenizer_config and tokenizer_config.args is not None
+                else {}
+            )
             tokenizer_instance = tokenizer_class(
-                slaf_array=slaf_array, **tokenizer_config["kwargs"]
+                slaf_array=slaf_array, **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,64 +147,54 @@ 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)
     from slaf.distributed.boundary import GroupBoundaryHandler
 
     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 +204,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)
@@ -330,7 +293,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(

slaf/integrations/scanpy.py

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

slaf/ml/aggregators.py

@@ -109,15 +109,16 @@ def apply(
                     )
                 )
                 .with_columns(
-                    pl.when(pl.col("log_value") > 0)
+                    # TODO: Decide whether scGPT expression bins should use log_value
+                    # or the transformed raw value column for normalization.
+                    pl.when(pl.col(vk) > 0)
                     .then(
-                        (
-                            pl.col("log_value")
-                            * n_expression_bins
-                            / pl.col("log_value").max().over(gk)
+                        1
+                        + (
+                            (pl.col(vk) * n_expression_bins / pl.col(vk).max().over(gk))
+                            .floor()
+                            .clip(0, n_expression_bins - 1)
                         )
-                        .floor()
-                        .clip(0, n_expression_bins - 1)
                     )
                     .otherwise(0)
                     .alias("expr_bin")