from transformers import pipeline

pipe = pipeline(
    "text-generation", model="bart1259/MiniCOTMath"
)
print(pipe("Input: (5 + 5)\n", max_new_tokens=100)[0]["generated_text"])

Outputs:

Input: (5 + 5)

Step 1:
(5 + 5)
(5 + 5) = 10

Step 2:
10
Final Result: 10
<end>

Input: (3 * 8)

Step 1:
(3 * 8)
(3

To end at the token, you can setup streaming like:

from transformers import AutoTokenizer, AutoModelForCausalLM, TextStreamer, StoppingCriteria
from transformers import StoppingCriteria

class StopCriteria(StoppingCriteria):
    def __call__(self, input_ids, scores, **kwargs):
        generated_text = tokenizer.decode(input_ids[0])
        return "<end>" in generated_text

    def __len__(self):
        return 1

    def __iter__(self):
        yield self

prompt = "Input: (5 + 5)\n"

tokenizer = AutoTokenizer.from_pretrained("bart1259/MiniCOTMath")
model = AutoModelForCausalLM.from_pretrained("bart1259/MiniCOTMath").cuda()

encoded_input = tokenizer(prompt, return_tensors='pt')
input_ids=encoded_input['input_ids'].cuda()

streamer = TextStreamer(tokenizer=tokenizer, skip_prompt=False)
_ = model.generate(
    input_ids,
    streamer=streamer,
    pad_token_id=tokenizer.eos_token_id,
    do_sample=True,
    temperature=0.25,
    max_new_tokens=256,
    stopping_criteria=StopCriteria()
)

Outputs:

Input: (5 + 5)

Step 1:
(5 + 5)
(5 + 5) = 10

Step 2:
10
Final Result: 10
<end>