Add confidence based noising

app.py

@@ -73,6 +73,21 @@ def noisify_answer(input_ids, answer_start, threshold=1.0, eot_weight=1.0):
             noised[idx] = val
     return noised
 
+# Add new noising function
+def confidence_guided_noising(input_ids, answer_start, confidences, eot_weight):
+    noised = input_ids.copy()
+    mixed_probs = token_probabilities.copy()
+    mixed_probs[eot_token_id] *= eot_weight
+    mixed_probs /= mixed_probs.sum()
+
+    for i, conf in enumerate(confidences[answer_start:]):
+        p_noise = 1.0 - conf
+        if rng.random() < p_noise:
+            idx = answer_start + i
+            noised[idx] = rng.choice(np.arange(vocab_size), p=mixed_probs)
+
+    return noised
+
 @spaces.GPU
 def generate_diffusion_text(input_ids, answer_start):
     with torch.no_grad():
@@ -81,18 +96,21 @@ def generate_diffusion_text(input_ids, answer_start):
         probs = torch.nn.functional.softmax(logits, dim=-1).squeeze()
         probs = torch.clamp(probs, min=1e-8, max=1.0)
         sampled = torch.multinomial(probs, num_samples=1).squeeze().tolist()
-    return input_ids[:answer_start] + sampled[answer_start:]
+
+        # Extract confidence of selected tokens
+        conf = probs[range(len(sampled)), sampled].cpu().numpy()
+    return sampled, conf  # ✅ NEW: Return confidence
 
 # --- Inference Wrapper ---
 
 
-def diffusion_chat(question, eot_weight, max_it, sharpness):
+# Modify diffusion_chat to use new noising conditionally
+def diffusion_chat(question, eot_weight, max_it, sharpness, use_confidence_noising):
     placeholder = "What do you know about the city of New York?"
     if question.strip() == "":
         question = placeholder
-    
-    print('started generation')
 
+    print('started generation')
     prompt = f"User: {question}\nAssistant:"
     input_ids = tokenizer.encode(prompt, add_special_tokens=False)
     answer_start = find_answer_start(input_ids, assistant_marker_ids)
@@ -112,7 +130,7 @@ def diffusion_chat(question, eot_weight, max_it, sharpness):
 
     for i in range(max_it):
         print('Generating output')
-        generated_tokens = generate_diffusion_text(current_tokens, answer_start)
+        generated_tokens, confidences = generate_diffusion_text(current_tokens, answer_start)
         current_tokens = generated_tokens
 
         decoded_ids = current_tokens[answer_start:]
@@ -141,7 +159,11 @@ def diffusion_chat(question, eot_weight, max_it, sharpness):
             break
 
         threshold = get_noising_schedule(i, max_it, sharpness=sharpness)
-        current_tokens = noisify_answer(generated_tokens, answer_start, threshold=threshold, eot_weight=eot_weight)
+        if use_confidence_noising:
+            current_tokens = confidence_guided_noising(generated_tokens, answer_start, confidences, eot_weight)
+        else:
+            current_tokens = noisify_answer(generated_tokens, answer_start, threshold=threshold, eot_weight=eot_weight)
+
         time.sleep(0.01)
 
     final_tokens = tokenizer.convert_ids_to_tokens(current_tokens[answer_start:])
@@ -162,7 +184,8 @@ demo = gr.Interface(
         gr.Textbox(label="User Question", lines=2, placeholder="What do you know about the city of New York?"),
         gr.Slider(0, 1, value=0.4, step=0.05, label="↓ = longer answers (EOT weight)"),
         gr.Slider(1, 512, value=64, step=1, label="↑ = more iterations"),
-        gr.Slider(1.0, 20.0, value=5.0, step=0.5, label="↓ = more noising (sharpness)")
+        gr.Slider(1.0, 20.0, value=5.0, step=0.5, label="↓ = more noising (sharpness)"),
+        gr.Checkbox(value=False, label="Use confidence-guided noising")  # ✅ NEW
     ],
     outputs=[gr.HTML(label="Diffusion Output")],
     title="Diffusion Language Model Chat",