Auto-add timm tag to timm-wrapper models. (#35794)

Works for fine-tuned or exported models:

```py
from transformers import AutoModelForImageClassification

checkpoint = "timm/vit_base_patch16_224.augreg2_in21k_ft_in1k"
model = AutoModelForImageClassification.from_pretrained(checkpoint)

model.push_to_hub("pcuenq/tw1")
```

The uploaded model will now show snippets for both the timm and the
transformers libraries.
This commit is contained in:
Pedro Cuenca 2025-01-21 14:34:45 +01:00 committed by GitHub
parent dc10f7906a
commit 678bd7f1ce
No known key found for this signature in database
GPG key ID: B5690EEEBB952194

View file

@ -81,6 +81,7 @@ class TimmWrapperPreTrainedModel(PreTrainedModel):
main_input_name = "pixel_values"
config_class = TimmWrapperConfig
_no_split_modules = []
model_tags = ["timm"]
# used in Trainer to avoid passing `loss_kwargs` to model forward
accepts_loss_kwargs = False