Error with model creation from timm "vit_base_patch32_224" : 'TimmWrapperConfig' object has no attribute 'num_hidden_layers'

[andrewcaunes]

Hi, thank you for your work !

Just tried to run some of the main demo in a clone of the 2.0 version of the lib and got an error, I guess the demo should be updated.
Fresh conda env with only "pip install vit_prisma"

Code ran :
`from vit_prisma.models.base_vit import HookedViT
from vit_prisma.configs.HookedViTConfig import HookedViTConfig
from vit_prisma.training import trainer
from vit_prisma.dataloaders.induction import InductionDataset

from vit_prisma.training.training_utils import PrismaCallback
from vit_prisma.models.base_vit import HookedViT

model_function = HookedViT.from_pretrained("vit_base_patch32_224",
center_writing_weights=True,
center_unembed=True,
fold_ln=True,
refactor_factored_attn_matrices=True,
)`

Error:
Traceback (most recent call last):
File "", line 198, in _run_module_as_main
File "", line 88, in _run_code
File "/projects/ViT-Prisma/project/test.py", line 26, in
model_function = HookedViT.from_pretrained("vit_base_patch32_224",
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/projects/ViT-Prisma/src/vit_prisma/models/base_vit.py", line 780, in from_pretrained
cfg = convert_pretrained_model_config(
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/projects/ViT-Prisma/src/vit_prisma/prisma_tools/loading_from_pretrained.py", line 633, in convert_pretrained_model_config
'n_layers' : hf_config.num_hidden_layers,
^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/miniconda3/envs/prisma_1/lib/python3.12/site-packages/transformers/configuration_utils.py", line 211, in getattribute
return super().getattribute(key)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
AttributeError: 'TimmWrapperConfig' object has no attribute 'num_hidden_layers'

[d3tk]

Also having this issue running the Demo notebook ViT Prisma Main Demo

[d3tk]

@andrewcaunes to solve this in the short term in vit_prism/prisma_tools/loading_from_pretrained.py I changed line 351 pretrained_config to be:

pretrained_config = {
        'n_layers' : hf_config.num_hidden_layers if not is_timm else len(model.blocks),
        'd_model' : hf_config.hidden_size if not is_timm else model.embed_dim,
        'd_head' : hf_config.hidden_size // hf_config.num_attention_heads if not is_timm else model.embed_dim // model.blocks[0].attn.num_heads,
        'model_name' : hf_config._name_or_path,
        'n_heads' : hf_config.num_attention_heads if not is_timm else model.blocks[0].attn.num_heads,
        'd_mlp' : hf_config.intermediate_size if not is_timm else 4 * model.embed_dim,
        'activation_name' : hf_config.hidden_act if not is_timm else model.blocks[0].mlp.act.__class__.__name__.lower(),
        'eps' : getattr(hf_config, 'layer_norm_eps', 1e-12),
        'original_architecture' : getattr(hf_config, 'architecture', getattr(hf_config, 'architectures', None)),
        'initializer_range' : hf_config.initializer_range,
        'n_channels' : hf_config.num_channels if not is_timm else model.patch_embed.proj.weight.shape[1],
        'patch_size' : ps if not is_timm else model.patch_embed.proj.kernel_size[0],
        'image_size' : hf_config.image_size if not is_timm else model.default_cfg['input_size'][1],
        'n_classes' :  getattr(hf_config, "num_classes", None) if not is_timm else model.default_cfg['num_classes'],
        'n_params' : sum(p.numel() for p in model.parameters() if p.requires_grad) if is_timm else None,
    }