Improve confidence guided noising and show number of tokens generated

app.py

@@ -110,35 +110,48 @@ def noisify_answer(input_ids, answer_start, threshold=1.0, clustering=0.5, noise
 
 
 # Add new noising function
-def confidence_guided_noising(input_ids, answer_start, confidences, noise_clipping, threshold=1.0, noise_start = 1.0):
+def confidence_guided_noising(input_ids, answer_start, confidences, noise_clipping, threshold=1.0, noise_start=1.0):
     noised = input_ids.copy()
     answer_len = len(input_ids) - answer_start
     num_to_noise = int(threshold * answer_len * noise_start)
-    
     if num_to_noise == 0:
         return noised
 
-    
-    raw_weights = 1.0 - np.array(confidences[answer_start:])
-
-    # Avoid zero-probability weights for selection
-    # If noise clipping == 1, all tokens have equal chance to be noised. 
-    # If noise_clipping == 0.00001, all tokens are noised according to the confidence of the past prediction
-    raw_weights = np.clip(raw_weights, a_min = noise_clipping, a_max = None)
+    all_indices = np.arange(answer_start, len(input_ids))
+    eos_indices = [i for i in all_indices if input_ids[i] == eos_token_id]
+    non_eos_indices = [i for i in all_indices if input_ids[i] != eos_token_id]
 
-    weights = raw_weights / raw_weights.sum()
+    num_non_eos_to_noise = int(num_to_noise * (len(non_eos_indices) / (len(non_eos_indices) + len(eos_indices) + 1e-5)))
+    num_eos_to_noise = num_to_noise - num_non_eos_to_noise
 
-    if num_to_noise > len(weights):
-        num_to_noise = len(weights)  # prevent oversampling
+    # === Non-EOS sampling ===
+    raw_weights_non_eos = 1.0 - np.array([confidences[i - answer_start] for i in non_eos_indices])
+    raw_weights_non_eos = np.clip(raw_weights_non_eos, a_min=noise_clipping, a_max=None)
+    weights_non_eos = raw_weights_non_eos / raw_weights_non_eos.sum() if raw_weights_non_eos.sum() > 0 else None
 
-    indices = rng.choice(
-        np.arange(answer_start, len(input_ids)),
-        size=num_to_noise,
+    chosen_non_eos = rng.choice(
+        non_eos_indices,
+        size=min(num_non_eos_to_noise, len(non_eos_indices)),
         replace=False,
-        p=weights
-    )
+        p=weights_non_eos
+    ) if weights_non_eos is not None else []
+
+    # === EOS sampling ===
+    if eos_indices:
+        raw_weights_eos = 1.0 - np.array([confidences[i - answer_start] for i in eos_indices])
+        raw_weights_eos = np.clip(raw_weights_eos, a_min=noise_clipping, a_max=None)
+        weights_eos = raw_weights_eos / raw_weights_eos.sum() if raw_weights_eos.sum() > 0 else None
+
+        chosen_eos = rng.choice(
+            eos_indices,
+            size=min(num_eos_to_noise, len(eos_indices)),
+            replace=False,
+            p=weights_eos
+        ) if weights_eos is not None else []
+    else:
+        chosen_eos = []
 
-    for idx in indices:
+    for idx in list(chosen_non_eos) + list(chosen_eos):
         noised[idx] = mask_token_id
 
     return noised
@@ -256,11 +269,19 @@ def diffusion_chat(question, max_it, pause_length, sharpness, clustering, noise_
         time.sleep(pause_length)
 
 
-    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) != eos_token_id]
-    final_output = tokenizer.convert_tokens_to_string(final_tokens)
+    answer_ids = current_tokens[answer_start:]
+    try:
+        eos_index = answer_ids.index(eos_token_id)
+        final_ids = answer_ids[:eos_index]
+    except ValueError:
+        final_ids = answer_ids
+    
+    num_tokens = len(final_ids)
+    final_output = tokenizer.decode(final_ids, skip_special_tokens=True)
+    
     print(final_output)
-    yield f"<b>Final Output (after {i+1} iterations):</b><br>" + final_output.replace('\n', '<br>')
+    yield f"<b>Final Output ({num_tokens} tokens after {i+1} iterations):</b><br>" + final_output.replace('\n', '<br>')
+
 
 # --- Gradio Interface ---
 print("Loading model...")
@@ -271,11 +292,11 @@ demo = gr.Interface(
     fn=diffusion_chat,
     inputs=[
         gr.Textbox(label="User Question", lines=2, placeholder="What do you know about the city of Amsterdam?"),
-        gr.Slider(1, 512, value=32, step=1, label="↑ = more iterations"),
+        gr.Slider(1, 512, value=64, step=1, label="↑ = more iterations"),
         gr.Slider(0.01, 5, value=0.01, step=0.01, label="↑ = longer pause (for visualization)"),
-        gr.Slider(1.0, 20.0, value=5.0, step=0.5, label="↓ = more noising (sharpness)"),
+        gr.Slider(1.0, 20.0, value=1.0, step=0.5, label="↓ = more noising (sharpness)"),
         gr.Slider(0.0, 1.0, value=0.0, step=0.05, label="↑ = more clustered noising (fewer, larger edits)"),
-        gr.Slider(0.0, 1.0, value=0.5, step=0.05, label="↑ = more noise (noise start)"),
+        gr.Slider(0.0, 1.0, value=0.2, step=0.05, label="↑ = more noise (noise start)"),
         gr.Checkbox(value=False, label="Use confidence-guided noising"),
         gr.Slider(0.01, 1.0, value=0.01, step=0.01, label="↓ = more confidence guidance (noise clipping)"),