search

EleutherAI / elk

Keeping language models honest by directly eliciting knowledge encoded in their activations.
MIT License
178 stars 33 forks source link

Use dill and Apache Arrow directly for caching & storing hidden states #203

Open norabelrose opened 1 year ago

norabelrose commented 1 year ago

Right now we save hidden states as a HF dataset so that we can use its magical caching abilities. We're clearly wrangling HF into supporting something it wasn't designed for and this does cause us to need to load the model from disk a bunch of times during a sweep, whereas we could load it only once if we had more control. I expect we could get at least a 2x speedup from this for large models, and the code could be cleaner. It would also remove the confusing "Downloading and preparing" messages that currently pop up when we extract hiddens

norabelrose commented 1 year ago

This is kinda low priority but would still like to do this at some point