Remove unnecessary print statements - Add MASK noising

app.py

@@ -115,6 +115,7 @@ def noisify_answer(input_ids, answer_start, threshold=1.0, eot_weight=1.0, clust
     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
+        noised[idx] = tokenizer.encode('MASK', add_special_tokens = False)
 
     return noised, noised_indices
 
@@ -166,8 +167,6 @@ def generate_diffusion_text(input_ids):
         logits = logits.clamp(min=-1e4, max=1e4)
         probs = torch.nn.functional.softmax(logits, dim=-1)[0]
         probs = torch.clamp(probs, min=1e-8, max=1.0)
-        print("probs", probs)
-        print("probs min:", probs.min().item(), "max:", probs.max().item(), "sum:", probs.sum().item())
         assert torch.all(torch.isfinite(probs)), "Non-finite values in probs!"
         assert (probs >= 0).all(), "Negative probs!"
         sampled = torch.multinomial(probs, num_samples=1).squeeze(-1).tolist()