Decode_n_tokens yield eos/eot token

nlpfollower · nlpfollower · commit 2ab9092bc54f · 2025-01-23T13:23:03.000-08:00
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:`