Towards Unified Sequence Representation I: Add SequenceProcessorV2 for packed hierarchical records

docs/api/processors.rst

@@ -24,6 +24,7 @@ Available Processors
 
 **Sequence Processors:**
 
+- :doc:`SequenceProcessorV2 <processors/pyhealth.processors.SequenceProcessorV2>`: Unified processor for hierarchical packed records (codes, continuous values, and timestamps at any nesting depth) without padding or truncation
 - ``SequenceProcessor``: For categorical sequences (e.g., medical codes like diagnoses, procedures)
 - ``NestedSequenceProcessor``: For nested categorical sequences (e.g., drug recommendation with visit history)
 - ``NestedMultiHotProcessor``: For nested categorical sequences as per-visit multi-hot vectors (e.g., generative EHR models)
@@ -312,9 +313,11 @@ How Processed Samples Are Stored
 ``set_task()`` and ``create_sample_dataset(..., in_memory=False)`` write processed
 samples to a disk cache (litdata). Each field is stored according to its processor:
 
-- Processors that always return one ``torch.Tensor`` declare ``stores_tensor = True``
-  (``sequence``, ``nested_sequence``, ``tensor``, ``timeseries``, the label
-  processors, ``image``, ``audio`` and others). Their outputs are stored as tensors.
+- Processors that always return a ``torch.Tensor`` — or a tuple of tensors with
+  the same length for every sample — declare ``stores_tensor = True``
+  (``sequence_v2``, ``sequence``, ``nested_sequence``, ``tensor``, ``timeseries``,
+  the label processors, ``image``, ``audio`` and others). Their outputs are
+  stored natively as tensors.
 - Every other field — ``raw``, ``text``, ``stagenet``, ``tuple_time_text``,
   ``graph``, ``temporal_timeseries``, ``time_image``, and keys that are not in the
   schema at all, such as ``patient_id`` or extra keys a task adds — is stored in
@@ -329,23 +332,27 @@ already holds pickled processors in ``schema.pkl``: only load caches you created
 trust.
 
 A custom processor should set ``stores_tensor = True`` only if ``process()`` always
-returns a single tensor; if such a field ever returns something else, writing the
-cache fails with an error naming the field.
+returns a single tensor or a fixed-length tuple of tensors; if such a field ever
+returns something else, writing the cache fails with an error naming the field.
 
 In PyHealth 2.0.2 and earlier, non-tensor fields were handed to litdata as they
 were. litdata infers the layout of every sample from the first one it writes, so a
 value whose length, shape or type varied between samples failed to write, or was
 read back wrong. Caches written by those versions are rebuilt automatically (the
 storage layout is part of the cache key).
 
-The default collate function treats any 2-tuple as a ``(time, values)`` pair of
-tensors, so batch raw fields holding 2-tuples with your own ``collate_fn``.
+Processors whose batches need custom packing (such as ``SequenceProcessorV2``)
+define ``collate(items)``, which ``get_dataloader`` uses automatically for those
+fields. For fields without a processor-level ``collate``, the default collate
+function treats any 2-tuple as a ``(time, values)`` pair of tensors, so batch raw
+fields holding 2-tuples with your own ``collate_fn``.
 
 Processor String Keys
 ---------------------
 
 Common string keys for automatic processor selection:
 
+- ``"sequence_v2"``: For unified hierarchical packed records (codes, continuous values, and timestamps)
 - ``"temporal_timeseries"``: For time-series data with preserved timestamps (use in place of ``"timeseries"`` when building ``UnifiedMultimodalEmbeddingModel``)
 - ``"sequence"``: For categorical sequences (medical codes)
 - ``"nested_sequence"``: For nested categorical sequences (visit history)
@@ -543,6 +550,7 @@ API Reference
     processors/pyhealth.processors.SampleProcessor
     processors/pyhealth.processors.DatasetProcessor
     processors/pyhealth.processors.SequenceProcessor
+    processors/pyhealth.processors.SequenceProcessorV2
     processors/pyhealth.processors.NestedSequenceProcessor
     processors/pyhealth.processors.NestedMultiHotProcessor
     processors/pyhealth.processors.NestedFloatsProcessor

pyhealth/datasets/sample_dataset.py

@@ -353,10 +353,11 @@ def encode(self, transformed: dict[str, Any]) -> dict[str, Any]:
         litdata infers one serializer per value from the first sample it
         writes and expects every later sample to have the same structure, so a
         list whose length varies between samples cannot be written as-is.
