Fix clamping and introduce top-k and top-p filtering

app.py

@@ -151,13 +151,40 @@ def confidence_guided_noising(input_ids, answer_start, confidences, noise_clippi
     noised_indices = sorted(noised_indices)
     return noised, noised_indices
 
+def filter_logits(logits, top_k=0, top_p=0.0):
+    """Filter logits per position for top-k / nucleus (top-p) sampling."""
+    logits = logits.clone()  # don't modify in-place
+    batch_size, seq_len, vocab_size = logits.shape
+
+    for i in range(seq_len):
+        token_logits = logits[0, i]
+
+        if top_k > 0:
+            top_values, _ = torch.topk(token_logits, top_k)
+            threshold = top_values[-1]
+            token_logits[token_logits < threshold] = float("-inf")
+
+        if top_p > 0.0:
+            sorted_logits, sorted_indices = torch.sort(token_logits, descending=True)
+            cumulative_probs = torch.softmax(sorted_logits, dim=-1).cumsum(dim=-1)
+
+            sorted_indices_to_remove = cumulative_probs > top_p
+            sorted_indices_to_remove[1:] = sorted_indices_to_remove[:-1].clone()
+            sorted_indices_to_remove[0] = 0  # always keep at least 1 token
+
+            token_logits[sorted_indices[sorted_indices_to_remove]] = float("-inf")
+
+        logits[0, i] = token_logits
+
+    return logits
 
 @spaces.GPU
 def generate_diffusion_text(input_ids):
     with torch.no_grad():
         input_tensor = torch.tensor([input_ids], dtype=torch.long).to(model.device)
         logits = model(input_ids=input_tensor)["logits"]
-        logits = logits.clamp(min=-1e4, max=1e4)
+        logits = filter_logits(logits, top_k=top_k, top_p=top_p) 
+        logits = logits.clamp(min=-1e8, max=1e4)
         probs = torch.nn.functional.softmax(logits, dim=-1)[0]
         probs = torch.clamp(probs, min=1e-8, max=1.0)
         assert torch.all(torch.isfinite(probs)), "Non-finite values in probs!"