Fit processors for schema fields that were not supplied

docs/api/processors.rst

@@ -276,6 +276,31 @@ For maximum control, you can pass pre-configured processor instances directly in
             return samples
 
 
+Reusing Fitted Processors
+-------------------------
+
+``set_task()`` and ``create_sample_dataset()`` accept already-fitted processors
+through ``input_processors`` and ``output_processors``. A supplied processor is
+used as it is and never refitted; every schema field you do not supply gets a
+processor fitted on the samples as usual. A common use is fitting on training
+samples only and reusing those processors for validation and test data, so test
+patients do not shape preprocessing:
+
+.. code-block:: python
+
+    from pyhealth.datasets import create_sample_dataset
+
+    train = create_sample_dataset(train_samples, input_schema, output_schema)
+    test = create_sample_dataset(
+        test_samples, input_schema, output_schema,
+        input_processors=train.input_processors,
+        output_processors=train.output_processors,
+    )
+
+In PyHealth 2.0.2 and earlier, supplying a processor for some fields fitted
+none for the rest of that schema: those fields reached the model as raw Python
+values, without an error. See ``examples/reuse_train_processors.py``.
+
 Processor String Keys
 ---------------------
 

examples/reuse_train_processors.py

@@ -0,0 +1,52 @@
+"""Fit processors on training samples and reuse them for test samples.
+
+Fitting processors only on training data keeps test patients from shaping
+preprocessing (vocabularies, and any statistics a processor learns). Pass the
+training dataset's processors when building the test dataset: supplied
+processors are reused as they are, never refitted.
+
+Runs in a few seconds on CPU with synthetic data; no download.
+
+Usage:
+    python examples/reuse_train_processors.py
+"""
+
+from pyhealth.datasets import create_sample_dataset
+
+INPUT_SCHEMA = {"codes": "sequence", "age": "tensor"}
+OUTPUT_SCHEMA = {"label": "binary"}
+
+train_samples = [
+    {"patient_id": "p1", "codes": ["I10", "E11"], "age": [61.0], "label": 1},
+    {"patient_id": "p2", "codes": ["J45"], "age": [34.0], "label": 0},
+    {"patient_id": "p3", "codes": ["E11", "N18"], "age": [72.0], "label": 1},
+]
+test_samples = [
+    # "C50" never occurs in training, so it maps to <unk>, as at deployment
+    {"patient_id": "p9", "codes": ["I10", "C50"], "age": [58.0], "label": 0},
+]
+
+
+def main():
+    train = create_sample_dataset(
+        train_samples, INPUT_SCHEMA, OUTPUT_SCHEMA, dataset_name="train"
+    )
+    codes = train.input_processors["codes"]
+    print("training vocabulary:", sorted(codes.code_vocab))
+
+    # Reuse every fitted processor. (A field left out of input_processors would
+    # get a processor fitted on the samples passed here, i.e. on test data.)
+    test = create_sample_dataset(
+        test_samples,
+        INPUT_SCHEMA,
+        OUTPUT_SCHEMA,
+        input_processors=train.input_processors,
+        output_processors=train.output_processors,
+        dataset_name="test",
+    )
+    print("same codes processor:", test.input_processors["codes"] is codes)
+    print("test sample:", test[0])
+
+
+if __name__ == "__main__":
+    main()

pyhealth/datasets/sample_dataset.py

@@ -49,6 +49,29 @@ class SampleBuilder:
     After saving the schema, `litdata.optimize` can be used with `builder.transform`
     to serialize and chunk pickled sample items into a directory that can be
     loaded via SampleDataset.
+
+    Pre-fitted processors passed as `input_processors` / `output_processors` are
+    used as they are; every schema field without one gets a fitted processor.
+
+    Examples:
+        >>> from pyhealth.datasets.sample_dataset import SampleBuilder
+        >>> from pyhealth.processors import SequenceProcessor
+        >>> samples = [
+        ...     {"patient_id": "p1", "codes": ["a", "b"], "age": [40.0], "label": 1},
+        ...     {"patient_id": "p2", "codes": ["c"], "age": [55.0], "label": 0},
+        ... ]
+        >>> codes = SequenceProcessor()
+        >>> codes.fit(samples, "codes")
+        >>> builder = SampleBuilder(
+        ...     input_schema={"codes": "sequence", "age": "tensor"},
+        ...     output_schema={"label": "binary"},
+        ...     input_processors={"codes": codes},
+        ... )
+        >>> builder.fit(samples)
+        >>> builder.input_processors["codes"] is codes
+        True
+        >>> type(builder.input_processors["age"]).__name__
+        'TensorProcessor'
     """
 
     def __init__(
@@ -61,10 +84,10 @@ def __init__(
         self.input_schema = input_schema
         self.output_schema = output_schema
         self._input_processors = (
-            input_processors if input_processors is not None else {}
+            dict(input_processors) if input_processors is not None else {}
         )
         self._output_processors = (
-            output_processors if output_processors is not None else {}
+            dict(output_processors) if output_processors is not None else {}
         )
         self._patient_to_index: Dict[str, List[int]] = {}
         self._record_to_index: Dict[str, List[int]] = {}
@@ -145,9 +168,9 @@ def fit(
             - Builds `patient_to_index` and `record_to_index` mappings by
               recording the sample indices associated with `patient_id` and
               `record_id`/`visit_id` fields.
-            - Instantiates and fits input/output processors from the provided
-              schemas (unless pre-fitted processors were supplied to the
-              constructor).
+            - Instantiates and fits a processor for every schema field that
+              has none. Pre-fitted processors supplied to the constructor are
+              used as they are; any remaining fields still get fitted ones.
         """
         # Validate the samples
         input_keys = set(self.input_schema.keys())
