Increment start_pos by encoded size in generate (#1462)

nlpfollower · nlpfollower · Jack-Khuu · web-flow · commit 0ee4aea4f14f · 2025-01-23T15:03:14.000-08:00
* Add encoded size to start_pos

* Only in chat mode

---------

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
@@ -1192,6 +1192,8 @@ def callback(x, *, done_generating=False):
                     max_seq_length=max_seq_length,
                     attention_backend=self.builder_args.attention_backend,
                 )
+                if generator_args.chat_mode:
+                    start_pos += encoded.size(0)
                 for token_tensor, metrics in generator_func:
                     if token_tensor is not None:
                         start_pos += token_tensor.size(0)

Original file line number	Diff line number	Diff line change
`@@ -1192,6 +1192,8 @@ def callback(x, *, done_generating=False):`
`1192`	`1192`	`max_seq_length=max_seq_length,`
`1193`	`1193`	`attention_backend=self.builder_args.attention_backend,`
`1194`	`1194`	`)`
	`1195`	`+ if generator_args.chat_mode:`
	`1196`	`+ start_pos += encoded.size(0)`
`1195`	`1197`	`for token_tensor, metrics in generator_func:`
`1196`	`1198`	`if token_tensor is not None:`
`1197`	`1199`	`start_pos += token_tensor.size(0)`