Fix

app.py

@@ -175,7 +175,7 @@ def diffusion_chat(question, eot_weight, max_it, sharpness, noise_clipping, use_
         generated_tokens, confidences = generate_diffusion_text(current_tokens)
 
         # Save full output for noising step
-        current_tokens = ori_input_tokens[answer_start] + generated_tokens[answer_start:]
+        current_tokens = ori_input_tokens[:answer_start] + generated_tokens[answer_start:]
 
         # --- GREEN HIGHLIGHT ---
         decoded_tokens = tokenizer.convert_ids_to_tokens(current_tokens[answer_start:])
@@ -194,6 +194,14 @@ def diffusion_chat(question, eot_weight, max_it, sharpness, noise_clipping, use_
         yield f"<b>Iteration {i+1}/{max_it} (after generation):</b><br>" + "".join(highlighted).replace('\n', '<br>')
         time.sleep(0.1)
 
+        # --- Early stopping ---
+        last_tokens.append(current_tokens)
+        if len(last_tokens) > 3:
+            last_tokens.pop(0)
+        if len(last_tokens) == 3 and last_tokens[0] == last_tokens[1] == last_tokens[2]:
+            yield f"<b>Stopped early after {i+1} iterations.</b>"
+            break
+
         # --- NOISING STEP ---
         threshold = get_noising_schedule(i, max_it, sharpness=sharpness)
         if use_confidence_noising:
@@ -226,13 +234,6 @@ def diffusion_chat(question, eot_weight, max_it, sharpness, noise_clipping, use_
         yield f"<b>Iteration {i+1}/{max_it} (after noising):</b><br>" + "".join(highlighted).replace('\n', '<br>')
         time.sleep(0.1)
 
-        # --- Early stopping ---
-        last_tokens.append(generated_tokens)
-        if len(last_tokens) > 3:
-            last_tokens.pop(0)
-        if len(last_tokens) == 3 and last_tokens[0] == last_tokens[1] == last_tokens[2]:
-            yield f"<b>Stopped early after {i+1} iterations.</b>"
-            break
 
     final_tokens = tokenizer.convert_ids_to_tokens(current_tokens[answer_start:])
     final_tokens = [tok for tok in final_tokens if tokenizer.convert_tokens_to_ids(tok) != eot_token_id]