Show red highlighting before actual noising

app.py

@@ -202,6 +202,8 @@ def diffusion_chat(question, eot_weight, max_it, sharpness, noise_clipping, use_
             yield f"<b>Stopped early after {i+1} iterations.</b>"
             break
 
+        previous_tokens = current_tokens.copy()
+
         # --- NOISING STEP ---
         threshold = get_noising_schedule(i, max_it, sharpness=sharpness)
         if use_confidence_noising:
@@ -218,7 +220,7 @@ def diffusion_chat(question, eot_weight, max_it, sharpness, noise_clipping, use_
         current_tokens = ori_input_tokens[:answer_start] + noised_answer[answer_start:]
 
         # --- RED HIGHLIGHT ---
-        decoded_tokens = tokenizer.convert_ids_to_tokens(current_tokens[answer_start:])
+        decoded_tokens = tokenizer.convert_ids_to_tokens(previous_tokens[answer_start:])
         highlighted = []
         for j, tok in enumerate(decoded_tokens):
             tok_id = tokenizer.convert_tokens_to_ids(tok)
@@ -231,7 +233,7 @@ def diffusion_chat(question, eot_weight, max_it, sharpness, noise_clipping, use_
             else:
                 highlighted.append(token_str)
 
-        yield f"<b>Iteration {i+1}/{max_it} (after noising):</b><br>" + "".join(highlighted).replace('\n', '<br>')
+        yield f"<b>Iteration {i+1}/{max_it} (before noising):</b><br>" + "".join(highlighted).replace('\n', '<br>')
         time.sleep(0.1)