Change highlighting

app.py

@@ -265,15 +265,12 @@ def diffusion_chat(question, max_it, pause_length, sharpness, clustering, noise_
             noised_answer, just_noised_indices = confidence_guided_noising(
                 current_tokens, answer_start, confidences, noise_clipping, threshold=threshold, noise_start=noise_start
             )
-            just_noised_indices = []
+            # just_noised_indices = []
         else:
             noised_answer, just_noised_indices = noisify_answer(
                 current_tokens, answer_start, threshold=threshold, clustering=clustering, noise_start = noise_start,
             )
 
-        # Compose full input again: prompt + noised answer
-        current_tokens = ori_input_tokens[:answer_start] + noised_answer[answer_start:]
-
         # --- RED HIGHLIGHT ---
         decoded_tokens = tokenizer.convert_ids_to_tokens(current_tokens[answer_start:])
         highlighted = []
@@ -288,8 +285,11 @@ def diffusion_chat(question, max_it, pause_length, sharpness, clustering, noise_
             else:
                 highlighted.append(token_str)
 
+        # Compose full input again: prompt + noised answer
+        current_tokens = ori_input_tokens[:answer_start] + noised_answer[answer_start:]
+        
         yield f"<b>Iteration {i+1}/{max_it} (before noising):</b><br>" + "".join(highlighted).replace('\n', '<br>')
-        time.sleep(pause_length)
+        # time.sleep(pause_length)
 
 
     answer_ids = current_tokens[answer_start:]