Update prediction.py

prediction.py

@@ -6,6 +6,7 @@ import joblib
 from transformers import AutoTokenizer, AutoModel
 from rdkit import Chem
 from rdkit.Chem import Descriptors
+from rdkit.Chem import AllChem
 from datetime import datetime
 from db import get_database  # This must be available in your repo
 
@@ -86,12 +87,20 @@ def compute_descriptors(smiles: str):
     ]
     return np.array(descriptors, dtype=np.float32)
 
+# Fingerprint computation
+def get_morgan_fingerprint(smiles, radius=2, n_bits=1280):
+    mol = Chem.MolFromSmiles(smiles)
+    if mol is None:
+        raise ValueError("Invalid SMILES string.")
+    fp = AllChem.GetMorganFingerprintAsBitVect(mol, radius, nBits=n_bits)
+    return np.array(fp, dtype=np.float32).reshape(1, -1)  # (1, 1280)
+
 # Embedding function
 def get_chemberta_embedding(smiles: str):
     inputs = tokenizer(smiles, return_tensors="pt")
     with torch.no_grad():
         outputs = chemberta(**inputs)
-    return outputs.last_hidden_state[:, 0, :]  # CLS token
+    return outputs.last_hidden_state[:, 0, :]  # CLS token (1, 768)
 
 # Save prediction to MongoDB
 def save_to_db(smiles, predictions):
@@ -119,11 +128,15 @@ def show():
                 return
 
             descriptors = compute_descriptors(smiles_input)
-            descriptors_tensor = torch.tensor(descriptors, dtype=torch.float32).unsqueeze(0)
+            descriptors_tensor = torch.tensor(descriptors, dtype=torch.float32).unsqueeze(0)  # (1, 10)
+
+            fingerprint = get_morgan_fingerprint(smiles_input)  # (1, 1280)
+            fingerprint_tensor = torch.tensor(fingerprint, dtype=torch.float32)  # (1, 1280)
 
-            embedding = get_chemberta_embedding(smiles_input)
+            embedding = get_chemberta_embedding(smiles_input)  # (1, 768)
 
-            combined = torch.cat([embedding, descriptors_tensor], dim=1).unsqueeze(1)  # (1, 1, 2058)
+            combined_input = torch.cat([embedding, descriptors_tensor, fingerprint_tensor], dim=1)  # (1, 2058)
+            combined = combined_input.unsqueeze(1)  # (1, 1, 2058)
 
             with torch.no_grad():
                 preds = model(combined)