Only noise until max_it-1

app.py

@@ -135,33 +135,34 @@ def diffusion_chat(question, max_it, pause_length, sharpness,
             break
 
         # NOISING
-        threshold = get_noising_schedule(i, max_it, sharpness=sharpness)
-        if use_confidence_noising:
-            noised_answer, just_noised_indices = confidence_guided_noising(
-                current_tokens, answer_start, tokenizer, confidences, noise_clipping,
-                threshold=threshold, noise_start=noise_start
-            )
-        elif use_permanent_unmasking:
-            noised_answer, just_noised_indices = noisify_answer_without_remasking(
-                current_tokens, answer_start, tokenizer, threshold=threshold,
-                noise_start=noise_start, unmasked_mask=unmasked_mask
-            )
-        else:
-            noised_answer, just_noised_indices = noisify_answer(
-                current_tokens, answer_start, tokenizer,
-                threshold=threshold, clustering=clustering, noise_start=noise_start
-            )
-        
-        for idx in range(answer_start, len(current_tokens)):
-            if noised_answer[idx] != mask_token_id:
-                unmasked_mask[idx] = True
-
-
-
-        yield render_html(f"Iteration {i+1}/{max_it} (before noising)",
-                          highlight_tokens(current_tokens[answer_start:], answer_start, just_noised_indices, color="red"))
-
-        current_tokens = ori_input_tokens[:answer_start] + noised_answer[answer_start:]
+        if i < max_it-1:
+            threshold = get_noising_schedule(i, max_it, sharpness=sharpness)
+            if use_confidence_noising:
+                noised_answer, just_noised_indices = confidence_guided_noising(
+                    current_tokens, answer_start, tokenizer, confidences, noise_clipping,
+                    threshold=threshold, noise_start=noise_start
+                )
+            elif use_permanent_unmasking:
+                noised_answer, just_noised_indices = noisify_answer_without_remasking(
+                    current_tokens, answer_start, tokenizer, threshold=threshold,
+                    noise_start=noise_start, unmasked_mask=unmasked_mask
+                )
+            else:
+                noised_answer, just_noised_indices = noisify_answer(
+                    current_tokens, answer_start, tokenizer,
+                    threshold=threshold, clustering=clustering, noise_start=noise_start
+                )
+            
+            for idx in range(answer_start, len(current_tokens)):
+                if noised_answer[idx] != mask_token_id:
+                    unmasked_mask[idx] = True
+    
+    
+    
+            yield render_html(f"Iteration {i+1}/{max_it} (before noising)",
+                              highlight_tokens(current_tokens[answer_start:], answer_start, just_noised_indices, color="red"))
+    
+            current_tokens = ori_input_tokens[:answer_start] + noised_answer[answer_start:]
 
     # Final output
     answer_ids = current_tokens[answer_start:]