Younes Belkada winglian commited on
Commit
7d55607
·
unverified ·
1 Parent(s): 7803f09

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

Browse files

* optimize tags

* chore: lint

---------

Co-authored-by: Wing Lian <[email protected]>

src/axolotl/core/trainer_builder.py CHANGED
@@ -793,6 +793,12 @@ class TrainerBuilderBase(abc.ABC):
793
  self.model = model
794
  self.tokenizer = tokenizer
795
 
 
 
 
 
 
 
796
  @property
797
  def model_ref(self):
798
  return self._model_ref
 
793
  self.model = model
794
  self.tokenizer = tokenizer
795
 
796
+ # in case the model supports tagging, add the axolotl tag.
797
+ # This makes sure the tag is correctly pushed even if a user calls
798
+ # model.push_to_hub instad of trainer.push_to_hub.
799
+ if hasattr(model, "add_model_tags"):
800
+ model.add_model_tags(["axolotl"])
801
+
802
  @property
803
  def model_ref(self):
804
  return self._model_ref