HF / FEAT: Optimize HF tags (#1425) [skip ci]

* optimize tags * chore: lint --------- Co-authored-by: Wing Lian <[email protected]>
axolotl-ai-cloud · Mar 21, 2024 · 7d55607 · 7d55607
1 parent 7803f09
commit 7d55607
Showing 1 changed file with 6 additions and 0 deletions.
diff --git a/src/axolotl/core/trainer_builder.py b/src/axolotl/core/trainer_builder.py
@@ -793,6 +793,12 @@ def __init__(self, cfg, model, tokenizer):
         self.model = model
         self.tokenizer = tokenizer
 
+        # in case the model supports tagging, add the axolotl tag.
+        # This makes sure the tag is correctly pushed even if a user calls
+        # model.push_to_hub instad of  trainer.push_to_hub.
+        if hasattr(model, "add_model_tags"):
+            model.add_model_tags(["axolotl"])
+
     @property
     def model_ref(self):
         return self._model_ref