calculator-6m

Model-calculator. See demo here.

Works well with simple calculations, but fails with complex ones.

Usage

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

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

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)

Correct Examples

find 2 + 3
step 2 + 3 = 5
answer 5

find (2 + 3) / 2
step 2 + 3 = 5
step 5 / 2 = 2.5
answer 2.5

find 0.2 + 0.4
step 0.2 + 0.4 = 0.6
answer 0.6

find 1000 + 1500
step 1000 + 1500 = 2500
answer 2500

find 10 * 20
step 10 * 20 = 200
answer 200

find 10 * 0.25
step 10 * 0.25 = 2.5
answer 2.5

Almost Correct Examples

find 0.5 + 0.25
step 0.5 + 0.25 = 0.78
answer 0.78

find 100 / 12 + 1
step 100 / 12 = 8.5
step 8.5 + 1 = 9.5
answer 9.5

find 999 / 102
step 999 / 102 = 9.8391
answer 9.8391