From 7d55607368e01387e55809612e3ea8843f7dc16b Mon Sep 17 00:00:00 2001
From: Younes Belkada <49240599+younesbelkada@users.noreply.github.com>
Date: Thu, 21 Mar 2024 16:55:56 +0100
Subject: [PATCH] HF / FEAT: Optimize HF tags (#1425) [skip ci]

* optimize tags

* chore: lint

---------

Co-authored-by: Wing Lian <wing.lian@gmail.com>
---
 src/axolotl/core/trainer_builder.py | 6 ++++++
 1 file changed, 6 insertions(+)

diff --git a/src/axolotl/core/trainer_builder.py b/src/axolotl/core/trainer_builder.py
index b23a8a124..53f6cca90 100644
--- a/src/axolotl/core/trainer_builder.py
+++ b/src/axolotl/core/trainer_builder.py
@@ -793,6 +793,12 @@ class TrainerBuilderBase(abc.ABC):
         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