Safe sampling

app.py

@@ -135,9 +135,13 @@ 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"]
-        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()
+        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()
 
         # Extract confidence of selected tokens
         conf = probs[range(len(sampled)), sampled].cpu().numpy()