Fix generation

app.py

@@ -171,15 +171,11 @@ def diffusion_chat(question, eot_weight, max_it, sharpness, noise_clipping, use_
     for i in range(max_it):
         print('Generating output')
 
-        # Compose full input: original prompt + current answer
-        full_input_tokens = ori_input_tokens[:answer_start] + current_tokens[answer_start:]
-        full_input_tokens = full_input_tokens[:256] + [pad_token] * max(0, 256 - len(full_input_tokens))
-
         # Model step
-        generated_tokens, confidences = generate_diffusion_text(full_input_tokens)
+        generated_tokens, confidences = generate_diffusion_text(current_tokens)
 
         # Save full output for noising step
-        current_tokens = generated_tokens
+        current_tokens = ori_input_tokens[answer_start] + generated_tokens[answer_start:]
 
         # --- GREEN HIGHLIGHT ---
         decoded_tokens = tokenizer.convert_ids_to_tokens(current_tokens[answer_start:])
@@ -202,17 +198,16 @@ def diffusion_chat(question, eot_weight, max_it, sharpness, noise_clipping, use_
         threshold = get_noising_schedule(i, max_it, sharpness=sharpness)
         if use_confidence_noising:
             noised_answer = confidence_guided_noising(
-                generated_tokens, answer_start, confidences, threshold, eot_weight, noise_clipping
+                current_tokens, answer_start, confidences, threshold, eot_weight, noise_clipping
             )
             just_noised_indices = []
         else:
             noised_answer, just_noised_indices = noisify_answer(
-                generated_tokens, answer_start, threshold=threshold, eot_weight=eot_weight, clustering=clustering
+                current_tokens, answer_start, threshold=threshold, eot_weight=eot_weight, clustering=clustering
             )
 
         # Compose full input again: prompt + noised answer
         current_tokens = ori_input_tokens[:answer_start] + noised_answer[answer_start:]
-        current_tokens = current_tokens[:256] + [pad_token] * max(0, 256 - len(current_tokens))
 
         # --- RED HIGHLIGHT ---
         decoded_tokens = tokenizer.convert_ids_to_tokens(current_tokens[answer_start:])