NanoThink-5M

João tem 3 maçãs e ganhou 2, quantas ele tem agora?

<THINK>
3 + 2 = 5
</THINK>
<ANSWER>
João tem 5 maçãs.
</ANSWER>

import torch
from tokenizers import Tokenizer
from model import NanoThink
from safetensors.torch import load_file

MODEL_PATH = "model.safetensors"
TOKENIZER_PATH = "tokenizer.json"


tokenizer = Tokenizer.from_file(TOKENIZER_PATH)

model = NanoThink(vocab_size=tokenizer.get_vocab_size())
model.load_state_dict(load_file(MODEL_PATH))
model.eval()

history = ""

while True:
    user_input = input("You: ")

    if user_input.lower() in ["get out", "exit", "quit"]:
        break

    prompt = history + f"\n<USER>\n{user_input}\n</USER>\n"

    input_ids = torch.tensor([tokenizer.encode(prompt).ids])

    output_ids = []

    for _ in range(120):
        logits = model(input_ids)
        next_token = torch.multinomial(torch.softmax(logits[0, -1], dim=-1), 1).item()

        input_ids = torch.cat([input_ids, torch.tensor([[next_token]])], dim=1)
        output_ids.append(next_token)

        text = tokenizer.decode(output_ids)

        if "</ANSWER>" in text:
            break

    output = tokenizer.decode(output_ids)

    
    if "<ANSWER>" in output:
        output = output.split("<ANSWER>")[1].split("</ANSWER>")[0]

    print("\n💬 Answer:")
    print(output.strip())
    print("\n" + "-"*50 + "\n")

    history += f"\n<USER>\n{user_input}\n</USER>\n<ANSWER>\n{output.strip()}\n</ANSWER>\n"