Repository: ThuCCSLab/JailbreakEval

Issue:
[BUG] The performance of LlamaGuard evaluator is unreasonably low
**Describe the bug**
When using LlamaGuard 1/2 evaluator, the result is unreasonably low.

**To Reproduce**
Steps to reproduce the behavior:
```shell
JailbreakEval --config config.yaml --dataset data/example.csv --output result.csv HFChat-inan2023llama-llamaguard HFChat-inan2023llama-llamaguard2
```
And the results are
```
Dataset: data/example.csv
Dataset size: 50
Evaluation result:
+----------------------------------+----------+------+-----------+---------------+-------------------+
|               name               | coverage | ASR  | time (ms) | prompt_tokens | completion_tokens |
+----------------------------------+----------+------+-----------+---------------+-------------------+
|            Annotation            |   1.00   | 0.62 |    N/A    |      N/A      |        N/A        |
| HFChat-inan2023llama-llamaguard  |   1.00   | 0.00 |   63542   |      N/A      |        N/A        |
| HFChat-inan2023llama-llamaguard2 |   1.00   | 0.00 |   36333   |      N/A      |        N/A        |
+----------------------------------+----------+------+-----------+---------------+-------------------+
Evaluation agreement with annotation:
+----------------------------------+----------+----------+--------+-----------+------+
|               name               | coverage | accuracy | recall | precision |  f1  |
+----------------------------------+----------+----------+--------+-----------+------+
| HFChat-inan2023llama-llamaguard  |   1.00   |   0.38   |  0.00  |    0.00   | 0.00 |
| HFChat-inan2023llama-llamaguard2 |   1.00   |   0.38   |  0.00  |    0.00   | 0.00 |
+----------------------------------+----------+----------+--------+-----------+------+
```
which is inconsistent with the paper.

**Expected behavior**
The LlamaGuard 1/2 Evaluator should have a good performance, as described in the paper.

**System (please complete the following information):**
 - OS: Ubuntu
 - Version: 22.04

**Additional context**
Add any other context about the problem here.

Discussion:

Solution patch:
diff --git a/jailbreakeval/services/chat/service_base.py b/jailbreakeval/services/chat/service_base.py
index efe4255..276cab4 100644
--- a/jailbreakeval/services/chat/service_base.py
+++ b/jailbreakeval/services/chat/service_base.py
@@ -13,7 +13,7 @@ class ChatService(ABC):
     def __init__(self, model: Optional[str] = None, chat_kwargs: Optional[Dict[str, Any]] = None, **kwargs) -> None:
         super().__init__()
         self.model = model
-        self.chat_kwargs = deepcopy(chat_kwargs) or {}
+        self.chat_kwargs = deepcopy(chat_kwargs or {})
 
     @classmethod
     def __init_subclass__(cls, service_type: Optional[str] = None, **kwargs):
diff --git a/jailbreakeval/services/chat/transformers_pipeline_service.py b/jailbreakeval/services/chat/transformers_pipeline_service.py
index b5da946..6eeb097 100644
--- a/jailbreakeval/services/chat/transformers_pipeline_service.py
+++ b/jailbreakeval/services/chat/transformers_pipeline_service.py
@@ -19,6 +19,8 @@ class TransformersPipelineChatService(ChatService, service_type="transformers-pi
     ) -> None:
         self.pipeline = pipeline or pipeline_factory(task="text-generation", model=model, **pipeline_kwargs)
         super().__init__(self.pipeline.model.name_or_path, chat_kwargs)
+        self.chat_kwargs["return_full_text"] = False
+        self.chat_kwargs["continue_final_message"] = False
 
     def chat(self, conversation: Sequence[Message], **override_chat_kwargs) -> str:
         chat_kwargs = copy.deepcopy(self.chat_kwargs)
