litgpt/tutorials/python-api.md

LitGPT Python API

This is a work-in-progress draft describing the current LitGPT Python API (experimental and subject to change).

Model loading

Download a model using the CLI:

litgpt download EleutherAI/pythia160-m

Then, load the model in Python:

from litgpt.api import LLM
llm = LLM.load("EleutherAI/pythia-160m", accelerator="cuda", devices=1)

Generate/Chat

Generate output using the .generate method:

text = llm.generate("What do Llamas eat?", top_k=1)
print(text)

What do Llamas eat?

"A lot of people, the Llamas, I was, a Llamas, a lama, a lama, a lama, a lama, a lama, a lama, a lama, a