ChEB-AI · aditya0by0 · Aug 15, 2026
diff --git a/chebai/preprocessing/collate.py b/chebai/preprocessing/collate.py
@@ -61,10 +61,13 @@ def __call__(self, data: List[Union[Dict, Tuple]]) -> XYData:
         """
         Collate ragged data samples (i.e., samples of unequal size, such as molecular sequences) into a batch.
 
-        Handles both fully and partially labeled data, where some samples may have `None` as their label. The indices
-        of non-null labels are stored in the `non_null_labels` field, which is used to filter out predictions for
-        unlabeled data during evaluation (e.g., F1, MSE). For models supporting partially labeled data, this method
-        ensures alignment between features and labels. Missing labels are passed as a loss keyword.
+        Handles both fully and partially labeled data by use of the following fields in the returned XYData:
+
+        `non_null_labels`: Stores batch row indices of samples where the whole `labels` field is not None, like [0, 2].
+            - Example: [[True, False], None, [False, None]] would result in `non_null_labels` = [0, 2].
+            - This is used to filter out predictions for unlabeled samples during evaluation.
+
+        `missing_labels`: Stores a per-sample, per-label-position boolean mask for unknown entries inside a label row, like the None in [1, None, 0].
 
         Args:
             data (List[Union[Dict, Tuple]]): List of ragged data samples. Each sample can be a dictionary or tuple
@@ -81,38 +84,49 @@ def __call__(self, data: List[Union[Dict, Tuple]]) -> XYData:
         if isinstance(data[0], tuple):
             # For legacy data
             x, y, idents = zip(*data)
-            missing_labels = None
         else:
             x, y, idents = zip(
                 *((d["features"], d["labels"], d.get("ident")) for d in data)
             )
-            missing_labels = [
-                d.get(
-                    "missing_labels",
-                    [False for _ in y[0]] if y[0] is not None else [False],
-                )
-                for d in data
-            ]
 
+        # Compute the per-sample, per-label-position boolean mask for unknown entries
+        # (e.g., the None in [1, None, 0]) on the *original* labels, before any
+        # filtering/padding is applied to `y`. Rows whose entire label is None are
+        # represented as all-False rows of the maximum label length.
+        if any(labels is not None for labels in y):
+            max_label_len = max(len(labels) for labels in y if labels is not None)
+            missing_labels = pad_sequence(
+                [
+                    torch.tensor([label is None for label in labels])
+                    if labels is not None
+                    else torch.zeros(max_label_len, dtype=torch.bool)
+                    for labels in y
+                ],
+                batch_first=True,
+            )
+        else:
+            missing_labels = torch.tensor([])
+
+        # Typical y: ([True, False], None, [True, None], [True])
         if any(x is not None for x in y):
-            # If any label is not None: (None, None, `1`, None)
+            # If any label is not None: (None, None, `[True, None]`, None)
             if any(x is None for x in y):
-                # If any label is None: (`None`, `None`, 1, `None`)
+                # If any label is None: (`None`, [True, False], [True], [False])
                 non_null_labels = [i for i, r in enumerate(y) if r is not None]
                 y = self.process_label_rows(
                     tuple(ye for i, ye in enumerate(y) if i in non_null_labels)
                 )
                 loss_kwargs["non_null_labels"] = non_null_labels
             else:
-                # If all labels are not None: (`0`, `2`, `1`, `3`)
+                # If all labels are not None: (`[True, False]`, `[False, True, True]`, `[False]`, `[True]`)
                 y = self.process_label_rows(y)
 
         else:
             # If all labels are None : (`None`, `None`, `None`, `None`)
             y = None
             loss_kwargs["non_null_labels"] = []
 
-        loss_kwargs["missing_labels"] = torch.tensor(missing_labels)
+        loss_kwargs["missing_labels"] = missing_labels
         # Calculate the lengths of each sequence, create a binary mask for valid (non-padded) positions
         lens = torch.tensor(list(map(len, x)))
         model_kwargs["mask"] = torch.arange(max(lens))[None, :] < lens[:, None]

diff --git a/chebai/preprocessing/reader.py b/chebai/preprocessing/reader.py
@@ -99,11 +99,6 @@ def _read_components(self, row: Dict[str, Any]) -> Dict[str, Any]:
         under the additional `missing_labels` keyword."""
         labels = self._get_raw_label(row)
         additional_kwargs = self._get_additional_kwargs(row)
