Update pipe.py

pipe.py

@@ -81,8 +81,6 @@ class SmilesDiffusionPipe(Pipeline):
     
             total_tokens = sequences.size(1)
     
-            # ---------------------------------------------------------
-            # KEEP THIS EXACTLY
             original_num_masked = (sequences == mask_token).sum().item()
             print(f"Original number of masked tokens: {original_num_masked}")
     
@@ -99,9 +97,6 @@ class SmilesDiffusionPipe(Pipeline):
     
             filled_tokens_confidence = {}
     
-            # -------------------------------
-            # NEW DIFFUSION LOGIC
-            # -------------------------------
             for _ in range(steps):
                 # Diffusion timestep (constant 1)
                 t_val = torch.ones(1, device=self.device)
@@ -130,9 +125,7 @@ class SmilesDiffusionPipe(Pipeline):
                     filled_tokens_confidence[pos[1].item()] = conf
                     confidences[i] = conf
     
-            # ---------------------------------------
-            # SORT & PRINT (KEEP EXACTLY)
-            # ---------------------------------------
+
             sorted_filled_tokens = sorted(
                 ((idx, conf) for idx, conf in filled_tokens_confidence.items() if conf is not None),
                 key=lambda x: x[0]
@@ -140,10 +133,9 @@ class SmilesDiffusionPipe(Pipeline):
     
             for idx, confidence in sorted_filled_tokens:
                 print(f"Final confidence for token at index {idx}: {confidence:.3f}")
-    
-            # ---------------------------------------
+
             # Decode
-            # ---------------------------------------
+
             decoded = self.decode([
                 token for token in sequences[0].tolist()
                 if token not in (voc.vocab.get('<s>'), voc.vocab.get('</s>'))