Update pipe.py

pipe.py

@@ -78,54 +78,67 @@ class SmilesDiffusionPipe(Pipeline):
         return torch.multinomial(probs, 1)
 
 
-    def unmask_partial_smiles(self, input_ids, model, voc, steps, k):
-        """
-        Iteratively unmask a single SMILES string using a diffusion model.
-        Always fills the masked token with highest confidence first until done.
-        """
-        pad_token = voc.vocab["[PAD]"]
-        mask_token = voc.vocab["[MASK]"]
-    
-        # Encode input SMILES
-        encoded = voc.encode(voc.tokenize(input_ids))
-        unmasked = torch.tensor(encoded, dtype=torch.long).unsqueeze(0).to(device)
-    
+    def unmask_partial_smiles(self, input_ids, model, voc, steps=30, k=3):
+        model.eval()
         with torch.no_grad():
-            while (unmasked == mask_token).any():
-                # Estimate diffusion time t
-                if pad_token is not None:
-                    valid_mask = unmasked != pad_token
-                    frac_masked = ((unmasked == mask_token).sum().float() / valid_mask.sum())
-                else:
-                    frac_masked = ((unmasked == mask_token).sum().float() / unmasked.size(1))
-                t_val = torch.tensor([frac_masked], device=device)
-    
-                # Forward pass
-                logits = model.model(unmasked, t=t_val)
-                probs = F.softmax(logits, dim=-1)
+            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)
+    
+            total_tokens = sequences.size(1)
+            original_num_masked = (sequences == mask_token).sum().item()
+            print(f"Original number of masked tokens: {original_num_masked}")
     
-                mask_positions = (unmasked == mask_token)
-                masked_probs = probs[mask_positions]
+            if total_tokens == 0:
+                raise ValueError("Encoded SMILES has no tokens.")
+            if original_num_masked == 0:
+                print("No masked tokens found. Returning original SMILES.")
+                input_ids = input_ids[1:-1]
+                return self.decode(input_ids)
     
-                if masked_probs.size(0) == 0:
-                    break  # safety check
+            filled_tokens_confidence = {}
+    
+            # --- ITERATION LOGIC: unmask one token at a time based on highest confidence ---
+            num_masks = (sequences == mask_token).sum().item()
+            for _ in range(num_masks):
+                # Estimate t per sequence
+                frac_masked = (sequences == mask_token).sum().item() / total_tokens
+                t = torch.tensor([frac_masked], device=self.device)
+    
+                logits = model(sequences, t=t)
+                probs = F.softmax(logits, dim=-1)
     
-                # Pick the mask with highest confidence
-                masked_confidence = masked_probs.max(dim=-1).values
-                best_idx = torch.argmax(masked_confidence)
+                mask_positions = (sequences == mask_token)
                 mask_indices = mask_positions.nonzero(as_tuple=False)
-                pos = mask_indices[best_idx]
     
-                # Unmask the selected token
-                sampled_id = torch.multinomial(masked_probs[best_idx], num_samples=1).item()
-                unmasked[0, pos[1]] = sampled_id
+                if len(mask_indices) == 0:
+                    print("All tokens filled.")
+                    break
+    
+                masked_probs = probs[mask_positions]  # (#masked, vocab)
+                sampled_ids = torch.multinomial(masked_probs, num_samples=1).squeeze(-1)
+                masked_confidence = torch.max(masked_probs, dim=-1).values
+    
+                best_idx = torch.argmax(masked_confidence).item()
+                chosen_pos = mask_indices[best_idx]       # [batch=0, position]
+                chosen_token = sampled_ids[best_idx].item()
+                chosen_conf = masked_confidence[best_idx].item()
+    
+                sequences[0, chosen_pos[1]] = chosen_token
+                filled_tokens_confidence[chosen_pos[1].item()] = chosen_conf
+    
+                print(f"Filled index {chosen_pos[1].item()} with token {chosen_token} (conf={chosen_conf:.4f})")
+    
+            # Print confidence for all filled tokens
+            for idx, confidence in sorted(filled_tokens_confidence.items(), key=lambda x: x[0]):
+                print(f"Final confidence for token at index {idx}: {confidence:.3f}")
+    
+            # Decode, removing <s> and </s>
+            decoded_tokens = [t for t in sequences[0].tolist() if t not in (voc.vocab.get('<s>'), voc.vocab.get('</s>'))]
     
-            # ---- 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]
+            if len(decoded_tokens) > 2:
+                decoded_tokens = decoded_tokens[1:-1]
     
-            return self.decode(decoded).replace(" ", "")
+            decoded = self.decode(decoded_tokens).replace(" ", "")
+            return decoded