Add support for greedy decoding (#5)

- Add support for greedy decoding (3d62ef09768490af8a8c7867405942fcc79749b8)


Co-authored-by: Aditya Tomar <adityastomar@users.noreply.huggingface.co>

modeling_llada2_moe.py

@@ -1431,6 +1431,14 @@ class LLaDA2MoeModelLM(LLaDA2MoePreTrainedModel, GenerationMixin):
         orig_shape = logits.shape[:-1]
         vocab_size = logits.shape[-1]
         logits = logits.reshape(-1, vocab_size)
+        
+        # Greedy mode: temperature = 0, no top-k/p
+        if temperature == 0.0 and (top_k in (None, 0)) and (top_p is None or top_p >= 1.0):
+            probs = F.softmax(logits, dim=-1)
+            token = logits.argmax(dim=-1, keepdim=True)
+            token_prob = probs.gather(-1, token)
+            return token.view(*orig_shape), token_prob.view(*orig_shape)
+        
         if temperature > 0 and temperature != 1.0:
             logits = logits / temperature
         logits = self._top_k_logits(logits, top_k)