Update generate.py

generate.py

@@ -1,6 +1,8 @@
 import torch
 
-def generate(model, input_ids, generation_config, **kwargs):
+def generate(model, generation_config, **kwargs):
+    input_ids = kwargs.get("input_ids") or kwargs.get("input")
+    
     generation_config = generation_config or model.generation_config  # default to the model generation config
     cur_length = input_ids.shape[1]
     max_length = generation_config.max_length or cur_length + generation_config.max_new_tokens