批量推理

[codingforwhat]

while True:
prompt = input("Prompt: ")
inputs = tokenizer(prompt, return_tensors="pt").to(args.device)
response = model.generate(input_ids=inputs["input_ids"],
max_length=inputs["input_ids"].shape[-1] + args.max_new_tokens)
response = response[0, inputs["input_ids"].shape[-1]:]
print("Response:", tokenizer.decode(response, skip_special_tokens=True))
怎么修改这个代码进行批量推理

prompts = []

inputs = tokenizer(prompts, padding=True, truncation=True, return_tensors="pt")
input_ids = inputs.input_ids.to(args.device)
attention_mask = inputs.attention_mask.to(args.device)

response = model.generate(
input_ids=input_ids,
attention_mask=attention_mask,
max_length=input_ids.shape[-1] + args.max_new_tokens
)
这样行吗

[zRzRzRzRzRzRzR]

我们有batch推理的demo呀