Repository: deepset-ai/haystack

Issue:
`HuggingFaceLocalGenerator` - stopwords support
Similarly to the 1.x Hugging Face Invocation Layer,
we want to provide out-of-the-box support for stopwords in `HuggingFaceLocalGenerator`.
The purpose is to allow stopping text generation when one of them is encountered.

Most of the code can be reutilized/migrated.


Discussion:

Resolving patch:
diff --git a/haystack/preview/components/generators/hugging_face/hugging_face_local.py b/haystack/preview/components/generators/hugging_face/hugging_face_local.py
--- a/haystack/preview/components/generators/hugging_face/hugging_face_local.py
+++ b/haystack/preview/components/generators/hugging_face/hugging_face_local.py
@@ -7,7 +7,18 @@
 
 with LazyImport(message="Run 'pip install transformers'") as transformers_import:
     from huggingface_hub import model_info
-    from transformers import pipeline
+    from transformers import (
+        pipeline,
+        StoppingCriteriaList,
+        StoppingCriteria,
+        PreTrainedTokenizer,
+        PreTrainedTokenizerFast,
+    )
+
+with LazyImport(
+    message="PyTorch is needed to run this component. Please install it by following the instructions at https://pytorch.org/"
+) as torch_import:
+    import torch
 
 logger = logging.getLogger(__name__)
 
@@ -45,6 +56,7 @@ def __init__(
         token: Optional[Union[str, bool]] = None,
         generation_kwargs: Optional[Dict[str, Any]] = None,
         pipeline_kwargs: Optional[Dict[str, Any]] = None,
+        stop_words: Optional[List[str]] = None,
     ):
         """
         :param model_name_or_path: The name or path of a Hugging Face model for text generation,
@@ -75,8 +87,13 @@ def __init__(
             In this dictionary, you can also include `model_kwargs` to specify the kwargs
             for model initialization:
             https://huggingface.co/docs/transformers/en/main_classes/model#transformers.PreTrainedModel.from_pretrained
+        :param stop_words: A list of stop words. If any one of the stop words is generated, the generation is stopped.
+            If you provide this parameter, you should not specify the `stopping_criteria` in `generation_kwargs`.
+            For some chat models, the output includes both the new text and the original prompt.
+            In these cases, it's important to make sure your prompt has no stop words.
         """
         transformers_import.check()
+        torch_import.check()
 
         pipeline_kwargs = pipeline_kwargs or {}
         generation_kwargs = generation_kwargs or {}
@@ -106,14 +123,28 @@ def __init__(
         if task == "text-generation":
             generation_kwargs.setdefault("return_full_text", False)
 
+        if stop_words and "stopping_criteria" in generation_kwargs:
+            raise ValueError(
+                "Found both the `stop_words` init parameter and the `stopping_criteria` key in `generation_kwargs`. "
+                "Please specify only one of them."
+            )
+
         self.pipeline_kwargs = pipeline_kwargs
         self.generation_kwargs = generation_kwargs
+        self.stop_words = stop_words
         self.pipeline = None
+        self.stopping_criteria_list = None
 
     def warm_up(self):
         if self.pipeline is None:
             self.pipeline = pipeline(**self.pipeline_kwargs)
 
+        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])
+
     def to_dict(self) -> Dict[str, Any]:
         """
         Serialize this component to a dictionary.
@@ -125,7 +156,10 @@ def to_dict(self) -> Dict[str, Any]:
             pipeline_kwargs_to_serialize["token"] = None
 
         return default_to_dict(
-            self, pipeline_kwargs=pipeline_kwargs_to_serialize, generation_kwargs=self.generation_kwargs
+            self,
+            pipeline_kwargs=pipeline_kwargs_to_serialize,
+            generation_kwargs=self.generation_kwargs,
+            stop_words=self.stop_words,
         )
 
     @classmethod
@@ -135,14 +169,56 @@ def from_dict(cls, data: Dict[str, Any]) -> "HuggingFaceLocalGenerator":
         """
         return default_from_dict(cls, data)
 
-    @component.output_types(replies=List[str], metadata=List[Dict[str, Any]])
+    @component.output_types(replies=List[str])
     def run(self, prompt: str):
         if self.pipeline is None:
             raise RuntimeError("The generation model has not been loaded. Please call warm_up() before running.")
 
-        replies = []
-        if prompt:
-            output = self.pipeline(prompt, **self.generation_kwargs)
-            replies = [o["generated_text"] for o in output if "generated_text" in o]
+        if not prompt:
+            return {"replies": []}
+
+        output = self.pipeline(prompt, stopping_criteria=self.stopping_criteria_list, **self.generation_kwargs)
+        replies = [o["generated_text"] for o in output if "generated_text" in o]
+
+        if self.stop_words:
+            # the output of the pipeline includes the stop word
+            replies = [reply.replace(stop_word, "").rstrip() for reply in replies for stop_word in self.stop_words]
 
         return {"replies": replies}
