lm-checkpoints 0.1.12

Simple library for loading checkpoints of language models.

lm-checkpoints

Simple library for dealing with language model checkpoints.

Install using pip install lm-checkpoints.

Example

Say you want to compute some metrics for all model checkpoints of Pythia 160m, but only seed 0.

from lm_checkpoints import PythiaCheckpoints

for ckpt in PythiaCheckpoints(size=160,seed=[0]):
    # Do something with ckpt.model or ckpt.tokenizer

Or if you only want to load steps 0, 1, 2, 4, 8, 16 for all available seeds:

from lm_checkpoints import PythiaCheckpoints

for ckpt in PythiaCheckpoints(size=160,step=[0, 1, 2, 4, 8, 16]):
    # Do something with ckpt.model or ckpt.tokenizer

Alternatively, you may want to load all final checkpoints of MultiBERTs:

from lm_checkpoints import MultiBERTCheckpoints

for ckpt in MultiBERTCheckpoints.final_checkpoints():
    # Do something with ckpt.model or ckpt.tokenizer

In case you don't want the checkpoints to fill up your space, use clean_cache=True to delete older checkpoints:

from lm_checkpoints import PythiaCheckpoints

for ckpt in PythiaCheckpoints(size=14,clean_cache=True):
    # Do something with ckpt.model or ckpt.tokenizer