---
datasets:
- georgiyozhegov/g.arithmetic
library_name: transformers
license: mit
pipeline_tag: text-generation
tags:
- trl
- sft
---

Model-calculator. See demo [here](https://huggingface.co/spaces/georgiyozhegov/calculator-demo).

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

# Usage

```python
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
```