Loading a pre-trained lens#

From the hugging face API#

First check if there is a pre-trained lens available in our spaces’ pre-trained lenses folder.

Once you have found a lens that you want to use, you can simply load it. A tuned lens is always associated with a model that was used to train it so first load the model and then the lens.

>>> import torch
>>> from tuned_lens import TunedLens
>>> from transformers import AutoModelForCausalLM
>>> model = AutoModelForCausalLM.from_pretrained('EleutherAI/pythia-160m-deduped-v0')
>>> tuned_lens = TunedLens.from_model_and_pretrained(model)

If you want to load from your own code space you can override the default by providing the correct environment variables see tuned_lens.load_artifacts.

From the a local folder#

If you have trained a lens and want to load it for inference simply pass the model used to train it and the folder you saved it to.

>>> lens = TunedLens.from_model(model)
>>> # Do some thing
>>> lens.save(directory_path)
>>> lens = TunedLens.from_model_and_pretrained(model, directory_path)

Note the folder structure must look as follows:

path/to/folder
β”œβ”€β”€ config.json
└── params.pt

If you saved the model using tuned_lens.save("path/to/folder") then this should already be the case.