-        Fields whose processor declares ``stores_tensor`` stay tensors; every
-        other field (including keys not in the schema, such as ``patient_id``)
-        goes into one pickled payload, which has the same structure in every
-        sample whatever the values look like. :class:`SampleDataset` decodes it
+        Fields whose processor declares ``stores_tensor`` stay tensors (or
+        tuples of tensors, one per ``schema()`` entry); every other field
+        (including keys not in the schema, such as ``patient_id``) goes into
+        one pickled payload, which has the same structure in every sample
+        whatever the values look like. :class:`SampleDataset` decodes it
         on read, so callers get back the original Python objects.
 
         Args:
@@ -366,20 +367,34 @@ def encode(self, transformed: dict[str, Any]) -> dict[str, Any]:
             A record with the tensor fields plus one bytes payload.
 
         Raises:
-            ValueError: If a field declared ``stores_tensor`` is not a tensor.
+            ValueError: If a field declared ``stores_tensor`` is not a tensor
+                or a tuple of tensors, one per entry of its processor's
+                ``schema()``.
         """
         tensor_fields = self._tensor_fields()
+        processors = {**self._input_processors, **self._output_processors}
         record: dict[str, Any] = {}
         objects: dict[str, Any] = {}
         for key, value in transformed.items():
             if key in tensor_fields:
-                if not isinstance(value, torch.Tensor):
-                    processor = self._input_processors.get(key) or self._output_processors.get(key)
+                is_tensors = isinstance(value, tuple) and all(
+                    isinstance(item, torch.Tensor) for item in value
+                )
+                processor = processors[key]
+                if not (isinstance(value, torch.Tensor) or is_tensors):
                     raise ValueError(
                         f"Field '{key}': {type(processor).__name__} declares "
                         f"stores_tensor=True but process() returned "
                         f"{type(value).__name__}. Set stores_tensor = False on the "
-                        "processor, or return a torch.Tensor."
+                        "processor, or return a torch.Tensor or a tuple of them."
+                    )
+                # litdata expects the same structure in every sample, and
+                # schema() fixes the tuple's length for all of them.
+                if is_tensors and len(value) != len(processor.schema()):
+                    raise ValueError(
+                        f"Field '{key}': {type(processor).__name__}.process() "
+                        f"returned {len(value)} tensors, but its schema() has "
+                        f"{len(processor.schema())} entries."
                     )
                 record[key] = value
             else:

pyhealth/datasets/utils.py

@@ -1,11 +1,13 @@
+import functools
 import hashlib
 import os
 import pickle
+from collections.abc import Callable, Mapping
 from datetime import datetime
 from typing import Any, Dict, List, Optional, Tuple
 
-import torch
 import litdata
+import torch
 from dateutil.parser import parse as dateutil_parse
 from torch.nn.utils.rnn import pad_sequence
 from torch.utils.data import DataLoader
@@ -250,25 +252,42 @@ def collate_fn_dict(batch: List[dict]) -> dict:
     return {key: [d[key] for d in batch] for key in batch[0]}
 
 
-def collate_fn_dict_with_padding(batch: List[dict]) -> dict:
+def collate_fn_dict_with_padding(
+    batch: list[dict], collators: Mapping[str, Callable] | None = None
+) -> dict:
     """Collates a batch of data into a dictionary with padding for tensor values.
 
     Args:
         batch: List of dictionaries, where each dictionary represents a data sample.
+        collators: Optional collate function per key, given the key's values
+            in the batch. Keys with one are collated only by it.
 
     Returns:
         A dictionary where each key corresponds to a list of values from the batch.
         Tensor values are padded to the same shape.
         Tuples of (time, values) from temporal processors are collated separately.
