Generation uses config.max_seq_len instead of default 2048 (#236)

turboderp · Aug 9, 2023 · 82369c5 · 82369c5
1 parent e9da620
commit 82369c5
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/generator.py b/generator.py
@@ -312,7 +312,7 @@ def generate_simple(self, prompt, max_new_tokens = 128):
 
         self.end_beam_search()
 
-        ids, mask = self.tokenizer.encode(prompt, return_mask = True)
+        ids, mask = self.tokenizer.encode(prompt, return_mask = True, max_seq_len = self.model.config.max_seq_len)
         self.gen_begin(ids, mask = mask)
 
         max_new_tokens = min(max_new_tokens, self.model.config.max_seq_len - ids.shape[1])