FEAT: add tagging support to axolotl

[younesbelkada]

Hi there,

I would like to introduce a feature request to axolotl; automatic tagging when users push the trained model on the Hub. That way, you can easily filter on 🤗 Hub models that has been trained with axolotl with a simple filter: https://huggingface.co/models?other=axolotl&sort=created

Similar feature as: huggingface/trl#1133 that we added in TRL library

cc @winglian

[younesbelkada]

@winglian , shall we both quickly test if the tags work? I suggest we simply init a xxxTrainer and just call trainer.push_to_hub() on a small model

[winglian]

@winglian , shall we both quickly test if the tags work? I suggest we simply init a xxxTrainer and just call trainer.push_to_hub() on a small model

that would be great if you could test that to verify please.

[younesbelkada]

Thanks @winglian
I just tested it out for all axolotl trainer classes using this script for completeness using the latest commit from this branch:

from transformers import TrainingArguments, AutoModelForCausalLM

from axolotl.core.trainer_builder import AxolotlTrainer, AxolotlMambaTrainer, OneCycleLRSchedulerTrainer, ReLoRATrainer

trainer_classes = [AxolotlTrainer, AxolotlMambaTrainer, OneCycleLRSchedulerTrainer, ReLoRATrainer]

model = AutoModelForCausalLM.from_pretrained("facebook/opt-125m")

for trainer_cls in trainer_classes:
    args = TrainingArguments(
        output_dir=f"test-axolotl-{trainer_cls.__name__.lower()}"
    )

    trainer = trainer_cls(
        args=args,
        model=model
    )

    trainer.push_to_hub()

And you can find below, the links to all pushed trainers:

• AxolotlTrainer: https://huggingface.co/ybelkada/test-axolotl-axolotltrainer
• AxolotlMambaTrainer: https://huggingface.co/ybelkada/test-axolotl-axolotlmambatrainer
• OneCycleLRSchedulerTrainer: https://huggingface.co/ybelkada/test-axolotl-onecyclelrschedulertrainer
• ReLoRATrainer: https://huggingface.co/ybelkada/test-axolotl-reloratrainer

As you can see, they all have the tag axolotl, with the last 3 trainers having each an extra tag (e.g. mamba, relora, etc.)

[tmm1]

This is really great, thanks @younesbelkada !

[younesbelkada]

Thanks @winglian @tmm1 !