feature: deferred loading and requirement pruning

garak/_plugins.py

@@ -400,6 +400,22 @@ def load_plugin(path, break_on_fail=True, config_root=_config) -> object:
             ) from ve
         else:
             return False
+
+    full_plugin_name = ".".join((category, module_name, plugin_class_name))
+
+    # check cache for optional imports
+    extra_dependency_names = PluginCache.instance()[category][full_plugin_name][
+        "extra_dependency_names"
+    ]
+    if len(extra_dependency_names) > 0:
+        for dependency_module_name in extra_dependency_names:
+            for dependency_path in [ # support both plain names and also multi-point names e.g. langchain.llms
+                ".".join(dependency_module_name.split(".")[: n + 1])
+                for n in range(dependency_module_name.count(".") + 1)
+            ]:
+                if importlib.util.find_spec(dependency_path) is None:
+                    _import_failed(dependency_path, full_plugin_name)
+
     module_path = f"garak.{category}.{module_name}"
     try:
         mod = importlib.import_module(module_path)
@@ -426,6 +442,7 @@ def load_plugin(path, break_on_fail=True, config_root=_config) -> object:
         if plugin_instance is None:
             plugin_instance = klass(config_root=config_root)
             PluginProvider.storeInstance(plugin_instance, config_root)
+
     except Exception as e:
         logging.warning(
             "Exception instantiating %s.%s: %s",
@@ -440,3 +457,20 @@ def load_plugin(path, break_on_fail=True, config_root=_config) -> object:
             return False
 
     return plugin_instance
+
+
+def load_optional_module(module_name: str):
+    try:
+        m = importlib.import_module(module_name)
+    except ModuleNotFoundError:
+        requesting_module = Path(inspect.stack()[1].filename).name.replace(".py", "")
+        _import_failed(module_name, requesting_module)
+    return m
+
+
+def _import_failed(import_module: str, calling_module: str):
+    msg = f"⛔ Plugin '{calling_module}' requires Python module '{import_module}' but this isn't installed/available."
+    hint = f"💡 Try 'pip install {import_module}' to get it."
+    logging.critical(msg)
+    print(msg + "\n" + hint)
+    raise ModuleNotFoundError(msg)

garak/buffs/base.py

@@ -27,6 +27,8 @@ class Buff(Configurable):
     doc_uri = ""
     lang = None  # set of languages this buff should be constrained to
     active = True
+    # list of strings naming modules required but not explicitly in garak by default
+    extra_dependency_names = []
 
     DEFAULT_PARAMS = {}
 

garak/detectors/base.py

@@ -27,6 +27,8 @@ class Detector(Configurable):
     accuracy = None
     active = True
     tags = []  # list of taxonomy categories per the MISP format
+    # list of strings naming modules required but not explicitly in garak by default
+    extra_dependency_names = []
 
     # support mainstream any-to-any large models
     # legal element for str list `modality['in']`: 'text', 'image', 'audio', 'video', '3d'

garak/exception.py

@@ -36,3 +36,7 @@ class ConfigFailure(GarakException):
 
 class PayloadFailure(GarakException):
     """Problem instantiating/using payloads"""
+
+
+class GeneratorBackoffExceptionPlaceholder(GarakException):
+    """Placeholder used for lazy-loaded exceptions"""

garak/generators/azure.py

@@ -82,6 +82,7 @@ def _validate_env_var(self):
         return super()._validate_env_var()
 
     def _load_client(self):
+        self._load_deps()
         if self.model_name in openai_model_mapping:
             self.model_name = openai_model_mapping[self.model_name]
 

garak/generators/base.py

@@ -44,6 +44,8 @@ class Generator(Configurable):
     supports_multiple_generations = (
         False  # can more than one generation be extracted per request?
     )
+    # list of strings naming modules required but not explicitly in garak by default
+    extra_dependency_names = []
 
     def __init__(self, name="", config_root=_config):
         self._load_config(config_root)
@@ -63,6 +65,29 @@ def __init__(self, name="", config_root=_config):
             f"🦜 loading {Style.BRIGHT}{Fore.LIGHTMAGENTA_EX}generator{Style.RESET_ALL}: {self.generator_family_name}: {self.name}"
         )
         logging.info("generator init: %s", self)
