From 3bc33030e27b6afe80f17a0963282bcc5bf7839a Mon Sep 17 00:00:00 2001 From: Razvan Dinu Date: Fri, 8 Dec 2023 10:40:43 +0200 Subject: [PATCH 1/4] Add NeMoGuardrails generator (WIP). --- garak/generators/guardrails.py | 40 ++++++++++++++++++++++++++++++++++ 1 file changed, 40 insertions(+) create mode 100644 garak/generators/guardrails.py diff --git a/garak/generators/guardrails.py b/garak/generators/guardrails.py new file mode 100644 index 000000000..47f236278 --- /dev/null +++ b/garak/generators/guardrails.py @@ -0,0 +1,40 @@ +"""NeMo Guardrails generator.""" + +from typing import List + +from garak.generators.base import Generator + + +class NeMoGuardrails(Generator): + """Generator wrapper for NeMo Guardrails.""" + + supports_multiple_generations = False + generator_family_name = "Guardrails" + + def __init__(self, name, generations=1): + try: + from nemoguardrails import RailsConfig, LLMRails + from nemoguardrails.logging.verbose import set_verbose + except ImportError: + raise Exception("You must first install NeMo Guardrails using `pip install nemoguardrails`.") + + self.name = name + self.fullname = f"Guardrails {self.name}" + + # Currently, we use the model_name as the path to the config + config = RailsConfig.from_path(self.name) + self.rails = LLMRails(config=config) + + super().__init__(name, generations=generations) + + def _call_model(self, prompt: str) -> List[str]: + print(f"{prompt}") + + result = self.rails.generate(prompt) + print(result) + print("--") + + return [result] + + +default_class = "NeMoGuardrails" From 26507680bad2162e41a7c063417f1419e03b8469 Mon Sep 17 00:00:00 2001 From: Leon Derczynski Date: Fri, 8 Dec 2023 10:38:00 +0100 Subject: [PATCH 2/4] tweak guardrails generator output type --- garak/generators/guardrails.py | 15 +++++++++------ 1 file changed, 9 insertions(+), 6 deletions(-) diff --git a/garak/generators/guardrails.py b/garak/generators/guardrails.py index 47f236278..4c9d1ce49 100644 --- a/garak/generators/guardrails.py +++ b/garak/generators/guardrails.py @@ -15,8 +15,10 @@ def __init__(self, name, generations=1): try: from nemoguardrails import RailsConfig, LLMRails from nemoguardrails.logging.verbose import set_verbose - except ImportError: - raise Exception("You must first install NeMo Guardrails using `pip install nemoguardrails`.") + except ImportError as e: + raise Exception( + "You must first install NeMo Guardrails using `pip install nemoguardrails`." + ) from e self.name = name self.fullname = f"Guardrails {self.name}" @@ -28,13 +30,14 @@ def __init__(self, name, generations=1): super().__init__(name, generations=generations) def _call_model(self, prompt: str) -> List[str]: - print(f"{prompt}") + # print(f"{prompt}") result = self.rails.generate(prompt) - print(result) - print("--") + # print(result) + # print(repr(result)) + # print("--") - return [result] + return result default_class = "NeMoGuardrails" From a83996983c81e2ece4c99ed1afa4868502c5bcf3 Mon Sep 17 00:00:00 2001 From: Leon Derczynski Date: Fri, 8 Dec 2023 12:45:31 +0100 Subject: [PATCH 3/4] spdx --- garak/generators/guardrails.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/garak/generators/guardrails.py b/garak/generators/guardrails.py index 4c9d1ce49..adba93fe0 100644 --- a/garak/generators/guardrails.py +++ b/garak/generators/guardrails.py @@ -1,3 +1,8 @@ +#!/usr/bin/env python3 + +# SPDX-FileCopyrightText: Copyright (c) 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + """NeMo Guardrails generator.""" from typing import List From 9a541eb51dc46ba6c5fb1dc07f669002a12775f2 Mon Sep 17 00:00:00 2001 From: Leon Derczynski Date: Tue, 12 Dec 2023 11:33:06 +0100 Subject: [PATCH 4/4] suppress sentence-transformer tqdms by eating stderr; rm debug output; narrow exception type --- garak/generators/guardrails.py | 17 ++++++++--------- 1 file changed, 8 insertions(+), 9 deletions(-) diff --git a/garak/generators/guardrails.py b/garak/generators/guardrails.py index adba93fe0..12239e460 100644 --- a/garak/generators/guardrails.py +++ b/garak/generators/guardrails.py @@ -5,6 +5,8 @@ """NeMo Guardrails generator.""" +from contextlib import redirect_stderr +import io from typing import List from garak.generators.base import Generator @@ -21,7 +23,7 @@ def __init__(self, name, generations=1): from nemoguardrails import RailsConfig, LLMRails from nemoguardrails.logging.verbose import set_verbose except ImportError as e: - raise Exception( + raise NameError( "You must first install NeMo Guardrails using `pip install nemoguardrails`." ) from e @@ -29,18 +31,15 @@ def __init__(self, name, generations=1): self.fullname = f"Guardrails {self.name}" # Currently, we use the model_name as the path to the config - config = RailsConfig.from_path(self.name) - self.rails = LLMRails(config=config) + with redirect_stderr(io.StringIO()) as f: # quieten the tqdm + config = RailsConfig.from_path(self.name) + self.rails = LLMRails(config=config) super().__init__(name, generations=generations) def _call_model(self, prompt: str) -> List[str]: - # print(f"{prompt}") - - result = self.rails.generate(prompt) - # print(result) - # print(repr(result)) - # print("--") + with redirect_stderr(io.StringIO()) as f: # quieten the tqdm + result = self.rails.generate(prompt) return result