debug predict.py

predict.py

@@ -24,12 +24,14 @@ def predict(smiles_list: list[str]) -> dict:
     Returns:
         dict: nested prediction dictionary, following {'<smiles>': {'<target>': <pred>}}
     """
+    print(f"Received {len(smiles_list)} SMILES strings")
     # preprocessing pipeline
     features, removed_idxs = preprocess_molecules(
         smiles_list,
         load_ecdf_path="assets/ecdfs.pkl",
         load_scaler_path="assets/scaler.pkl",
     )
+    print(f"{len(removed_idxs)} molecules removed during cleaning")
 
     # setup model
     model = Tox21RFClassifier(seed=42)
@@ -37,13 +39,16 @@ def predict(smiles_list: list[str]) -> dict:
 
     # make predicitons
     predictions = defaultdict(dict)
+    # make smiles list with same num_samples as features
+    clean_smiles = [smi for i, smi in enumerate(smiles_list) if i not in removed_idxs]
+    no_pred_smiles = [smi for i, smi in enumerate(smiles_list) if i in removed_idxs]
 
-    for i, smiles in enumerate(smiles_list):
-        for target in model.tasks:
-            predictions[smiles][target] = (
-                0.0
-                if i in removed_idxs
-                else model.predict(target, features[i].reshape(1, -1))
-            )
+    for target in model.tasks:
+        target_pred = model.predict(target, features)
+        for i, smiles in enumerate(clean_smiles):
+            predictions[smiles][target] = target_pred[i]
+
+        for smiles in no_pred_smiles:
+            predictions[smiles][target] = 0.0
 
     return predictions