+        self._load_deps()
+
+    def _load_deps(self):
+        # load external dependencies. should be invoked at construction and
+        # in _client_load (if used)
+        for extra_dependency in self.extra_dependency_names:
+            extra_dep_name = extra_dependency.replace(".", "_").replace("-", "_")
+            if (
+                not hasattr(self, extra_dep_name)
+                or getattr(self, extra_dep_name) is None
+            ):
+                setattr(
+                    self,
+                    extra_dep_name,
+                    garak._plugins.load_optional_module(extra_dependency),
+                )
+
+    def _clear_deps(self):
+        # unload external dependencies from class. should be invoked before
+        # serialisation, esp. in _clear_client (if used)
+        for extra_dependency in self.extra_dependency_names:
+            extra_dep_name = extra_dependency.replace(".", "_")
+            setattr(self, extra_dep_name, None)
 
     def _call_model(
         self, prompt: str, generations_this_call: int = 1
@@ -101,7 +126,7 @@ def _prune_skip_sequences(self, outputs: List[str | None]) -> List[str | None]:
         )
         rx_missing_final = re.escape(self.skip_seq_start) + ".*?$"
         rx_missing_start = ".*?" + re.escape(self.skip_seq_end)
-        
+
         if self.skip_seq_start == "":
             complete_seqs_removed = [
                 (

garak/generators/cohere.py

@@ -10,17 +10,20 @@
 from typing import List, Union
 
 import backoff
-import cohere
 import tqdm
 
 from garak import _config
+import garak._plugins
+from garak.exception import GeneratorBackoffExceptionPlaceholder
 from garak.generators.base import Generator
 
 
 COHERE_GENERATION_LIMIT = (
     5  # c.f. https://docs.cohere.com/reference/generate 18 may 2023
 )
 
+cohere_exception = GeneratorBackoffExceptionPlaceholder
+
 
 class CohereGenerator(Generator):
     """Interface to Cohere's python library for their text2text model.
@@ -38,22 +41,27 @@ class CohereGenerator(Generator):
         "presence_penalty": 0.0,
         "stop": [],
     }
+    extra_dependency_names = ["cohere"]
 
     supports_multiple_generations = True
     generator_family_name = "Cohere"
 
     def __init__(self, name="command", config_root=_config):
+
         self.name = name
         self.fullname = f"Cohere {self.name}"
 
         super().__init__(self.name, config_root=config_root)
 
+        global cohere_exception
+        cohere_exception = self.cohere.error.CohereAPIError
+
         logging.debug(
             "Cohere generation request limit capped at %s", COHERE_GENERATION_LIMIT
         )
-        self.generator = cohere.Client(self.api_key)
+        self.generator = self.cohere.Client(self.api_key)
 
-    @backoff.on_exception(backoff.fibo, cohere.error.CohereAPIError, max_value=70)
+    @backoff.on_exception(backoff.fibo, cohere_exception, max_value=70)
     def _call_cohere_api(self, prompt, request_size=COHERE_GENERATION_LIMIT):
         """as of jun 2 2023, empty prompts raise:
         cohere.error.CohereAPIError: invalid request: prompt must be at least 1 token long

garak/generators/groq.py

@@ -39,6 +39,7 @@ class GroqChat(OpenAICompatible):
     generator_family_name = "Groq"
 
     def _load_client(self):
+        self._load_deps()
         self.client = openai.OpenAI(base_url=self.uri, api_key=self.api_key)
         if self.name in ("", None):
             raise ValueError(

garak/generators/guardrails.py

@@ -16,27 +16,21 @@ class NeMoGuardrails(Generator):
 
     supports_multiple_generations = False
     generator_family_name = "Guardrails"
+    extra_dependency_names = ["nemoguardrails"]
 
     def __init__(self, name="", config_root=_config):
-        # another class that may need to skip testing due to non required dependency
-        try:
-            from nemoguardrails import RailsConfig, LLMRails
-            from nemoguardrails.logging.verbose import set_verbose
-        except ImportError as e:
-            raise NameError(
-                "You must first install NeMo Guardrails using `pip install nemoguardrails`."
-            ) from e
 
         self.name = name
         self._load_config(config_root)
         self.fullname = f"Guardrails {self.name}"
 
+        super().__init__(self.name, config_root=config_root)
+
+        set_verbose = self.nemoguardrails.logging.verbose.set_verbose
         # Currently, we use the model_name as the path to the config
         with redirect_stderr(io.StringIO()) as f:  # quieten the tqdm
-            config = RailsConfig.from_path(self.name)
-            self.rails = LLMRails(config=config)
-
-        super().__init__(self.name, config_root=config_root)
+            config = self.nemoguardrails.RailsConfig.from_path(self.name)
+            self.rails = self.nemoguardrails.LLMRails(config=config)
 
     def _call_model(
         self, prompt: str, generations_this_call: int = 1

garak/generators/huggingface.py

@@ -21,7 +21,6 @@
 
 import backoff
 import torch
-from PIL import Image
 
 from garak import _config
 from garak.exception import ModelNameMissingError, GarakException
@@ -70,6 +69,7 @@ def __init__(self, name="", config_root=_config):
         self._load_client()
 
     def _load_client(self):
+        self._load_deps()
         if hasattr(self, "generator") and self.generator is not None:
             return
 
@@ -104,6 +104,7 @@ def _load_client(self):
         self._set_hf_context_len(self.generator.model.config)
 
     def _clear_client(self):
+        self._clear_deps()
         self.generator = None
 
     def _format_chat_prompt(self, prompt: str) -> List[dict]:
@@ -158,19 +159,15 @@ class OptimumPipeline(Pipeline, HFCompatible):
     generator_family_name = "NVIDIA Optimum Hugging Face 🤗 pipeline"
     supports_multiple_generations = True
     doc_uri = "https://huggingface.co/blog/optimum-nvidia"
+    extra_dependency_names = ["optimum-nvidia"]
 
     def _load_client(self):
+        self._load_deps()
         if hasattr(self, "generator") and self.generator is not None:
             return
 
-        try:
-            from optimum.nvidia.pipelines import pipeline
-            from transformers import set_seed
-        except Exception as e:
-            logging.exception(e)
-            raise GarakException(
-                f"Missing required dependencies for {self.__class__.__name__}"
-            )
+        pipeline = self.optimum.nvidia.pipelines.pipeline
+        from transformers import set_seed
 
         if self.seed is not None:
             set_seed(self.seed)
@@ -205,6 +202,7 @@ class ConversationalPipeline(Pipeline, HFCompatible):
     supports_multiple_generations = True
 
     def _load_client(self):
+        self._load_deps()
         if hasattr(self, "generator") and self.generator is not None:
             return
 
@@ -454,6 +452,7 @@ class Model(Pipeline, HFCompatible):
     supports_multiple_generations = True
 
     def _load_client(self):
+        self._load_deps()
         if hasattr(self, "model") and self.model is not None:
             return
 
@@ -501,6 +500,7 @@ def _load_client(self):
         self.generation_config.pad_token_id = self.model.config.eos_token_id
 
     def _clear_client(self):
+        self._clear_deps()
         self.model = None
         self.config = None
         self.tokenizer = None
@@ -575,6 +575,8 @@ class LLaVA(Generator, HFCompatible):
     NB. This should be use with strict modality matching - generate() doesn't
     support text-only prompts."""
 
+    extra_dependency_names = ["PIL"]
+
     DEFAULT_PARAMS = Generator.DEFAULT_PARAMS | {
         "max_tokens": 4000,
         # "exist_tokens + max_new_tokens < 4K is the golden rule."
@@ -626,7 +628,7 @@ def generate(
 
         text_prompt = prompt["text"]
         try:
-            image_prompt = Image.open(prompt["image"])
+            image_prompt = self.PIL.Image.open(prompt["image"])
         except FileNotFoundError:
             raise FileNotFoundError(f"Cannot open image {prompt['image']}.")
         except Exception as e:

garak/generators/langchain.py

@@ -8,9 +8,6 @@
 import logging
 from typing import List, Union
 
-
-import langchain.llms
-
 from garak import _config
 from garak.generators.base import Generator
 
@@ -43,7 +40,7 @@ class LangChainLLMGenerator(Generator):
         "presence_penalty": 0.0,
         "stop": [],
     }
-
+    extra_dependency_names = ["langchain.llms"]
     generator_family_name = "LangChain"
 
     def __init__(self, name="", config_root=_config):
@@ -53,14 +50,7 @@ def __init__(self, name="", config_root=_config):
 
         super().__init__(self.name, config_root=config_root)
 
-        try:
-            # this might need some special handling to allow tests
-            llm = getattr(langchain.llms, self.name)()
-        except Exception as e:
-            logging.error("Failed to import Langchain module: %s", repr(e))
-            raise e
-
-        self.generator = llm
+        self.generator = getattr(self.langchain_llms, self.name)()
 
     def _call_model(
         self, prompt: str, generations_this_call: int = 1