-        if labels is not None:
-            if any(label is None for label in labels):
-                additional_kwargs["missing_labels"] = [
-                    label is None for label in labels
-                ]
         return dict(
             features=self._get_raw_data(row),
             labels=labels,

diff --git a/tests/unit/collators/testRaggedCollator.py b/tests/unit/collators/testRaggedCollator.py
@@ -73,20 +73,50 @@ def test_call_with_missing_entire_labels(self) -> None:
         data: List[Dict] = [
             {"features": [1, 2], "labels": [True, False], "ident": "sample1"},
             {"features": [3, 4, 5], "labels": None, "ident": "sample2"},
-            {"features": [6], "labels": [True], "ident": "sample3"},
+            {"features": [7], "labels": [True, None], "ident": "sample3"},
+            {"features": [6], "labels": [True], "ident": "sample4"},
+            {"features": [8, 9], "labels": [None, None], "ident": "sample5"},
         ]
 
         result: XYData = self.collator(data)
 
         # https://github.com/ChEB-AI/python-chebai/pull/48#issuecomment-2324393829
-        expected_x = torch.tensor([[1, 2, 0], [3, 4, 5], [6, 0, 0]])
+        expected_x = torch.tensor(
+            [
+                [1, 2, 0],
+                [3, 4, 5],
+                [7, 0, 0],
+                [6, 0, 0],
+                [8, 9, 0],
+            ]
+        )
         expected_y = torch.tensor(
-            [[True, False], [True, False]]
+            [
+                [True, False],
+                [True, False],
+                [True, False],
+                [False, False],
+            ]
         )  # True -> 1, False -> 0
         expected_mask_for_x = torch.tensor(
-            [[True, True, False], [True, True, True], [True, False, False]]
+            [
+                [True, True, False],
+                [True, True, True],
+                [True, False, False],
+                [True, False, False],
+                [True, True, False],
+            ]
+        )
+        expected_lens_for_x = torch.tensor([2, 3, 1, 1, 2])
+        expected_missing_labels = torch.tensor(
+            [
+                [False, False],  # sample1 has no missing labels
+                [False, False],  # sample2 has no missing labels (entire label is None)
+                [False, True],  # sample3 has a missing label at index 1
+                [False, False],  # sample4 has no missing labels
+                [True, True],  # sample5 has missing labels at both indices
+            ]
         )
-        expected_lens_for_x = torch.tensor([2, 3, 1])
 
         self.assertTrue(
             torch.equal(result.x, expected_x),
@@ -110,19 +140,26 @@ def test_call_with_missing_entire_labels(self) -> None:
         )
         self.assertEqual(
             result.additional_fields["loss_kwargs"]["non_null_labels"],
-            [0, 2],
+            [0, 2, 3, 4],
             "The non-null labels list does not match the expected output.",
         )
         self.assertEqual(
             len(result.additional_fields["loss_kwargs"]["non_null_labels"]),
-            result.y.shape[1],
+            result.y.shape[0],
             "The length of non null labels list must match with target label variable size",
         )
         self.assertEqual(
             result.additional_fields["idents"],
-            ("sample1", "sample2", "sample3"),
+            ("sample1", "sample2", "sample3", "sample4", "sample5"),
             "The identifiers do not match the expected output when labels are missing.",
         )
+        self.assertTrue(
+            torch.equal(
+                result.additional_fields["loss_kwargs"]["missing_labels"],
+                expected_missing_labels,
+            ),
+            "The missing labels tensor does not match the expected output when labels are missing.",
+        )
 
     def test_call_with_none_in_labels(self) -> None:
         """
@@ -132,18 +169,32 @@ def test_call_with_none_in_labels(self) -> None:
             {"features": [1, 2], "labels": [None, True], "ident": "sample1"},
             {"features": [3, 4, 5], "labels": [True, False], "ident": "sample2"},
             {"features": [6], "labels": [True], "ident": "sample3"},
+            {"features": [7, 8], "labels": [None, None], "ident": "sample4"},
         ]
 
         result: XYData = self.collator(data)
 
-        expected_x = torch.tensor([[1, 2, 0], [3, 4, 5], [6, 0, 0]])
+        expected_x = torch.tensor([[1, 2, 0], [3, 4, 5], [6, 0, 0], [7, 8, 0]])
         expected_y = torch.tensor(
-            [[False, True], [True, False], [True, False]]
+            [[False, True], [True, False], [True, False], [False, False]]
         )  # None -> False
         expected_mask_for_x = torch.tensor(
-            [[True, True, False], [True, True, True], [True, False, False]]
+            [
+                [True, True, False],
+                [True, True, True],
+                [True, False, False],
+                [True, True, False],
+            ]
+        )
+        expected_lens_for_x = torch.tensor([2, 3, 1, 2])
+        expected_missing_labels = torch.tensor(
+            [
+                [True, False],  # sample1 has a missing label at index 0
+                [False, False],  # sample2 has no missing labels
+                [False, False],  # sample3 has no missing labels
+                [True, True],  # sample4 has missing labels at both indices
+            ]
         )
-        expected_lens_for_x = torch.tensor([2, 3, 1])
 
         self.assertTrue(
             torch.equal(result.x, expected_x),
@@ -167,9 +218,16 @@ def test_call_with_none_in_labels(self) -> None:
         )
         self.assertEqual(
             result.additional_fields["idents"],
-            ("sample1", "sample2", "sample3"),
+            ("sample1", "sample2", "sample3", "sample4"),
             "The identifiers do not match the expected output when labels contain None.",
         )
+        self.assertTrue(
+            torch.equal(
+                result.additional_fields["loss_kwargs"]["missing_labels"],
+                expected_missing_labels,
+            ),
+            "The missing labels tensor does not match the expected output when labels contain None.",
+        )
 
     def test_call_with_empty_data(self) -> None:
         """

diff --git a/tests/unit/mock_data/tox_mock_data.py b/tests/unit/mock_data/tox_mock_data.py
@@ -394,10 +394,6 @@ def data_in_dict_format() -> List[Dict]:
         for dict_ in data_list:
             dict_["features"] = Tox21ChallengeMockData.FEATURE_OF_SMILES
             dict_["group"] = None
-            if any(label is None for label in dict_["labels"]):
-                dict_["missing_labels"] = [
-                    True if label is None else False for label in dict_["labels"]
-                ]
 
         return data_list
 
@@ -509,9 +505,5 @@ def get_setup_processed_output_data() -> List[Dict]:
                     "group": None,
                 }
             )
-            if any(label is None for label in dict_["labels"]):
-                complete_list[-1]["missing_labels"] = [
-                    True if label is None else False for label in dict_["labels"]
-                ]
 
         return complete_list