MicroBananaMind-v1

MicroBananaMind-v1

MicroBananaMind-v1 is a very small causal language model trained from scratch on FineWeb-Edu, FineMath, and Cosmopedia-v2.

The model has 902,272 parameters and uses a custom 1536-token byte-level BPE tokenizer with digit-aware tokenization It is our smallest model ever that is not just a TinyStories model.

Model Details

Tokenizer

MicroBananaMind-v1 uses our digit-aware 1536-token tokenizer.

Training Data

Training setup:

We recommend using a temperature of 0 or 0.1

Usage

This model uses custom architecture code, so load it with trust_remote_code=True.

pip install -U transformers safetensors torch

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

model_id = "BananaMind/MicroBananaMind-v1"

tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    trust_remote_code=True,
    torch_dtype=torch.bfloat16 if torch.cuda.is_available() and torch.cuda.is_bf16_supported() else torch.float16,
).cuda().eval()

prompt = "The color of the sky is "
input_ids = tokenizer(prompt, return_tensors="pt").input_ids.to(model.device)

with torch.no_grad():
    output = model.generate(
        input_ids=input_ids,
        max_new_tokens=64,
        do_sample=False,
        repetition_penalty=1.1,
        pad_token_id=tokenizer.eos_token_id,
        eos_token_id=tokenizer.eos_token_id,
    )

print(tokenizer.decode(output[0], skip_special_tokens=True))

License

Apache 2.0