Merge pull request #352 from NVIDIA/fix/qa-fixes-4

Fix a few issues related to jailbreak detection heuristics.
NVIDIA · Feb 23, 2024 · 0bd4257 · 0bd4257
2 parents 65f0e9e + 88a47a6
commit 0bd4257
Show file tree

Hide file tree

Showing 4 changed files with 12 additions and 5 deletions.
diff --git a/nemoguardrails/embeddings/embedding_providers/fastembed.py b/nemoguardrails/embeddings/embedding_providers/fastembed.py
@@ -37,7 +37,7 @@ class FastEmbedEmbeddingModel(EmbeddingModel):
     """
 
     def __init__(self, embedding_model: str):
-        from fastembed.embedding import FlagEmbedding as Embedding
+        from fastembed import TextEmbedding as Embedding
 
         # Enabling a short form model name for all-MiniLM-L6-v2.
         if embedding_model == "all-MiniLM-L6-v2":

diff --git a/nemoguardrails/library/jailbreak_detection/Dockerfile b/nemoguardrails/library/jailbreak_detection/Dockerfile
@@ -19,6 +19,9 @@ RUN pip install -r requirements.txt
 # Set the device on which the model should load e.g., "cpu", "cuda:0", etc.
 ENV JAILBREAK_CHECK_DEVICE=cpu
 
+# Predownload the GPT2 model.
+RUN python -c "from transformers import GPT2LMHeadModel, GPT2TokenizerFast; GPT2LMHeadModel.from_pretrained('gpt2-large'); GPT2TokenizerFast.from_pretrained('gpt2-large');"
+
 # Expose a port for the server
 EXPOSE 1337
 

diff --git a/nemoguardrails/library/jailbreak_detection/Dockerfile-GPU b/nemoguardrails/library/jailbreak_detection/Dockerfile-GPU
@@ -22,6 +22,9 @@ RUN pip install -r requirements.txt
 # Set the device on which the model should load e.g., "cpu", "cuda:0", etc.
 ENV JAILBREAK_CHECK_DEVICE=cuda:0
 
+# Predownload the GPT2 model.
+RUN python -c "from transformers import GPT2LMHeadModel, GPT2TokenizerFast; GPT2LMHeadModel.from_pretrained('gpt2-large'); GPT2TokenizerFast.from_pretrained('gpt2-large');"
+
 # Expose a port for the server
 EXPOSE 1337
 

diff --git a/nemoguardrails/library/jailbreak_detection/actions.py b/nemoguardrails/library/jailbreak_detection/actions.py
@@ -17,10 +17,6 @@
 from typing import Optional
 
 from nemoguardrails.actions import action
-from nemoguardrails.library.jailbreak_detection.heuristics.checks import (
-    check_jailbreak_length_per_perplexity,
-    check_jailbreak_prefix_suffix_perplexity,
-)
 from nemoguardrails.library.jailbreak_detection.request import (
     jailbreak_detection_heuristics_request,
 )
@@ -43,6 +39,11 @@ async def jailbreak_detection_heuristics(
     prompt = context.get("user_message")
 
     if not jailbreak_api_url:
+        from nemoguardrails.library.jailbreak_detection.heuristics.checks import (
+            check_jailbreak_length_per_perplexity,
+            check_jailbreak_prefix_suffix_perplexity,
+        )
+
         log.warning(
             "No jailbreak heuristics endpoint set. Running in-process, NOT RECOMMENDED FOR PRODUCTION."
         )