Skip to main content

Ready-to-go example code

Quick start​

This page provides a minimal example to get you started with Ai2 models. We'll use OLMo-2-0425-1B, our smallest model that can run on a CPU, making it capable of running on a very wide variety of hardware. Note that while this model will run on modern CPUs, it may still run very slowly. For optimizing runtime performance on CPUs, see our quantization guide.

Before running the code below, make sure you have:

  1. Python 3.8+ installed
  2. A virtual environment set up (we will use venv in these docs, but conda is fine too) 
    python -m venv olmo-env
    source olmo-env/bin/activate  # On Windows: olmo-env\Scripts\activate
  3. PyTorch and Transformers installed (tested with PyTorch 2.7.0 and transformers 4.51.3) 
    pip install torch transformers

For detailed setup instructions, see our Running Locally guide.

Hello world: getting started​

Get started with Olmo in just a few lines of code:

from transformers import AutoTokenizer, AutoModelForCausalLM
import torch

model = AutoModelForCausalLM.from_pretrained("allenai/OLMo-2-0425-1B")
tokenizer = AutoTokenizer.from_pretrained("allenai/OLMo-2-0425-1B")
model.to("cuda" if torch.cuda.is_available() else "cpu")

outputs = model.generate(**tokenizer("Hello", return_tensors="pt"), max_length=20)
print(tokenizer.decode(outputs[0], skip_special_tokens=True))

Example output:​

Hello, I'm trying to figure out how many more miles I need to run to reach my goal.

This simple example demonstrates how to:

  1. Load the model and tokenizer
  2. Move the model to GPU if available
  3. Generate text from a prompt
  4. Decode and print the output