Spaces:
Sleeping
Sleeping
File size: 2,694 Bytes
acb544e |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 |
from typing import Dict, List, Optional
import structlog
from llm_guard import input_scanners, output_scanners
from llm_guard.input_scanners.base import Scanner as InputScanner
from llm_guard.output_scanners.base import Scanner as OutputScanner
from llm_guard.vault import Vault
from .config import ScannerConfig
from .util import get_resource_utilization
LOGGER = structlog.getLogger(__name__)
def get_input_scanners(scanners: List[ScannerConfig], vault: Vault) -> List[InputScanner]:
"""
Load input scanners from the configuration file.
"""
input_scanners_loaded = []
for scanner in scanners:
LOGGER.debug("Loading input scanner", scanner=scanner.type, **get_resource_utilization())
input_scanners_loaded.append(
_get_input_scanner(
scanner.type,
scanner.params,
vault=vault,
)
)
return input_scanners_loaded
def get_output_scanners(scanners: List[ScannerConfig], vault: Vault) -> List[OutputScanner]:
"""
Load output scanners from the configuration file.
"""
output_scanners_loaded = []
for scanner in scanners:
LOGGER.debug("Loading output scanner", scanner=scanner.type, **get_resource_utilization())
output_scanners_loaded.append(
_get_output_scanner(
scanner.type,
scanner.params,
vault=vault,
)
)
return output_scanners_loaded
def _get_input_scanner(
scanner_name: str,
scanner_config: Optional[Dict],
*,
vault: Vault,
):
if scanner_config is None:
scanner_config = {}
if scanner_name == "Anonymize":
scanner_config["vault"] = vault
if scanner_name in [
"Anonymize",
"BanTopics",
"Code",
"Gibberish",
"Language",
"PromptInjection",
"Toxicity",
]:
scanner_config["use_onnx"] = True
return input_scanners.get_scanner_by_name(scanner_name, scanner_config)
def _get_output_scanner(
scanner_name: str,
scanner_config: Optional[Dict],
*,
vault: Vault,
):
if scanner_config is None:
scanner_config = {}
if scanner_name == "Deanonymize":
scanner_config["vault"] = vault
if scanner_name in [
"BanTopics",
"Bias",
"Code",
"FactualConsistency",
"Gibberish",
"Language",
"LanguageSame",
"MaliciousURLs",
"NoRefusal",
"Relevance",
"Sensitive",
"Toxicity",
]:
scanner_config["use_onnx"] = True
return output_scanners.get_scanner_by_name(scanner_name, scanner_config)
|