Decode_n_tokens yield eos/eot token (#1474)

nlpfollower · nlpfollower · Jack-Khuu · web-flow · commit 53a100462cde · 2025-02-07T10:10:39.000-08:00
Co-authored-by: nlpfollower &lt;nlp@nlpfollower.nlp&gt;
Co-authored-by: Jack-Khuu &lt;jack.khuu.7@gmail.com&gt;
diff --git a/torchchat/generate.py b/torchchat/generate.py
@@ -576,6 +576,7 @@ def decode_n_tokens(
                         **sampling_kwargs,
                     )
                     input_pos += 1
+                    yield cur_token.clone(), next_prob.clone()
                     break
 
         if not encountered_eos:

Original file line number	Diff line number	Diff line change
`@@ -576,6 +576,7 @@ def decode_n_tokens(`
`576`	`576`	`**sampling_kwargs,`
`577`	`577`	`)`
`578`	`578`	`input_pos += 1`
	`579`	`+ yield cur_token.clone(), next_prob.clone()`
`579`	`580`	`break`
`580`	`581`
`581`	`582`	`if not encountered_eos:`