ChEB-AI · aditya0by0 · Aug 8, 2026 · Aug 10, 2026 · Aug 10, 2026 · Aug 11, 2026
diff --git a/chebai_graph/models/architectures/base.py b/chebai_graph/models/architectures/base.py
@@ -26,7 +26,20 @@ def _get_prediction_and_labels(
         Returns:
             tuple[torch.Tensor, torch.Tensor]: Tuple of (predictions, labels).
         """
-        return torch.sigmoid(output), labels.int()
+        valid_label_mask = data["loss_kwargs"]["valid_label_mask"]
+        predictions = torch.sigmoid(output)
+        labels = labels.int()
+
+        if valid_label_mask is not None:
+            labels[~valid_label_mask] = -1  # Mark invalid labels as -1
+            # https://lightning.ai/docs/torchmetrics/stable/classification/auroc#multilabelauroc
+            # -1 as we torchmetrics ignores -1 labels in multilabel metrics
+            # metric = MultilabelAUROC(
+            #    num_labels=labels.shape[1],
+            #    ignore_index=-1,
+            # )
+
+        return predictions, labels
 
     def _process_labels_in_batch(self, batch: XYData) -> torch.Tensor | None:
         """

diff --git a/chebai_graph/preprocessing/bin/bace/AtomCharge/indices_one_hot.txt b/chebai_graph/preprocessing/bin/bace/AtomCharge/indices_one_hot.txt
@@ -0,0 +1,3 @@
+0
+1
+-1
diff --git a/chebai_graph/preprocessing/bin/bace/AtomFunctionalGroup/indices_one_hot.txt b/chebai_graph/preprocessing/bin/bace/AtomFunctionalGroup/indices_one_hot.txt
@@ -0,0 +1,44 @@
+RING_6
+difluoromethyl
+RING_5
+ether
+alkyne
+alkyl
+primary_amine
+fluoro
+amide
+graph_fg
+NO_FG
+alkene
+tertiary_carbon
+alkene_carbon
+nitrile
+trifluoromethyl
+ketone
+chloro
+quaternary_carbon
+sulfonyl
+hydroxyl
+bromo
+secondary_amine
+ester
+aldoxime
+sulfide
+iodo
+tertiary_amine
+RING_4
+secondary_ketimine
+ketoxime
+RING_3
+carboxylate
+acetal
+sulfinyl
+RING_7
+RING_16
+RING_15
+nitro
+primary_ketimine
+sulfonate_ester
+RING_8
+RING_14
+RING_13
diff --git a/chebai_graph/preprocessing/bin/bace/AtomHybridization/indices_one_hot.txt b/chebai_graph/preprocessing/bin/bace/AtomHybridization/indices_one_hot.txt
@@ -0,0 +1,3 @@
+SP2
+SP3
+SP
diff --git a/chebai_graph/preprocessing/bin/bace/AtomNodeLevel/indices_one_hot.txt b/chebai_graph/preprocessing/bin/bace/AtomNodeLevel/indices_one_hot.txt
@@ -0,0 +1,3 @@
+atom_node_lvl
+fg_node_lvl
+graph_node_level
diff --git a/chebai_graph/preprocessing/bin/bace/AtomNumHs/indices_one_hot.txt b/chebai_graph/preprocessing/bin/bace/AtomNumHs/indices_one_hot.txt
@@ -0,0 +1,4 @@
+0
+1
+3
+2
diff --git a/chebai_graph/preprocessing/bin/bace/AtomType/indices_one_hot.txt b/chebai_graph/preprocessing/bin/bace/AtomType/indices_one_hot.txt
@@ -0,0 +1,8 @@
+6
+7
+8
+9
+17
+16
+35
+53
diff --git a/chebai_graph/preprocessing/bin/bace/BondLevel/indices_one_hot.txt b/chebai_graph/preprocessing/bin/bace/BondLevel/indices_one_hot.txt
@@ -0,0 +1,4 @@
+within_atoms_lvl
+atom_fg_lvl
+within_fg_lvl
+to_graphNode_lvl
diff --git a/chebai_graph/preprocessing/bin/bace/BondType/indices_one_hot.txt b/chebai_graph/preprocessing/bin/bace/BondType/indices_one_hot.txt
@@ -0,0 +1,4 @@
+SINGLE
+AROMATIC
+TRIPLE
+DOUBLE
diff --git a/chebai_graph/preprocessing/bin/bace/NumAtomBonds/indices_one_hot.txt b/chebai_graph/preprocessing/bin/bace/NumAtomBonds/indices_one_hot.txt
@@ -0,0 +1,4 @@
+3
+2
+4
+1
diff --git a/chebai_graph/preprocessing/bin/bbbp/AtomCharge/indices_one_hot.txt b/chebai_graph/preprocessing/bin/bbbp/AtomCharge/indices_one_hot.txt
@@ -0,0 +1,4 @@
+0
+1
+-1
+2
diff --git a/chebai_graph/preprocessing/bin/bbbp/AtomFunctionalGroup/indices_one_hot.txt b/chebai_graph/preprocessing/bin/bbbp/AtomFunctionalGroup/indices_one_hot.txt
@@ -0,0 +1,80 @@
+alkyl
+ester
+chloro
+RING_6
+tertiary_amine
+graph_fg
+tertiary_carbon
+primary_amine
+dichloromethyl
+sulfonyl
+ketone
+secondary_amine
+hydroxyl
+ether
+quaternary_carbon
+amide
+NO_FG
+carboxyl
+nitro
+alkene
+alkene_carbon
+secondary_ketimine
+secondary_aldimine
+carboxylate
+iodo
+aldoxime
+bromo
+alkyne
+fluoro
+carbamate
+trifluoromethyl
+sulfide
+trichloromethyl
+hemiacetal
+primary_ketimine
+azide
+4_ammonium_ion
+sulfonate_ester
+ketoxime
+nitroso
+sulfhydryl
+difluoromethyl
+fluorochloromethyl
+chlorobromomethyl
+phosphono
+carbodithio
+disulfide
+difluorochloromethyl
+carbonate_ester
+phosphate
+RING_5
+thiolester
+RING_7
+nitrile
+phosphoryl
+aldehyde
+sulfinyl
+thioketone
+acetal
+RING_4
+sulfonic_acid
+imide
+RING_14
+RING_24
+amidine
+RING_3
+RING_9
+RING_19
+RING_16
+RING_8
+isonitrile
+RING_23
+RING_12
+RING_21
+RING_17
+RING_30
+RING_33
+RING_36
+azo
+RING_15
diff --git a/chebai_graph/preprocessing/bin/bbbp/AtomHybridization/indices_one_hot.txt b/chebai_graph/preprocessing/bin/bbbp/AtomHybridization/indices_one_hot.txt
@@ -0,0 +1,4 @@
+SP3
+SP2
+S
+SP
diff --git a/chebai_graph/preprocessing/bin/bbbp/AtomNodeLevel/indices_one_hot.txt b/chebai_graph/preprocessing/bin/bbbp/AtomNodeLevel/indices_one_hot.txt
@@ -0,0 +1,3 @@
+atom_node_lvl
+fg_node_lvl
+graph_node_level
diff --git a/chebai_graph/preprocessing/bin/bbbp/AtomNumHs/indices_one_hot.txt b/chebai_graph/preprocessing/bin/bbbp/AtomNumHs/indices_one_hot.txt
@@ -0,0 +1,4 @@
+2
+0
+1
+3
diff --git a/chebai_graph/preprocessing/bin/bbbp/AtomType/indices_one_hot.txt b/chebai_graph/preprocessing/bin/bbbp/AtomType/indices_one_hot.txt
@@ -0,0 +1,13 @@
+6
+17
+8
+7
+16
+11
+53
+35
+9
+1
+15
+5
+20
diff --git a/chebai_graph/preprocessing/bin/bbbp/BondLevel/indices_one_hot.txt b/chebai_graph/preprocessing/bin/bbbp/BondLevel/indices_one_hot.txt
@@ -0,0 +1,4 @@
+within_atoms_lvl
+atom_fg_lvl
+within_fg_lvl
+to_graphNode_lvl
diff --git a/chebai_graph/preprocessing/bin/bbbp/BondType/indices_one_hot.txt b/chebai_graph/preprocessing/bin/bbbp/BondType/indices_one_hot.txt
@@ -0,0 +1,4 @@
+DOUBLE
+SINGLE
+AROMATIC
+TRIPLE
diff --git a/chebai_graph/preprocessing/bin/bbbp/NumAtomBonds/indices_one_hot.txt b/chebai_graph/preprocessing/bin/bbbp/NumAtomBonds/indices_one_hot.txt
@@ -0,0 +1,5 @@
+2
+4
+1
+3
+0
diff --git a/chebai_graph/preprocessing/bin/clin_tox/AtomCharge/indices_one_hot.txt b/chebai_graph/preprocessing/bin/clin_tox/AtomCharge/indices_one_hot.txt
@@ -0,0 +1,6 @@
+0
+-1
+1
+3
+2
+-2
diff --git a/chebai_graph/preprocessing/bin/clin_tox/AtomFunctionalGroup/indices_one_hot.txt b/chebai_graph/preprocessing/bin/clin_tox/AtomFunctionalGroup/indices_one_hot.txt
@@ -0,0 +1,84 @@
+alkyl
+chloro
+RING_6
+ether
+secondary_amine
+fluoro
+graph_fg
+tertiary_carbon
+hydroxyl
+ketone
+sulfonyl
+carboxylate
+RING_5
+NO_FG
+amide
+sulfide
+RING_4
+quaternary_carbon
+alkyne
+alkene
+primary_amine
+carboxyl
+trifluoromethyl
+ester
+alkene_carbon
+RING_20
+sulfhydryl
+nitrile
+RING_3
+RING_29
+RING_7
+4_ammonium_ion
+acetal
+RING_9
+iodo
+RING_38
+sulfinyl
+difluoromethyl
+imide
+RING_16
+tertiary_amine
+nitro
+RING_8
+secondary_ketimine
+nitrate
+sulfonic_acid
+carbamate
+secondary_aldimine
+bromo
+thioketone
+RING_14
+azo
+RING_33
+trichloromethyl
+RING_15
+ketoxime
+phosphodiester
+RING_24
+RING_22
+aldehyde
+RING_18
+RING_21
+nitroso
+dichloromethyl
+carbonate_ester
+silyl_ether
+thiolester
+borono
+aldoxime
+sulfonate_ester
+RING_23
+RING_19
+RING_25
+carbodithio
+disulfide
+RING_36
+RING_31
+fluorochloromethyl
+RING_12
+phosphate
+azide
+chlorobromomethyl
+bromodifluoromethyl
+phosphino
diff --git a/chebai_graph/preprocessing/bin/clin_tox/AtomHybridization/indices_one_hot.txt b/chebai_graph/preprocessing/bin/clin_tox/AtomHybridization/indices_one_hot.txt
@@ -0,0 +1,7 @@
+SP3
+SP2
+SP
+S
+SP3D2
+UNSPECIFIED
+SP3D
diff --git a/chebai_graph/preprocessing/bin/clin_tox/AtomNodeLevel/indices_one_hot.txt b/chebai_graph/preprocessing/bin/clin_tox/AtomNodeLevel/indices_one_hot.txt
@@ -0,0 +1,3 @@
+atom_node_lvl
+fg_node_lvl
+graph_node_level
diff --git a/chebai_graph/preprocessing/bin/clin_tox/AtomNumHs/indices_one_hot.txt b/chebai_graph/preprocessing/bin/clin_tox/AtomNumHs/indices_one_hot.txt
@@ -0,0 +1,4 @@
+3
+0
+2
+1
diff --git a/chebai_graph/preprocessing/bin/clin_tox/AtomType/indices_one_hot.txt b/chebai_graph/preprocessing/bin/clin_tox/AtomType/indices_one_hot.txt
@@ -0,0 +1,30 @@
+6
+17
+7
+8
+9
+16
+15
+53
+80
+1
+13
+35
+79
+0
+27
+22
+78
+33
+34
+30
+14
+5
+29
+24
+81
+43
+83
+25
+20
+26
diff --git a/chebai_graph/preprocessing/bin/clin_tox/BondLevel/indices_one_hot.txt b/chebai_graph/preprocessing/bin/clin_tox/BondLevel/indices_one_hot.txt
@@ -0,0 +1,4 @@
+within_atoms_lvl
+atom_fg_lvl
+within_fg_lvl
+to_graphNode_lvl
diff --git a/chebai_graph/preprocessing/bin/clin_tox/BondType/indices_one_hot.txt b/chebai_graph/preprocessing/bin/clin_tox/BondType/indices_one_hot.txt
@@ -0,0 +1,5 @@
+SINGLE
+AROMATIC
+DOUBLE
+TRIPLE
+DATIVE
diff --git a/chebai_graph/preprocessing/bin/clin_tox/NumAtomBonds/indices_one_hot.txt b/chebai_graph/preprocessing/bin/clin_tox/NumAtomBonds/indices_one_hot.txt
@@ -0,0 +1,6 @@
+1
+3
+2
+4
+0
+6
-Original file line number
+Diff line change
@@ -0,0 +1,3 @@
+    SP2
+    SP3
+    SP
Original file line number	Diff line number	Diff line change
		@@ -0,0 +1,30 @@
		6
		17
		7
		8
		9
		16
		15
		53
		80
		1
		13
		35
		79
		0
		27
		22
		78
		33
		34
		30
		14
		5
		29
		24
		81
		43
		83
		25
		20
		26