Configurable model list for stormtrooper processor

digitalmethodsinitiative · Oct 11, 2024 · d399921 · d399921
1 parent db5b649
commit d399921
Showing 1 changed file with 9 additions and 2 deletions.
diff --git a/processors/machine_learning/annotate_text.py b/processors/machine_learning/annotate_text.py
@@ -40,6 +40,11 @@ class TextClassifier(BasicProcessor):
             "type": UserInput.OPTION_TOGGLE,
             "default": False,
             "help": "Enable LLM-powered text classification",
+        },
+        "dmi-service-manager.stormtrooper_models": {
+            "type": UserInput.OPTION_TEXT,
+            "default": "google/flan-t5-large,tiiaue/falcon-7b-instruct",
+            "help": "Comma-separated list of models that can be selected"
         }
     }
 
@@ -53,8 +58,6 @@ class TextClassifier(BasicProcessor):
             "type": UserInput.OPTION_CHOICE,
             "default": "google/flan-t5-large",
             "options": {
-                "google/flan-t5-large": "google/flan-t5-large",
-                "tiiaue/falcon-7b-instruct": "tiiaue/falcon-7b-instruct"
             },
             "help": "Large Language Model to use"
         },
@@ -97,6 +100,10 @@ def get_options(cls, parent_dataset=None, user=None):
         :return dict:  Processor options
         """
         options = cls.options
+
+        models = config.get("dmi-service-manager.stormtrooper_models", user=user).split(",")
+        options["model"]["options"] = {m: m for m in models}
+
         if parent_dataset is None:
             return options