Update prediction.py

prediction.py

@@ -40,54 +40,66 @@ def compute_descriptors(smiles: str):
     ]
     return np.array(descriptors, dtype=np.float32)
 
-# Transformer regression model definition
+# Transformer regression model definition (must match training)
 class TransformerRegressor(nn.Module):
-    def __init__(self, input_dim, hidden_dim, num_layers, output_dim):
+    def __init__(self, input_dim, embedding_dim, ff_dim, num_layers, output_dim):
         super().__init__()
-        self.feat_proj = nn.Linear(input_dim, hidden_dim)
-        encoder_layer = nn.TransformerEncoderLayer(d_model=hidden_dim, nhead=8)
+        self.feat_proj = nn.Linear(input_dim, embedding_dim)
+        encoder_layer = nn.TransformerEncoderLayer(
+            d_model=embedding_dim,
+            nhead=8,
+            dim_feedforward=ff_dim,
+            dropout=0.1,
+            batch_first=True
+        )
         self.transformer_encoder = nn.TransformerEncoder(encoder_layer, num_layers=num_layers)
         self.regression_head = nn.Sequential(
-            nn.Linear(hidden_dim, 128),
+            nn.Linear(embedding_dim, 256),
             nn.ReLU(),
-            nn.Linear(128, 64),
+            nn.Linear(256, 128),
             nn.ReLU(),
-            nn.Linear(64, output_dim)
+            nn.Linear(128, output_dim)
         )
 
     def forward(self, x):
         x = self.feat_proj(x)
         x = self.transformer_encoder(x)
-        x = x.mean(dim=1)  # Global average pooling
+        x = x.mean(dim=1)
         return self.regression_head(x)
 
 # Model hyperparameters (must match training)
-input_dim = 768  # Output size of ChemBERTa model
-hidden_dim = 256
+embedding_dim = 768
+descriptor_dim = 1290  # Based on earlier errors. If unsure, use 1290
+input_dim = embedding_dim + descriptor_dim  # 768 + 1290 = 2058
+ff_dim = 1024
 num_layers = 2
-output_dim = 6  # Number of properties
+output_dim = 6
 
 # Load trained model
 device = torch.device("cpu")
-model = TransformerRegressor(input_dim, hidden_dim, num_layers, output_dim)
+model = TransformerRegressor(input_dim, embedding_dim, ff_dim, num_layers, output_dim)
 model.load_state_dict(torch.load("transformer_model.pt", map_location=device))
 model.eval()
 
 # Prediction function
 def predict_properties(smiles: str):
     try:
-        # Validate SMILES and compute descriptors
-        _ = compute_descriptors(smiles)
+        # Compute descriptors
+        descriptors = compute_descriptors(smiles)
+        descriptors_tensor = torch.tensor(descriptors, dtype=torch.float32).unsqueeze(0)
 
-        # ChemBERTa embedding (CLS token)
+        # Get ChemBERTa embedding (CLS token)
         inputs = tokenizer(smiles, return_tensors="pt")
         with torch.no_grad():
             outputs = embedding_model(**inputs)
-        embedding = outputs.last_hidden_state[:, 0, :]  # Extracting the [CLS] token (shape: (1, 768))
+        embedding = outputs.last_hidden_state[:, 0, :]  # (1, 768)
+
+        # Combine features
+        combined = torch.cat([embedding, descriptors_tensor], dim=1).unsqueeze(1)  # Shape: (1, 1, 2058)
 
-        # Forward pass through model
+        # Forward pass
         with torch.no_grad():
-            preds = model(embedding)
+            preds = model(combined)
 
         preds_np = preds.numpy()
 
@@ -98,7 +110,6 @@ def predict_properties(smiles: str):
             for i in range(output_dim)
         ], axis=1)
 
-        # Create dictionary of results
         results = {key: round(val, 4) for key, val in zip(keys, preds_rescaled.flatten())}
         return results