Increment start_pos by encoded size in generate (#1462)

* Add encoded size to start_pos * Only in chat mode --------- Co-authored-by: nlpfollower <[email protected]> Co-authored-by: Jack-Khuu <[email protected]>
pytorch · Jan 23, 2025 · 0ee4aea · 0ee4aea
1 parent 3ce9c8e
commit 0ee4aea
Showing 1 changed file with 2 additions and 0 deletions.
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)