Implement auto_max_new_tokens for ExLlama

oobabooga · Aug 2, 2023 · 32a2bbe · 32a2bbe
1 parent e931844
commit 32a2bbe
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 1 deletion.
diff --git a/modules/exllama.py b/modules/exllama.py
@@ -94,11 +94,15 @@ def generate_with_streaming(self, prompt, state):
         # Tokenizing the input
         ids = self.generator.tokenizer.encode(prompt)
         ids = ids[:, -get_max_prompt_length(state):]
+        if state['auto_max_new_tokens']:
+            max_new_tokens = state['truncation_length'] - ids.shape[-1]
+        else:
+            max_new_tokens = state['max_new_tokens']
 
         self.generator.gen_begin_reuse(ids)
         initial_len = self.generator.sequence[0].shape[0]
         has_leading_space = False
-        for i in range(state['max_new_tokens']):
+        for i in range(max_new_tokens):
             token = self.generator.gen_single_token()
             if i == 0 and self.generator.tokenizer.tokenizer.IdToPiece(int(token)).startswith('▁'):
                 has_leading_space = True

diff --git a/modules/loaders.py b/modules/loaders.py
@@ -151,6 +151,7 @@
         'repetition_penalty_range',
         'seed',
         'ban_eos_token',
+        'auto_max_new_tokens',
     },
     'AutoGPTQ': {
         'temperature',