@@ -171,17 +194,17 @@ def fit(
             if record_id is not None:
                 self._record_to_index.setdefault(record_id, []).append(i)
 
-        # Fit processors if they were not provided
-        if not self._input_processors:
-            for key, spec in self.input_schema.items():
-                processor = self._get_processor_instance(spec)
-                processor.fit(samples, key)
-                self._input_processors[key] = processor
-        if not self._output_processors:
-            for key, spec in self.output_schema.items():
-                processor = self._get_processor_instance(spec)
-                processor.fit(samples, key)
-                self._output_processors[key] = processor
+        # Fit a processor for every schema field that has none. Supplied
+        # (pre-fitted) processors are used as they are, never refitted.
+        for schema, processors in (
+            (self.input_schema, self._input_processors),
+            (self.output_schema, self._output_processors),
+        ):
+            for key, spec in schema.items():
+                if key not in processors:
+                    processor = self._get_processor_instance(spec)
+                    processor.fit(samples, key)
+                    processors[key] = processor
 
         self._fitted = True
 

tests/core/test_partial_processors.py

@@ -0,0 +1,85 @@
+"""Supplying some processors must still fit the others.
+
+SampleBuilder used to fit input (or output) processors only when none were
+supplied. Supplying a processor for one field left every other field of that
+side without a processor, so its raw value (e.g. a Python list) reached the
+model without any error.
+"""
+
+import pickle
+import unittest
+
+import torch
+
+from pyhealth.datasets import create_sample_dataset
+from pyhealth.datasets.sample_dataset import SampleBuilder
+from pyhealth.processors import (
+    BinaryLabelProcessor,
+    SequenceProcessor,
+    TensorProcessor,
+)
+
+
+def _samples():
+    return [
+        {"patient_id": "p1", "codes": ["a", "b"], "age": [40.0], "label": 1},
+        {"patient_id": "p2", "codes": ["c"], "age": [55.0], "label": 0},
+    ]
+
+
+class TestPartialProcessors(unittest.TestCase):
+    def setUp(self):
+        self.codes = SequenceProcessor()
+        self.codes.fit(_samples(), "codes")
+
+    def test_missing_input_processor_is_fitted(self):
+        builder = SampleBuilder(
+            input_schema={"codes": "sequence", "age": "tensor"},
+            output_schema={"label": "binary"},
+            input_processors={"codes": self.codes},
+        )
+        builder.fit(_samples())
+        self.assertIs(builder.input_processors["codes"], self.codes)
+        self.assertIsInstance(builder.input_processors["age"], TensorProcessor)
+        out = builder.transform({"sample": pickle.dumps(_samples()[0])})
+        self.assertIsInstance(out["age"], torch.Tensor)
+
+    def test_supplied_processor_is_not_refitted(self):
+        vocab = dict(self.codes.code_vocab)
+        samples = _samples() + [
+            {"patient_id": "p3", "codes": ["z"], "age": [60.0], "label": 1}
+        ]
+        builder = SampleBuilder(
+            input_schema={"codes": "sequence", "age": "tensor"},
+            output_schema={"label": "binary"},
+            input_processors={"codes": self.codes},
+        )
+        builder.fit(samples)
+        self.assertEqual(self.codes.code_vocab, vocab)  # "z" was not added
+
+    def test_missing_output_processor_is_fitted(self):
+        samples = [dict(s, label2=s["label"]) for s in _samples()]
+        label = BinaryLabelProcessor()
+        label.fit(samples, "label")
+        builder = SampleBuilder(
+            input_schema={"codes": "sequence"},
+            output_schema={"label": "binary", "label2": "binary"},
+            output_processors={"label": label},
+        )
+        builder.fit(samples)
+        self.assertIs(builder.output_processors["label"], label)
+        self.assertIn("label2", builder.output_processors)
+
+    def test_create_sample_dataset_with_partial_processors(self):
+        dataset = create_sample_dataset(
+            samples=_samples(),
+            input_schema={"codes": "sequence", "age": "tensor"},
+            output_schema={"label": "binary"},
+            input_processors={"codes": self.codes},
+            dataset_name="test_partial_processors",
+        )
+        self.assertIsInstance(dataset[0]["age"], torch.Tensor)
+
+
+if __name__ == "__main__":
+    unittest.main()