+
+    Examples:
+        >>> import torch
+        >>> batch = [{"x": torch.tensor([1, 2]), "id": "a"}, {"x": torch.tensor([3]), "id": "b"}]
+        >>> collate_fn_dict_with_padding(batch)["x"]
+        tensor([[1, 2],
+                [3, 0]])
+        >>> collate_fn_dict_with_padding(batch)["id"]
+        ['a', 'b']
+        >>> collate_fn_dict_with_padding(batch, collators={"x": len})["x"]
+        2
     """
     collated = {}
     keys = batch[0].keys()
 
     for key in keys:
         values = [sample[key] for sample in batch]
 
+        if collators is not None and key in collators:
+            collated[key] = collators[key](values)
         # Check if this is a temporal feature tuple (time, values)
-        if isinstance(values[0], tuple) and len(values[0]) == 2:
+        elif isinstance(values[0], tuple) and len(values[0]) == 2:
             # Handle (time, values) tuples from processors
             time_tensors = [v[0] for v in values]
             value_tensors = [v[1] for v in values]
@@ -403,13 +422,25 @@ def get_dataloader(
     from .sample_dataset import SampleDataset
 
     if isinstance(dataset, SampleDataset):
+        # Fields whose processor defines collate() are batched by it.
+        processors = {**dataset.input_processors, **dataset.output_processors}
+        collators = {
+            key: processor.collate
+            for key, processor in processors.items()
+            if callable(getattr(processor, "collate", None))
+        }
+        collate_fn = (
+            functools.partial(collate_fn_dict_with_padding, collators=collators)
+            if collators
+            else collate_fn_dict_with_padding
+        )
         return _SeededDataLoader(
             dataset,
             shuffle_samples=shuffle,
             seed=seed,
             batch_size=batch_size,
             generator=generator,
-            collate_fn=collate_fn_dict_with_padding,
+            collate_fn=collate_fn,
         )
     if hasattr(dataset, "set_shuffle"):  # another iterable dataset, e.g. litdata's
         dataset.set_shuffle(shuffle)

pyhealth/medcode/cross_map.py

@@ -187,6 +187,19 @@ def unmapped_codes(self):
         """Source codes seen so far that had no target in this mapping."""
         return getattr(self.mapping, "unmapped_codes", frozenset())
 
+    @property
+    def has_many_targets(self):
+        """Whether some source code maps to more than one target code.
+
+        False means the mapping is single-valued (x-to-one): `map()` returns
+        at most one code. This describes the table, so it is True even if only
+        a few codes have several targets.
+        """
+        if isinstance(self.mapping, _ICDMappingsBackend):
+            # icd-mappings returns a single (primary) target for every code.
+            return False
+        return any(len(targets) > 1 for targets in self.mapping.values())
+
     def __repr__(self):
         return f"CrossMap(source_vocabulary={self.s_vocab}, source_class={self.s_class} target_vocabulary={self.t_vocab}, target_class={self.t_class})"
 

pyhealth/processors/__init__.py

@@ -37,6 +37,14 @@ def get_processor(name: str):
 )
 from .raw_processor import RawProcessor
 from .sequence_processor import SequenceProcessor
+from .sequence_processor_v2 import (
+    Impute,
+    LevelConfig,
+    Normalize,
+    Records,
+    Resample,
+    SequenceProcessorV2,
+)
 from .signal_processor import SignalProcessor
 from .stagenet_processor import (
     StageNetProcessor,
@@ -76,6 +84,12 @@ def get_processor(name: str):
     "NestedSequenceProcessor",
     "RawProcessor",
     "SequenceProcessor",
+    "SequenceProcessorV2",
+    "Records",
+    "LevelConfig",
+    "Impute",
+    "Normalize",
+    "Resample",
     "SignalProcessor",
     "StageNetProcessor",
     "StageNetTensorProcessor",

pyhealth/processors/base_processor.py

@@ -70,11 +70,16 @@ class FeatureProcessor(Processor):
     warns unless it is fitted on a training split
     (``set_task(task, split=PatientSplit(...))``).
 
-    Set ``stores_tensor = True`` only if ``process()`` always returns a single
-    ``torch.Tensor``. Such fields are written to the disk cache as tensors;
-    every other output (lists, tuples, dicts, strings, None, objects) is
-    stored as one pickled value per sample, so it may vary freely in length
-    and shape between samples.
+    Set ``stores_tensor = True`` only if ``process()`` always returns a
+    ``torch.Tensor``, or a tuple with one tensor per :meth:`schema` entry
+    (each tensor may change shape between samples). Such fields are written
+    to the disk cache as tensors; every other output (lists, other tuples,
+    dicts, strings, None, objects) is stored as one pickled value per sample,
+    so it may vary freely in length and shape between samples.
+
+    A processor whose batches need more than padding can define
+    ``collate(items)``, taking the ``process()`` outputs of one batch;
+    :func:`~pyhealth.datasets.get_dataloader` then uses it for that field.
 
     Examples:
         >>> import torch