Update pipe.py

pipe.py

@@ -80,52 +80,49 @@ class SmilesDiffusionPipe(Pipeline):
 
     def unmask_partial_smiles(self, input_ids, model, voc, steps, k):
         """
-        Iteratively unmask a SMILES sequence using a trained diffusion model.
-        Always fills the masked token with highest confidence first, until all masks are gone.
+        Iteratively unmask a single SMILES string using a diffusion model.
+        Always fills the masked token with highest confidence first until done.
         """
-        model.eval()
-        with torch.no_grad():
-            sequences = torch.tensor(
-                input_ids, dtype=torch.long, device=self.device
-            ).unsqueeze(0)
-    
-            mask_token = voc.vocab["[MASK]"]
-            pad_token = voc.vocab.get("[PAD]", None)
+        pad_token = voc.vocab["[PAD]"]
+        mask_token = voc.vocab["[MASK]"]
     
-            while True:
-                mask_positions = sequences == mask_token
-                num_masked = mask_positions.sum().item()
-                if num_masked == 0:
-                    break  # stop only when all masks are filled
+        # Encode input SMILES
+        encoded = voc.encode(voc.tokenize(input_ids))
+        unmasked = torch.tensor(encoded, dtype=torch.long).unsqueeze(0).to(device)
     
-                # ---- estimate diffusion time t ----
+        with torch.no_grad():
+            while (unmasked == mask_token).any():
+                # Estimate diffusion time t
                 if pad_token is not None:
-                    valid = sequences != pad_token
-                    frac_masked = mask_positions.sum().float() / valid.sum()
+                    valid_mask = unmasked != pad_token
+                    frac_masked = ((unmasked == mask_token).sum().float() / valid_mask.sum())
                 else:
-                    frac_masked = mask_positions.sum().float() / sequences.size(1)
-                t = torch.tensor([frac_masked], device=self.device)
+                    frac_masked = ((unmasked == mask_token).sum().float() / unmasked.size(1))
+                t_val = torch.tensor([frac_masked], device=device)
     
-                # ---- forward pass ----
-                logits = model(sequences, t=t)
+                # Forward pass
+                logits = model.model(unmasked, t=t_val)
                 probs = F.softmax(logits, dim=-1)
     
-                mask_indices = mask_positions.nonzero(as_tuple=False)
+                mask_positions = (unmasked == mask_token)
                 masked_probs = probs[mask_positions]
     
-                # ---- pick the mask with highest confidence ----
+                if masked_probs.size(0) == 0:
+                    break  # safety check
+    
+                # Pick the mask with highest confidence
                 masked_confidence = masked_probs.max(dim=-1).values
                 best_idx = torch.argmax(masked_confidence)
-    
-                sampled_id = torch.multinomial(masked_probs[best_idx], num_samples=1).item()
+                mask_indices = mask_positions.nonzero(as_tuple=False)
                 pos = mask_indices[best_idx]
     
-                # ---- fill that mask ----
-                sequences[pos[0], pos[1]] = sampled_id
+                # Unmask the selected token
+                sampled_id = torch.multinomial(masked_probs[best_idx], num_samples=1).item()
+                unmasked[0, pos[1]] = sampled_id
     
             # ---- decode ----
             decoded = [
-                t for t in sequences[0].tolist()
+                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: