Update pipe.py

pipe.py

@@ -84,57 +84,51 @@ class SmilesDiffusionPipe(Pipeline):
         Iteratively unmasks a single SMILES string using a trained diffusion model.
         High-confidence tokens are unmasked first, sampled multinomially.
         """
-        steps = 30
         model.eval()
+        pad_token = voc.vocab["[PAD]"]
+        mask_token = voc.vocab["[MASK]"]
+    
+        # encode input_ids (assuming input_ids is a SMILES string)
+        encoded = self.encode(self.tokenize_smiles(input_ids))
+        encoded_tensor = torch.tensor(encoded, dtype=torch.long).unsqueeze(0).to(self.device)
+    
         with torch.no_grad():
-            unmasked = torch.tensor(input_ids, dtype=torch.long, device=self.device).unsqueeze(0)
-            
-            mask_token = voc.vocab["[MASK]"]
-            pad_token = voc.vocab.get("[PAD]", None)
-    
-            num_masked = (unmasked == mask_token).sum().item()
-            if num_masked == 0:
-                return self.decode(input_ids[1:-1])
-    
-            steps = min(steps, num_masked)
-    
-            for _ in range(steps):
-                # Estimate diffusion time t
-                if pad_token is not None:
-                    valid = unmasked != pad_token
-                    frac_masked = (unmasked == mask_token).sum() / valid.sum()
-                else:
-                    frac_masked = (unmasked == mask_token).sum() / unmasked.size(1)
-                t = frac_masked.unsqueeze(0).to(self.device)
-    
-                # Forward pass
-                logits = model(unmasked, t=t)
-                probs = F.softmax(logits, dim=-1)
+            unmasked = encoded_tensor.clone()
+            num_masks = (unmasked == mask_token).sum().item()
     
-                # Mask positions
-                mask_positions = unmasked == mask_token
-                mask_indices = mask_positions.nonzero(as_tuple=False)
-                if mask_indices.numel() == 0:
-                    break
+            for _ in range(num_masks):
+                # compute diffusion time t
+                t_val = estimate_t_per_sequence(unmasked, mask_token, pad_token).to(self.device)
     
+                # forward pass
+                logits = model(unmasked, t=t_val)
+                probs = F.softmax(logits, dim=-1)
+    
+                # masked positions
+                mask_positions = (unmasked == mask_token)
                 masked_probs = probs[mask_positions]
     
-                # Multinomial sample for each masked token
+                if masked_probs.numel() == 0:
+                    break
+    
+                # multinomial sample per masked token
                 sampled_ids = torch.multinomial(masked_probs, num_samples=1).squeeze(-1)
                 masked_confidence = masked_probs.max(dim=-1).values
     
-                # Select highest-confidence mask to fill
+                # pick the highest-confidence mask
                 best_idx = torch.argmax(masked_confidence)
+                mask_indices = mask_positions.nonzero(as_tuple=False)
                 pos = mask_indices[best_idx]
     
+                # unmask
                 unmasked[0, pos[1]] = sampled_ids[best_idx]
     
-            # Decode
-            decoded = [
-                t for t in unmasked[0].tolist()
-                if t not in (pad_token, voc.vocab.get("<s>"), voc.vocab.get("</s>"))
-            ]
-            if len(decoded) > 2:
-                decoded = decoded[1:-1]
+            # decode and remove padding / special tokens
+            decoded_tokens = [t for t in unmasked[0].tolist() 
+                              if t not in (pad_token, voc.vocab.get("<s>"), voc.vocab.get("</s>"))]
+            if len(decoded_tokens) > 2:
+                decoded_tokens = decoded_tokens[1:-1]
+    
+            predicted_smi = self.decode(decoded_tokens).replace(" ", "")
     
-            return self.decode(decoded).replace(" ", "")
+        return predicted_smi