import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

tokenizer = AutoTokenizer.from_pretrained("georgiyozhegov/calculator-8m")
model = AutoModelForCausalLM.from_pretrained("georgiyozhegov/calculator-8m")

prompt = "find 2 + 3\nstep"

inputs = tokenizer(prompt, return_tensors="pt", return_token_type_ids=False)

with torch.no_grad():
    outputs = model.generate(
        input_ids=inputs["input_ids"],
        attention_mask=inputs["attention_mask"],
        max_length=32,
        do_sample=True,
        top_k=50,
        top_p=0.98
    )

# Cut the rest
count = 0
for index, token in enumerate(outputs[0]):
    if token == 6: count += 1
    if count >= 2: break

output = tokenizer.decode(outputs[0][:index])
print(output)