Update generate.py

generate.py

@@ -11,7 +11,7 @@ def custom_generate(
     self,
     input_ids,
     attention_mask=None,
-    max_length=None,
+    max_new_tokens=None,
     min_length=None,
     do_sample=None,
     early_stopping=None,
@@ -47,7 +47,7 @@ def custom_generate(
     with torch.no_grad():
         finished_generating = torch.zeros(len(input_ids), dtype=torch.bool, device=device)
 
-        while not finished_generating.all() and input_ids.shape[1] < max_length:
+        for cur_token_idx in range(max_new_tokens):
             # Sample the next token
             new_ids = self(
                 input_ids[~finished_generating],
@@ -86,6 +86,9 @@ def custom_generate(
                 # Check if the end token is generated
                 if new_ids_sampled == self.tokenizer.convert_tokens_to_ids("</s>"):
                     finished_generating[answer_idx] = 1
+            
+            if finished_generating.all():
+                break
 
             if streamer is not None:
                 streamer.put(new_ids_sampled)
@@ -98,7 +101,7 @@ def generate(
     self,
     input_ids,
     attention_mask=None,
-    max_length=None,
+    max_new_tokens=None,
     min_length=None,
     do_sample=None,
     early_stopping=None,
@@ -169,7 +172,7 @@ def generate(
         self,
         input_ids=input_ids,
         attention_mask=attention_mask,
-        max_length=max_length,
+        max_new_tokens=max_new_tokens,
         min_length=min_length,
         do_sample=do_sample,
         early_stopping=early_stopping,