Load directly huggingface/transformers models

[jbesomi]

Hi,
thank you for the nice tool!
I would like to understand why you are not loading the models directly from the transformers package.

Which part of the transformers model do you need to adapt to make it compatible with bertviz? I would like to add other huggingface models such as DistilBERT.

Did you consider a way to load the model directly from transformers?

Thank you in advance for your kind response.

[jessevig]

Hi Jonathan,
You are right! At least for the attention-head view and the model view I should not be relying on a forked version of transformers. The original version of the transformers package didn't give access to the attention weights so this is why I create a modified version in Bertviz. Since then, transformers has made the attention weights available, but I haven't had time to incorporate this into BertViz. The Neuron View still requires access to the query and key vectors, which aren't made available in the transformers package, so I would still have to do something special for that visualization. So my plan going forward is to load the models directly from transformers for the Head and Model views, enabling use of the latest models. Hopefully I can get to that in the next few weeks.

[jessevig]

Update BertViz now allows you to load the model directly from the transformers package, at least for the model and head views. So you should be able to visualize any model that the transformers library supports. Please give it a try!