Update generate.py

generate.py

@@ -1,15 +1,16 @@
 import torch
 
-def generate(model, model_inputs, generation_config, **kwargs):
-    cur_length = model_inputs["input_ids"].shape[1]
+def generate(model, input_ids, generation_config, **kwargs):
+    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
     
     while cur_length < max_length:
-        logits = model(model_inputs["input_ids"]).logits
+        logits = model(input_ids).logits
         next_token_logits = logits[:, -1, :]
         next_tokens = torch.argmax(next_token_logits)
-        model_inputs["input_ids"] = torch.cat((model_inputs["input_ids"], next_tokens), dim=-1)
+        model_inputs["input_ids"] = torch.cat((input_ids, next_tokens), dim=-1)
         cur_length += 1
 
-    return model_inputs["input_ids"]
+    return input_ids