@@ -26,4 +28,4 @@ class TransformersPipelineChatService(ChatService, service_type="transformers-pi
             if k in chat_kwargs:
                 chat_kwargs.pop(k)
         result = self.pipeline(text_inputs=conversation, **chat_kwargs, **override_chat_kwargs)
-        return result[0]["generated_text"][-1]["content"]
+        return result[0]["generated_text"]
diff --git a/jailbreakeval/services/text_classification/service_base.py b/jailbreakeval/services/text_classification/service_base.py
index a4f0256..1976285 100644
--- a/jailbreakeval/services/text_classification/service_base.py
+++ b/jailbreakeval/services/text_classification/service_base.py
@@ -11,7 +11,7 @@ class TextClassificationService(ABC):
     def __init__(self, model: Optional[str] = None, predict_kwargs: Optional[Dict[str, Any]] = None, **kwargs) -> None:
         super().__init__()
         self.model = model
-        self.predict_kwargs = deepcopy(predict_kwargs) or {}
+        self.predict_kwargs = deepcopy(predict_kwargs or {})
 
     @classmethod
     def __init_subclass__(cls, service_type: Optional[str] = None, **kwargs) -> None:
diff --git a/pyproject.toml b/pyproject.toml
index e1653f8..890551e 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -26,7 +26,7 @@ dependencies = [
     "numpy",
     "tqdm",
     "jinja2",
-    "transformers>=4.38.0",
+    "transformers>=4.45.0",
     "torch",
     "accelerate",
     "sentencepiece",


Test patch:
diff --git a/tests/evaluators/test_chat_evaluator.py b/tests/evaluators/test_chat_evaluator.py
index 4ba5543..bb1662f 100644
--- a/tests/evaluators/test_chat_evaluator.py
+++ b/tests/evaluators/test_chat_evaluator.py
@@ -31,6 +31,8 @@ def test_transformers_pipeline(mock_transformers_pipeline_call, mock_chat_servic
         ],
         temperature=0.0,
         max_new_tokens=20,
+        return_full_text=False,
+        continue_final_message=False,
     )
     assert not result
 
diff --git a/tests/services/chat/test_transformers_pipeline_service.py b/tests/services/chat/test_transformers_pipeline_service.py
index ad704f6..08ab905 100644
--- a/tests/services/chat/test_transformers_pipeline_service.py
+++ b/tests/services/chat/test_transformers_pipeline_service.py
@@ -42,13 +42,17 @@ def test_predict(mock_transformers_pipeline_call):
 
     service = TransformersPipelineChatService(model=MODEL)
     response = service.chat(MESSAGES)
-    mock_transformers_pipeline_call.assert_called_with(text_inputs=MESSAGES)
+    mock_transformers_pipeline_call.assert_called_with(
+        text_inputs=MESSAGES, return_full_text=False, continue_final_message=False
+    )
     assert isinstance(response, str)
     assert response == RESPONSE
 
     # override
     response = service.chat(MESSAGES, key="value")
-    mock_transformers_pipeline_call.assert_called_with(text_inputs=MESSAGES, key="value")
+    mock_transformers_pipeline_call.assert_called_with(
+        text_inputs=MESSAGES, key="value", return_full_text=False, continue_final_message=False
+    )
     assert isinstance(response, str)
     assert response == RESPONSE
 
@@ -56,13 +60,22 @@ def test_predict(mock_transformers_pipeline_call):
     service = TransformersPipelineChatService(model=MODEL, chat_kwargs={"hello": "world", "foo": "bar"})
     response = service.chat(MESSAGES)
 
-    mock_transformers_pipeline_call.assert_called_with(text_inputs=MESSAGES, hello="world", foo="bar")
+    mock_transformers_pipeline_call.assert_called_with(
+        text_inputs=MESSAGES, hello="world", foo="bar", return_full_text=False, continue_final_message=False
+    )
     assert isinstance(response, str)
     assert response == RESPONSE
 
     # gen_kwargs + override
-    response = service.chat(MESSAGES, foo="baz", key="value")
+    response = service.chat(MESSAGES, foo="baz", key="value", return_full_text=False, continue_final_message=False)
 
-    mock_transformers_pipeline_call.assert_called_with(text_inputs=MESSAGES, hello="world", foo="baz", key="value")
+    mock_transformers_pipeline_call.assert_called_with(
+        text_inputs=MESSAGES,
+        hello="world",
+        foo="baz",
+        key="value",
+        return_full_text=False,
+        continue_final_message=False,
+    )
     assert isinstance(response, str)
     assert response == RESPONSE
diff --git a/tests/testing_utils.py b/tests/testing_utils.py
index 48cac7b..2d53419 100644
--- a/tests/testing_utils.py
+++ b/tests/testing_utils.py
@@ -99,7 +99,7 @@ def transformers_chat_pipeline_response(
         response = [response]
     chat_id = str(uuid.uuid4())
     for i, response in enumerate(response):
-        yield [{"mock_id": f"mock-{chat_id}-{i}", "generated_text": [{"role": "assistance", "content": response}]}]
+        yield [{"mock_id": f"mock-{chat_id}-{i}", "generated_text": response}]
 
 
 # %%