feat: make upload tags configurable (#506)

config.default.toml

@@ -99,6 +99,13 @@ study_checkpoint_dir = "checkpoints"
 # Maximum size for individual safetensors files generated when exporting a model.
 max_shard_size = "5GB"
 
+# Tags to add to the Hugging Face repository when uploading the model.
+upload_tags = [
+    "uncensored",
+    "decensored",
+    "abliterated",
+]
+
 # System prompt to use when prompting the model.
 system_prompt = "You are a helpful assistant."
 

config.nohumor.toml

@@ -3,6 +3,10 @@
 
 max_response_length = 300
 
+upload_tags = [
+    "no-humor",
+]
+
 [scorer.KeywordRate]
 score_name = "Responses with humor"
 

config.noslop.toml

@@ -3,6 +3,10 @@
 
 max_response_length = 300
 
+upload_tags = [
+    "no-slop",
+]
+
 system_prompt = "You are a professional writer."
 
 [scorer.KeywordRate]

src/heretic/config.py

@@ -541,6 +541,15 @@ class Settings(BaseSettings):
         description='Which reproducibility information to add to the Hugging Face repository: "full", "basic", "none", or unset to prompt the user.',
     )
 
+    upload_tags: list[str] = Field(
+        default=[
+            "uncensored",
+            "decensored",
+            "abliterated",
+        ],
+        description="Tags to add to the Hugging Face repository when uploading the model.",
+    )
+
     ignore_mismatches: bool | None = Field(
         default=None,
         description="Whether to attempt to reproduce the model even if there are environment mismatches, or unset to prompt the user.",

src/heretic/main.py

@@ -1172,11 +1172,15 @@ def run():
                                     card.data = ModelCardData()
                                 if card.data.tags is None:
                                     card.data.tags = []
-                                card.data.tags.append("heretic")
-                                card.data.tags.append("uncensored")
-                                card.data.tags.append("decensored")
-                                card.data.tags.append("abliterated")
-                                if reproducibility_information != "none":
+                                if "heretic" not in card.data.tags:
+                                    card.data.tags.append("heretic")
+                                for tag in settings.upload_tags:
+                                    if tag not in card.data.tags:
+                                        card.data.tags.append(tag)
+                                if (
+                                    reproducibility_information != "none"
+                                    and "reproducible" not in card.data.tags
+                                ):
                                     card.data.tags.append("reproducible")
 
                                 # Must be a Hugging Face Hub repository ID,

tests/test_config.py

@@ -134,5 +134,30 @@ class ModelValidateInitOnlyTests(unittest.TestCase):
         self.assertEqual(settings.max_memory, {"0": "20GiB"})
 
 
+class SettingsUploadTagsTests(unittest.TestCase):
+    def test_default_upload_tags(self) -> None:
+        with patch("sys.argv", ["heretic"]):
+            settings = Settings(model="dummy")
+            self.assertEqual(
+                settings.upload_tags,
+                ["uncensored", "decensored", "abliterated"],
+            )
+
+    def test_custom_upload_tags_override(self) -> None:
+        with patch("sys.argv", ["heretic"]):
+            settings = Settings(model="dummy", upload_tags=["custom-tag", "experiment"])
+            self.assertEqual(settings.upload_tags, ["custom-tag", "experiment"])
+
+    def test_empty_upload_tags(self) -> None:
+        with patch("sys.argv", ["heretic"]):
+            settings = Settings(model="dummy", upload_tags=[])
+            self.assertEqual(settings.upload_tags, [])
+
+    def test_upload_tags_from_cli(self) -> None:
+        with patch("sys.argv", ["heretic", "--upload-tags", '["cli-tag", "test"]']):
+            settings = Settings(model="dummy")
+            self.assertEqual(settings.upload_tags, ["cli-tag", "test"])
+
+
 if __name__ == "__main__":
     unittest.main()