Update pipe.py

pipe.py

@@ -76,58 +76,73 @@ class SmilesDiffusionPipe(Pipeline):
         model.eval()
         with torch.no_grad():
             sequences = torch.tensor(input_ids, dtype=torch.long, device=self.device).unsqueeze(0)
-
+    
             total_tokens = sequences.size(1)
             original_num_masked = (sequences == voc.vocab['[MASK]']).sum().item()
             print(f"Original number of masked tokens: {original_num_masked}")
-
+    
             if total_tokens == 0:
                 raise ValueError("Encoded SMILES has no tokens.")
             if original_num_masked == 0:
                 print("No masked tokens found. Returning original SMILES.")
                 return self.decode(input_ids)
-
+    
             filled_tokens_confidence = {}
-
+    
+            # -------------------------------
+            # Iterative filling: one token per step
+            # -------------------------------
             for step in range(steps):
-                t = torch.tensor([(sequences == voc.vocab['[MASK]']).sum().item() / total_tokens], device=self.device)
+    
+                # Time step based on % masked
+                t = torch.tensor([(sequences == voc.vocab['[MASK]']).sum().item() / total_tokens],
+                                 device=self.device)
+    
                 logits = model(sequences, t=t)
                 probs = F.softmax(logits, dim=-1)
-
+    
                 mask_indices = (sequences == voc.vocab['[MASK]']).nonzero(as_tuple=True)[1]
+    
                 if len(mask_indices) == 0:
                     print("All tokens filled.")
                     break
-
-                confidences = []
-                filled_indices = []
-
+    
+                # For each masked position, get highest probability among all tokens
+                candidate_confs = []
+                candidate_samples = []
+    
                 for idx in mask_indices:
-                    token_probs = probs[0, idx]
-                    sampled_token = self.multinomial_sample(token_probs)
-                    confidence = token_probs[sampled_token].item()
-                    sequences[0, idx] = sampled_token
-                    confidences.append(confidence)
-                    filled_indices.append(idx)
-                    filled_tokens_confidence[idx] = confidence
-
-                avg_conf = sum(confidences) / len(confidences)
-
-                for idx, conf in zip(filled_indices, confidences):
-                    if conf < avg_conf:
-                        sequences[0, idx] = voc.vocab['[MASK]']
-                        filled_tokens_confidence[idx] = None
-
+                    p = probs[0, idx]
+    
+                    # sample using multinomial
+                    sampled_token = self.multinomial_sample(p)
+                    confidence = p[sampled_token].item()
+    
+                    candidate_confs.append(confidence)
+                    candidate_samples.append((idx.item(), sampled_token, confidence))
+    
+                # Pick the token with highest confidence
+                chosen_idx, chosen_token, chosen_conf = max(candidate_samples, key=lambda x: x[2])
+    
+                # Fill it permanently
+                sequences[0, chosen_idx] = chosen_token
+                filled_tokens_confidence[chosen_idx] = chosen_conf
+    
+                print(f"[Step {step+1}] Filled index {chosen_idx} with token {chosen_token} (conf={chosen_conf:.4f})")
+    
+            # Print final per-token confidence summary
             sorted_filled_tokens = sorted(
-                ((idx, conf) for idx, conf in filled_tokens_confidence.items() if conf is not None),
+                ((idx, conf) for idx, conf in filled_tokens_confidence.items()),
                 key=lambda x: x[0]
             )
-
+    
             for idx, confidence in sorted_filled_tokens:
                 print(f"Final confidence for token at index {idx}: {confidence:.3f}")
-
+    
+            # Decode without <s> or </s>
             decoded = self.decode([
-                token for token in sequences[0].tolist()
-                if token not in (voc.vocab.get('<s>'), voc.vocab.get('</s>'))
+                t for t in sequences[0].tolist()
+                if t not in (voc.vocab.get('<s>'), voc.vocab.get('</s>'))
             ]).replace(" ", "")
+    
             return decoded