From d399921b7538bc06264b122c8a0d523826bb91a5 Mon Sep 17 00:00:00 2001 From: Stijn Peeters Date: Fri, 11 Oct 2024 17:12:21 +0200 Subject: [PATCH] Configurable model list for stormtrooper processor --- processors/machine_learning/annotate_text.py | 11 +++++++++-- 1 file changed, 9 insertions(+), 2 deletions(-) diff --git a/processors/machine_learning/annotate_text.py b/processors/machine_learning/annotate_text.py index 954963de4..59016e077 100644 --- 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