Update generate.py

generate.py

@@ -47,8 +47,6 @@ def custom_generate(
     with torch.no_grad():
         finished_generating = torch.zeros(len(input_ids), dtype=torch.bool, device=device)
 
-        if max_new_tokens is None:
-            max_new_tokens = 50  # Default value if not specified
         for cur_token_idx in range(max_new_tokens):
             # Sample the next token
             new_ids = self(
@@ -95,8 +93,7 @@ def custom_generate(
             if streamer is not None:
                 streamer.put(new_ids_sampled)
 
-        generated_token_ids = input_ids.tolist()
-        return generated_token_ids, attention_mask
+        return input_ids
 
 
 def generate(
@@ -170,7 +167,7 @@ def generate(
     self.rm_initialized = True
     self.original_mode = False
 
-    generated_token_ids, attention_mask = custom_generate(
+    generated_token_ids = custom_generate(
         self,
         input_ids=input_ids,
         attention_mask=attention_mask,
@@ -207,4 +204,4 @@ def generate(
         **model_kwargs,
     )
 
-    return generated_token_ids, attention_mask
+    return generated_token_ids