From bac2cd3c345def3599a408b25a77b43fc7e43a5f Mon Sep 17 00:00:00 2001 From: aditya0b0 Date: Sat, 15 Aug 2026 21:57:48 +0200 Subject: [PATCH] move missing label mask to collator func and handle cases where all labels are None [None, None] --- chebai/preprocessing/collate.py | 46 +++++++----- chebai/preprocessing/reader.py | 5 -- tests/unit/collators/testRaggedCollator.py | 84 ++++++++++++++++++---- tests/unit/mock_data/tox_mock_data.py | 8 --- 4 files changed, 101 insertions(+), 42 deletions(-) diff --git a/chebai/preprocessing/collate.py b/chebai/preprocessing/collate.py index 06b4ff34..93d146a6 100644 --- 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,30 +84,41 @@ 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: @@ -112,7 +126,7 @@ def __call__(self, data: List[Union[Dict, Tuple]]) -> XYData: 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 index 664a8d8f..19e492e6 100644 --- 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 index d9ab2b1d..3c68825b 100644 --- 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 index fcf5633f..7567d6b2 100644 --- 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