Visualize weight conversions

[degenfabian]

Description

This PR visualizes how the weights are loaded into TransformerLens from HuggingFace for all the different models. The user can access this visualization via two static functions from the class WeightConversionUtils. The function, model_info takes in the name of the model and prints out the visualization, whereas the function model_info_cfg takes in a HookedTransformerConfig but prints out the same visualization.

The visualizations works by implementing a __repr__ function for all the different conversion steps, and then calling them for all the different conversion steps performed by a given model architecture.

I also implemented a little fix, because the Gemma architecture was returning None for the blocks_conversion function, which resulted in the weights not being properly loaded into the different blocks of the Gemma architecture

I also had to move some imports to exactly where they were needed, because I would otherwise get circular import errors. I'm happy to revert that change if that is a problem that was only occurring for me locally.

There is no specific issue attached to this PR.

Type of change

• New feature (non-breaking change which adds functionality)

Screenshots

This is what the output looks like for the function call

WeightConversionUtils.model_info("bert-base-cased")

Checklist:

• I have commented my code, particularly in hard-to-understand areas
• I have made corresponding changes to the documentation
• My changes generate no new warnings
• I have added tests that prove my fix is effective or that my feature works
• New and existing unit tests pass locally with my changes
• I have not rewritten tests relating to key interfaces which would affect backward compatibility