fixed bug

.gitignore

@@ -3,4 +3,5 @@ results.csv
 predict copy.py
 debug.py
 __pycache__
-tox21_test.csv
+tox21_test.csv
+predictions.json

predict.py

@@ -3,6 +3,7 @@ import csv
 import subprocess
 import pandas as pd
 import logging
+import json
 
 from src.preprocess import create_clean_smiles
 
@@ -22,9 +23,17 @@ def predict(smiles_list):
         clean_smiles, valid_mask = create_clean_smiles(smiles_list)
 
         # Mapping from cleaned to original for valid ones
-        cleaned_to_original = {
-            clean: orig for clean, orig, valid in zip(clean_smiles, smiles_list, valid_mask) if valid
-        }
+        originals_valid = [orig for orig, ok in zip(smiles_list, valid_mask) if ok]
+
+        # sanity check (optional but nice to have)
+        if len(originals_valid) != len(clean_smiles):
+            raise ValueError(
+                f"Mismatch: {len(originals_valid)} valid originals vs {len(clean_smiles)} cleaned SMILES"
+            )
+
+        # map cleaned → original
+        cleaned_to_original = dict(zip(clean_smiles, originals_valid))
+        print(len(cleaned_to_original.keys()))
 
         # tox21 targets
         TARGET_NAMES = [
@@ -59,17 +68,10 @@ def predict(smiles_list):
         predictions = {}
         with open("./src/preds.csv", "r", newline="") as f:
             reader = csv.DictReader(f)
+            rows = list(reader)
             target_names = [col for col in reader.fieldnames if col != "smiles"]
 
-            missing = [t for t in TARGET_NAMES if t not in target_names]
-            extra = [t for t in target_names if t not in TARGET_NAMES]
-
-            if missing:
-                logging.error(f"❌ Missing target columns in preds.csv: {missing}")
-            if extra:
-                logging.warning(f"⚠ Warning: Extra columns in preds.csv not expected: {extra}")
-
-            for row in reader:
+            for row in rows:
                 clean_smi = row["smiles"]
                 original_smi = cleaned_to_original.get(clean_smi, clean_smi)
                 pred_dict = {t: float(row[t]) for t in target_names}