Closed SumanthRH closed 7 months ago
Okay so I ran some more tests on using bfloat16
and it works for primary attributions, neuron factors, output token scores, token rankings and for visualizing the evolution of a token. Hoping to see this merged @jalammar !
Thanks @SumanthRH! Merging this into a branch to update the CI/CD python versions.
Which version of python did you test this on?
I tested this on python 3.10!
What does this PR do?
I've been playing around with the library recently and noticed that some of the files were a bit outdated. This is a simple PR to
kwargs
to the model during initialization. This will allow us to, for example, use bfloat16 dtype for memory efficiency. I made a small change in theLM.to()
method for dtype casting.Ex:
lm = ecco.from_pretrained("distilgpt2", activations=True, torch_dtype=torch.bfloat16)
I've tested attributions and neuron factors notebooks with the new changes and things seem to work well