fix multilabel datareader bug (#205)

deepmodeling · Jan 10, 2024 · 0575fb0 · 0575fb0
1 parent 6066381
commit 0575fb0
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/unimol_tools/unimol_tools/data/datareader.py b/unimol_tools/unimol_tools/data/datareader.py
@@ -65,9 +65,10 @@ def read_data(self, data=None, is_train=True, **params):
                 else:
                     for i in range(label.shape[1]):
                         data[target_col_prefix + str(i)] = label[:,i]
+
+            _ = data.pop('target')
             data = pd.DataFrame(data).rename(columns={smiles_col: 'SMILES'})
-            if 'target' in data:
-                data = data.drop(columns=['target'])
+
         elif isinstance(data, list):
             # load from smiles list
             data = pd.DataFrame(data, columns=['SMILES'])