Remove random noising

app.py

@@ -20,11 +20,6 @@ pad_token = tokenizer.pad_token_id or tokenizer.eos_token_id
 eot_token_id = tokenizer.eos_token_id
 assistant_marker_ids = tokenizer.encode("Assistant:", add_special_tokens=False)
 
-# --- Load token probabilities ---
-with open("token_probabilities.json") as f:
-    token_probs_dict = json.load(f)
-token_probabilities = np.array([token_probs_dict[str(i)] for i in range(len(token_probs_dict))], dtype=np.float32)
-
 # def load_model():
 #     ckpt_path = hf_hub_download(
 #         repo_id="ruurd/tini_bi_m",
@@ -87,7 +82,7 @@ def get_noising_schedule(i, max_it, sharpness=5.0):
     x = i / max_it
     return (np.exp(-sharpness * x) - np.exp(-sharpness)) / (1 - np.exp(-sharpness))
 
-def noisify_answer(input_ids, answer_start, threshold=1.0, eot_weight=1.0, mask_weight=0.0, clustering=0.5, noise_start = 1.0):
+def noisify_answer(input_ids, answer_start, threshold=1.0, clustering=0.5, noise_start = 1.0):
     noised = input_ids.copy()
     answer_len = len(noised) - answer_start
     num_to_noise = int(threshold * answer_len * noise_start)
@@ -96,19 +91,6 @@ def noisify_answer(input_ids, answer_start, threshold=1.0, eot_weight=1.0, mask_
     if num_to_noise == 0:
         return noised, []
 
-    mixed_probs = token_probabilities.copy()
-
-    # Apply EOT weighting 
-    mixed_probs[eot_token_id] *= eot_weight
-
-    # Scale all other probabilities so they sum to 1 - mask_weight
-    total_other = mixed_probs.sum() - mixed_probs[mask_token_id]
-    scale = (1.0 - mask_weight) / total_other
-    mixed_probs *= scale
-    
-    # Set mask_token_id to mask_weight explicitly
-    mixed_probs[mask_token_id] = mask_weight
-
     num_clusters = max(1, int((1 - clustering) * num_to_noise))
     cluster_size = max(1, int(num_to_noise / num_clusters))
 
@@ -121,15 +103,14 @@ def noisify_answer(input_ids, answer_start, threshold=1.0, eot_weight=1.0, mask_
 
     noised_indices = sorted(list(noised_indices))[:num_to_noise]
 
-    noise = rng.choice(np.arange(vocab_size), size=len(noised_indices), p=mixed_probs)
-    for idx, val in zip(noised_indices, noise):
-        noised[idx] = val
+    for idx in noised_indices:
+        noised[idx] = mask_token_id
 
     return noised, noised_indices
 
 
 # Add new noising function
-def confidence_guided_noising(input_ids, answer_start, confidences, noise_clipping, threshold=1.0, eot_weight = 1.0, mask_weight = 0.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)
@@ -158,22 +139,8 @@ def confidence_guided_noising(input_ids, answer_start, confidences, noise_clippi
         p=weights
     )
 
-    mixed_probs = token_probabilities.copy()
-
-    # Apply EOT weighting 
-    mixed_probs[eot_token_id] *= eot_weight
-
-    # Scale all other probabilities so they sum to 1 - mask_weight
-    total_other = mixed_probs.sum() - mixed_probs[mask_token_id]
-    scale = (1.0 - mask_weight) / total_other
-    mixed_probs *= scale
-    
-    # Set mask_token_id to mask_weight explicitly
-    mixed_probs[mask_token_id] = mask_weight
-
-    noise = rng.choice(np.arange(vocab_size), size=num_to_noise, p=mixed_probs)
-    for idx, val in zip(indices, noise):
-        noised[idx] = val
+    for idx in indices:
+        noised[idx] = mask_token_id
 
     return noised
 
@@ -194,7 +161,7 @@ def generate_diffusion_text(input_ids):
     return sampled, conf 
 
 # --- Inference Wrapper ---
-def diffusion_chat(question, eot_weight, mask_weight, max_it, pause_length, sharpness, clustering, noise_start, use_confidence_noising, noise_clipping):
+def diffusion_chat(question, max_it, pause_length, sharpness, clustering, noise_start, use_confidence_noising, noise_clipping):
     placeholder = "What do you know about the city of New York?"
     placeholder = ""
     if question.strip() == "":
@@ -215,7 +182,7 @@ def diffusion_chat(question, eot_weight, mask_weight, max_it, pause_length, shar
 
     ori_input_tokens = input_ids
     current_tokens, just_noised_indices = noisify_answer(
-                input_ids, answer_start, threshold=1.0, eot_weight=eot_weight, mask_weight=mask_weight, clustering=clustering, noise_start = 1.0,
+                input_ids, answer_start, threshold=1.0, clustering=clustering, noise_start = 1.0,
             )
     yield f"<b>Iteration 0 (initial noise):</b><br>" + tokenizer.decode(current_tokens[answer_start:], skip_special_tokens=True).replace('\n', '<br>')
     time.sleep(pause_length)
@@ -262,12 +229,12 @@ def diffusion_chat(question, eot_weight, mask_weight, max_it, pause_length, shar
         threshold = get_noising_schedule(i, max_it, sharpness=sharpness)
         if use_confidence_noising:
             noised_answer = confidence_guided_noising(
-                current_tokens, answer_start, confidences, noise_clipping, threshold=threshold, eot_weight=eot_weight, mask_weight=mask_weight, noise_start=noise_start
+                current_tokens, answer_start, confidences, noise_clipping, threshold=threshold, noise_start=noise_start
             )
             just_noised_indices = []
         else:
             noised_answer, just_noised_indices = noisify_answer(
-                current_tokens, answer_start, threshold=threshold, eot_weight=eot_weight, mask_weight=mask_weight, clustering=clustering, noise_start = noise_start,
+                current_tokens, answer_start, threshold=threshold, clustering=clustering, noise_start = noise_start,
             )
 
         # Compose full input again: prompt + noised answer
@@ -306,8 +273,6 @@ demo = gr.Interface(
     fn=diffusion_chat,
     inputs=[
         gr.Textbox(label="User Question", lines=2, placeholder="What do you know about the city of New York?"),
-        gr.Slider(0, 1, value=0.5, step=0.05, label="↓ = longer answers (EOT weight)"),
-        gr.Slider(0, 1, value=0.5, step=0.05, label="↓ = more random answers (MASK weight)"),
         gr.Slider(1, 512, value=32, 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)"),