FEATURE REQUEST: load_lens(model_name, layer_number) function implementation

[msakarvadia]

Have some way to load a trained attention lens (for a given layer for a given model). Len's for 1 layer for GPT2-small ~2GB, so not sure where to store trained weights (google drive).

Adding a function to allow users to easily load trained lenses will make it easier to use this repo.

[123epsilon]

So one thing to consider is that this is supported via torch.load and there's already a naming convention established with the checkpoint files. One thing we could do is, prepackaged a set of trained lenses and upload them to huggingface, then use a helper function with the signature you described to pull the appropriate weights and load them, how does that sound?

@msakarvadia

[msakarvadia]

Ah yes, that sounds perfect @123epsilon. We can check back in on this once we have a set of pretrained lenses.