Repository: deepset-ai/haystack

Issue:
`HuggingFaceLocalGenerator` - test stopwords handling in a e2e test
Related to #6049 (stopping text generation when a stopword is encountered)

The existing unit tests are not fully comprehensive.
It's difficult to test all the mechanisms for stopping text generation without performing inference.
(On the other hand, I'm noticing that we keep adding integration tests requiring inference, at the risk of making the CI slow again.)

**So, we would like to test also stopwords handling within an e2e test.**



Discussion:
The mechanism that stops the word generation is inside the Hugginface pipeline. We create a structure defining the stop tokens and criteria, `StoppingCriteriaList`, pass it to the Huggingface pipeline, and it's the pipeline's responsibility to use it to stop the generation.

here, we use the `stop_words` to create the  `StoppingCriteriaList` structure (see line [146](https://github.com/deepset-ai/haystack/blob/main/haystack/components/generators/hugging_face_local.py#L146))

```python
if self.stop_words and self.stopping_criteria_list is None:
    stop_words_criteria = StopWordsCriteria(tokenizer=self.pipeline.tokenizer, stop_words=self.stop_words, device=self.pipeline.device)
    self.stopping_criteria_list = StoppingCriteriaList([stop_words_criteria])
```

and here, we pass it to the pipeline (see line [232](https://github.com/deepset-ai/haystack/blob/main/haystack/components/generators/hugging_face_local.py#L232))

```pyhon
output = self.pipeline(prompt, stopping_criteria=self.stopping_criteria_list, **updated_generation_kwargs)
```


I would say enough tests are:
- make sure the `StopWordsCriteria`structure creation  works 
- make sure that passing it to the HF pipeline works, i.e. the pipeline runs and outputs something, doesn't raise any Exception because of the structure we pass
- test the `class StopWordsCriteria(StoppingCriteria)` (which it's already being done I can review it)

We can't control and shouldn't test whether the generation is stopped, it's on Hugginface's side using `StopWordsCriteria(StoppingCriteria)` within the model/pipeline to stop the generation. 

If we test for that, we would be testing Hugginface's code, and also, due to the non-deterministic behaviour of the text generation, this might be a hard test to assert.

Resolving patch:
diff --git a/haystack/components/generators/hugging_face_local.py b/haystack/components/generators/hugging_face_local.py
--- a/haystack/components/generators/hugging_face_local.py
+++ b/haystack/components/generators/hugging_face_local.py
@@ -143,7 +143,7 @@ def warm_up(self):
         if self.pipeline is None:
             self.pipeline = pipeline(**self.huggingface_pipeline_kwargs)
 
-        if self.stop_words and self.stopping_criteria_list is None:
+        if self.stop_words:
             stop_words_criteria = StopWordsCriteria(
                 tokenizer=self.pipeline.tokenizer, stop_words=self.stop_words, device=self.pipeline.device
             )
diff --git a/haystack/utils/hf.py b/haystack/utils/hf.py
--- a/haystack/utils/hf.py
+++ b/haystack/utils/hf.py
@@ -265,7 +265,12 @@ def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor, **kwa
                     return True
             return False
 
-        def is_stop_word_found(self, generated_text_ids: torch.Tensor, stop_id: torch.Tensor) -> bool:
+        @staticmethod
+        def is_stop_word_found(generated_text_ids: torch.Tensor, stop_id: torch.Tensor) -> bool:
+            """
+            Performs phrase matching, checking if a sequence of stop tokens appears in a continuous
+            or sequential order within the generated text.
+            """
             generated_text_ids = generated_text_ids[-1]
             len_generated_text_ids = generated_text_ids.size(0)
             len_stop_id = stop_id.size(0)


Test patch:
diff --git a/test/components/generators/test_hugging_face_local_generator.py b/test/components/generators/test_hugging_face_local_generator.py
--- a/test/components/generators/test_hugging_face_local_generator.py
+++ b/test/components/generators/test_hugging_face_local_generator.py
@@ -301,7 +301,7 @@ def test_warm_up(self, pipeline_mock):
         )
 
     @patch("haystack.components.generators.hugging_face_local.pipeline")
-    def test_warm_up_doesn_reload(self, pipeline_mock):
+    def test_warm_up_doesnt_reload(self, pipeline_mock):
         generator = HuggingFaceLocalGenerator(
             model="google/flan-t5-base", task="text2text-generation", token=Secret.from_token("fake-api-token")
         )
@@ -385,41 +385,18 @@ def test_run_fails_without_warm_up(self):
         with pytest.raises(RuntimeError, match="The generation model has not been loaded."):
             generator.run(prompt="irrelevant")
 
-    def test_stop_words_criteria(self):
+    def test_stop_words_criteria_with_a_mocked_tokenizer(self):
         """
-        Test that StopWordsCriteria will check stop word tokens in a continuous and sequential order
+        Test that StopWordsCriteria will caught stop word tokens in a continuous and sequential order in the input_ids
         """
-        # input ids for "unambiguously"
-        stop_words_id = torch.tensor([[73, 24621, 11937]])
-
-        # input ids for "This is ambiguously, but is unrelated."
-        input_ids1 = torch.tensor([[100, 19, 24621, 11937, 6, 68, 19, 73, 3897, 5]])
-        # input ids for "This is unambiguously"
-        input_ids2 = torch.tensor([[100, 19, 73, 24621, 11937]])
-
-        # We used to implement stop words algorithm using the torch.isin function like this:
-        # `all(torch.isin(stop_words_id, input_ids1)[0])`
-        # However, this algorithm is not correct as it will return True for presence of "unambiguously" in input_ids1
-        # and True for presence of "unambiguously" in input_ids2. This is because the algorithm will check
-        # if the stop word tokens are present in the input_ids, but it does not check if the stop word tokens are
-        # present in a continuous/sequential order.
-
-        # In "This is ambiguously, but is unrelated." sentence the "un" token comes from "unrelated" and the
-        # "ambiguously" token comes from "ambiguously". The algorithm will return True for presence of
-        # "unambiguously" in input_ids1 which is not correct.
-
+        stop_words_id = torch.LongTensor([[73, 24621, 11937]])  # "unambiguously"
+        # "This is ambiguously, but is unrelated."
+        input_ids_one = torch.LongTensor([[100, 19, 24621, 11937, 6, 68, 19, 73, 3897, 5]])
+        input_ids_two = torch.LongTensor([[100, 19, 73, 24621, 11937]])  # "This is unambiguously"
         stop_words_criteria = StopWordsCriteria(tokenizer=Mock(spec=PreTrainedTokenizerFast), stop_words=["mock data"])
-        # because we are mocking the tokenizer, we need to set the stop words manually
         stop_words_criteria.stop_ids = stop_words_id
-
-        # this is the correct algorithm to check if the stop word tokens are present in a continuous and sequential order
-        # For the input_ids1, the stop word tokens are present BUT not in a continuous order
-        present_and_continuous = stop_words_criteria(input_ids1, scores=None)
-        assert not present_and_continuous
-
-        # For the input_ids2, the stop word tokens are both present and in a continuous order
-        present_and_continuous = stop_words_criteria(input_ids2, scores=None)
-        assert present_and_continuous
+        assert not stop_words_criteria(input_ids_one, scores=None)
+        assert stop_words_criteria(input_ids_two, scores=None)
 
     @patch("haystack.components.generators.hugging_face_local.pipeline")
     @patch("haystack.components.generators.hugging_face_local.StopWordsCriteria")
@@ -428,32 +405,52 @@ def test_warm_up_set_stopping_criteria_list(
         self, pipeline_mock, stop_words_criteria_mock, stopping_criteria_list_mock
     ):
         """
-        Test that warm_up method sets the `stopping_criteria_list` attribute
-        if `stop_words` is provided
+        Test that warm_up method sets the `stopping_criteria_list` attribute if `stop_words` is provided
         """
         generator = HuggingFaceLocalGenerator(
-            model="google/flan-t5-base", task="text2text-generation", stop_words=["coca", "cola"]
+            model="google/flan-t5-small", task="text2text-generation", stop_words=["coca", "cola"]
         )
-
         generator.warm_up()
-
         stop_words_criteria_mock.assert_called_once()
         stopping_criteria_list_mock.assert_called_once()
-
         assert hasattr(generator, "stopping_criteria_list")
 
     def test_run_stop_words_removal(self):
-        """
-        Test that stop words are removed from the generated text
-        (does not test stopping text generation)
-        """
+        """Test that stop words are removed from the generated text (does not test stopping text generation)"""
         generator = HuggingFaceLocalGenerator(
-            model="google/flan-t5-base", task="text2text-generation", stop_words=["world"]
+            model="google/flan-t5-small", task="text2text-generation", stop_words=["world"]
         )
-
-        # create the pipeline object (simulating the warm_up)
         generator.pipeline = Mock(return_value=[{"generated_text": "Hello world"}])
-
         results = generator.run(prompt="irrelevant")
-
         assert results == {"replies": ["Hello"]}
+
+    @pytest.mark.integration
+    def test_stop_words_criteria_using_hf_tokenizer(self):
+        """
+        Test that StopWordsCriteria catches stop word tokens in a continuous and sequential order in the input_ids
+        using a real Huggingface tokenizer.
+        """
+        from transformers import AutoTokenizer
+
+        model_name = "google/flan-t5-small"
+        tokenizer = AutoTokenizer.from_pretrained(model_name)
+        criteria = StopWordsCriteria(tokenizer=tokenizer, stop_words=["unambiguously"])
+
+        text_one = "This is ambiguously, but is unrelated."
+        generated_text_ids = tokenizer.encode(text_one, add_special_tokens=False, return_tensors="pt")
+        assert criteria(generated_text_ids, scores=None) is False
+
+        text_two = "This is unambiguously"
+        generated_text_ids = tokenizer.encode(text_two, add_special_tokens=False, return_tensors="pt")
+        assert criteria(generated_text_ids, scores=None) is True
+
+    @pytest.mark.integration
+    def test_hf_pipeline_runs_with_our_criteria(self):
+        """Test that creating our own StopWordsCriteria and passing it to a Huggingface pipeline works."""
+        generator = HuggingFaceLocalGenerator(
+            model="google/flan-t5-small", task="text2text-generation", stop_words=["unambiguously"]
+        )
+        generator.warm_up()
+        results = generator.run(prompt="something that triggers something")
+        assert results["replies"] != []
+        assert generator.stopping_criteria_list is not None