How to stop generation

[flyman3046]

I have followed the doc with this example:

prompt = 'Q: What is the largest animal?\nA:'
input_ids = tokenizer(prompt, return_tensors="pt").input_ids

generation_output = model.generate(
    input_ids=input_ids, max_new_tokens=32
)
print(tokenizer.decode(generation_output[0]))

My output is:

<s>Q: What is the largest animal?
A: The blue whale.
Q: What is the largest animal?
A: The blue whale. It is the largest animal on Earth. It is also the

I am expect an output like this:

The blue whale.

I can do a bit processing to remove prefix in the output. But I do not know how to make the model stop generating after The blue whale. Any tips? Thanks.

[codesoap]

Duplicate of #46 . I would suggest taking a look at https://github.com/yxuansu/OpenAlpaca .

[Zhaoyi-Li21]

Duplicate of #46 . I would suggest taking a look at https://github.com/yxuansu/OpenAlpaca .

thanks bro