+
+
+class StopWordsCriteria(StoppingCriteria):
+    """
+    Stops text generation if any one of the stop words is generated.
+
+    Note: When a stop word is encountered, the generation of new text is stopped.
+    However, if the stop word is in the prompt itself, it can stop generating new text
+    prematurely after the first token. This is particularly important for LLMs designed
+    for dialogue generation. For these models, like for example mosaicml/mpt-7b-chat,
+    the output includes both the new text and the original prompt. Therefore, it's important
+    to make sure your prompt has no stop words.
+    """
+
+    def __init__(
+        self,
+        tokenizer: Union[PreTrainedTokenizer, PreTrainedTokenizerFast],
+        stop_words: List[str],
+        device: Union[str, torch.device] = "cpu",
+    ):
+        super().__init__()
+        encoded_stop_words = tokenizer(stop_words, add_special_tokens=False, padding=True, return_tensors="pt")
+        self.stop_ids = encoded_stop_words.input_ids.to(device)
+
+    def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs) -> bool:
+        for stop_id in self.stop_ids:
+            found_stop_word = self.is_stop_word_found(input_ids, stop_id)
+            if found_stop_word:
+                return True
+        return False
+
+    def is_stop_word_found(self, generated_text_ids: torch.Tensor, stop_id: torch.Tensor) -> bool:
+        generated_text_ids = generated_text_ids[-1]
+        len_generated_text_ids = generated_text_ids.size(0)
+        len_stop_id = stop_id.size(0)
+        result = all(generated_text_ids[len_generated_text_ids - len_stop_id :].eq(stop_id))
+        return result


Test patch:
diff --git a/test/preview/components/generators/hugging_face/test_hugging_face_local_generator.py b/test/preview/components/generators/hugging_face/test_hugging_face_local_generator.py
--- a/test/preview/components/generators/hugging_face/test_hugging_face_local_generator.py
+++ b/test/preview/components/generators/hugging_face/test_hugging_face_local_generator.py
@@ -1,8 +1,12 @@
 from unittest.mock import patch, Mock
 
 import pytest
+import torch
 
-from haystack.preview.components.generators.hugging_face.hugging_face_local import HuggingFaceLocalGenerator
+from haystack.preview.components.generators.hugging_face.hugging_face_local import (
+    HuggingFaceLocalGenerator,
+    StopWordsCriteria,
+)
 
 
 class TestHuggingFaceLocalGenerator:
@@ -122,6 +126,18 @@ def test_init_set_return_full_text(self):
 
         assert generator.generation_kwargs == {"return_full_text": False}
 
+    @pytest.mark.unit
+    def test_init_fails_with_both_stopwords_and_stoppingcriteria(self):
+        with pytest.raises(
+            ValueError,
+            match="Found both the `stop_words` init parameter and the `stopping_criteria` key in `generation_kwargs`",
+        ):
+            HuggingFaceLocalGenerator(
+                task="text2text-generation",
+                stop_words=["coca", "cola"],
+                generation_kwargs={"stopping_criteria": "fake-stopping-criteria"},
+            )
+
     @pytest.mark.unit
     @patch("haystack.preview.components.generators.hugging_face.hugging_face_local.model_info")
     def test_to_dict_default(self, model_info_mock):
@@ -135,6 +151,7 @@ def test_to_dict_default(self, model_info_mock):
             "init_parameters": {
                 "pipeline_kwargs": {"model": "google/flan-t5-base", "task": "text2text-generation", "token": None},
                 "generation_kwargs": {},
+                "stop_words": None,
             },
         }
 
@@ -146,6 +163,7 @@ def test_to_dict_with_parameters(self):
             device="cuda:0",
             token="test-token",
             generation_kwargs={"max_new_tokens": 100},
+            stop_words=["coca", "cola"],
         )
         data = component.to_dict()
 
@@ -159,6 +177,7 @@ def test_to_dict_with_parameters(self):
                     "device": "cuda:0",
                 },
                 "generation_kwargs": {"max_new_tokens": 100, "return_full_text": False},
+                "stop_words": ["coca", "cola"],
             },
         }
 
@@ -228,7 +247,9 @@ def test_run(self):
 
         results = generator.run(prompt="What's the capital of Italy?")
 
-        generator.pipeline.assert_called_once_with("What's the capital of Italy?", max_new_tokens=100)
+        generator.pipeline.assert_called_once_with(
+            "What's the capital of Italy?", max_new_tokens=100, stopping_criteria=None
+        )
         assert results == {"replies": ["Rome"]}
 
     @pytest.mark.unit
@@ -256,3 +277,79 @@ def test_run_fails_without_warm_up(self):
 
         with pytest.raises(RuntimeError, match="The generation model has not been loaded."):
             generator.run(prompt="irrelevant")
+
+    @pytest.mark.unit
+    def test_stop_words_criteria(self):
+        """
+        Test that StopWordsCriteria will check stop word tokens in a continuous and sequential order
+        """
+        # 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_criteria = StopWordsCriteria(tokenizer=Mock(), 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
+
+    @pytest.mark.unit
+    @patch("haystack.preview.components.generators.hugging_face.hugging_face_local.pipeline")
+    @patch("haystack.preview.components.generators.hugging_face.hugging_face_local.StopWordsCriteria")
+    @patch("haystack.preview.components.generators.hugging_face.hugging_face_local.StoppingCriteriaList")
+    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
+        """
+        generator = HuggingFaceLocalGenerator(
+            model_name_or_path="google/flan-t5-base", 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")
+
+    @pytest.mark.unit
+    def test_run_stop_words_removal(self):
+        """
+        Test that stop words are removed from the generated text
+        (does not test stopping text generation)
+        """
+        generator = HuggingFaceLocalGenerator(
+            model_name_or_path="google/flan-t5-base", 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"]}