Repository: deepset-ai/haystack

Issue:
Add stop word sequence to PromptTemplate
**Is your feature request related to a problem? Please describe.**
As a `PromptNode` user, I want to be able to direct the model to stop generation if it encounters specific word sequences in the generated text. 

**Describe the solution you'd like**
An easy-to-use API that allows me to specify stop word sequences both programmatically and from YAML configuration files. 

**Describe alternatives you've considered**
Currently, the only alternative is to get the result of `PromptNode` text generation, find one of the stop words and cut off the generated text after the stop words. However, even with this workaround, I can not achieve stopping of text generation by simply specifying the stop words in the YAML config.

**Additional context**
N/A


Discussion:

Resolving patch:
diff --git a/haystack/nodes/prompt/prompt_node.py b/haystack/nodes/prompt/prompt_node.py
--- a/haystack/nodes/prompt/prompt_node.py
+++ b/haystack/nodes/prompt/prompt_node.py
@@ -7,7 +7,7 @@
 
 import requests
 import torch
-from transformers import pipeline, AutoModelForSeq2SeqLM
+from transformers import pipeline, AutoModelForSeq2SeqLM, StoppingCriteria, StoppingCriteriaList, AutoTokenizer
 
 from haystack import MultiLabel
 from haystack.errors import OpenAIError, OpenAIRateLimitError
@@ -206,6 +206,20 @@ def supports(cls, model_name_or_path: str) -> bool:
         return False
 
 
+class StopWordsCriteria(StoppingCriteria):
+    """
+    Stops text generation if any one of the stop words is generated.
+    """
+
+    def __init__(self, model_name_or_path: str, stop_words: List[str]):
+        super().__init__()
+        tokenizer = AutoTokenizer.from_pretrained(model_name_or_path)
+        self.stop_words = tokenizer.encode(stop_words, add_special_tokens=False, return_tensors="pt")
+
+    def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs) -> bool:
+        return any(torch.isin(input_ids[-1], self.stop_words[-1]))
+
+
 class HFLocalInvocationLayer(PromptModelInvocationLayer):
     """
     A subclass of the PromptModelInvocationLayer class. It loads a pre-trained model from Hugging Face and
@@ -310,7 +324,8 @@ def invoke(self, *args, **kwargs):
         Note: Only kwargs relevant to Text2TextGenerationPipeline are passed to Hugging Face as model_input_kwargs.
         Other kwargs are ignored.
         """
-        output = []
+        output: List[Dict[str, str]] = []
+        stop_words = kwargs.pop("stop", None)
         if kwargs and "prompt" in kwargs:
             prompt = kwargs.pop("prompt")
 
@@ -322,8 +337,19 @@ def invoke(self, *args, **kwargs):
                 for key in ["return_tensors", "return_text", "clean_up_tokenization_spaces", "truncation"]
                 if key in kwargs
             }
+            if stop_words:
+                sw = StopWordsCriteria(model_name_or_path=self.model_name_or_path, stop_words=stop_words)
+                model_input_kwargs["stopping_criteria"] = StoppingCriteriaList([sw])
             output = self.pipe(prompt, max_length=self.max_length, **model_input_kwargs)
-        return [o["generated_text"] for o in output]
+        generated_texts = [o["generated_text"] for o in output if "generated_text" in o]
+
+        if stop_words:
+            # Although HF generates text until stop words are encountered unfortunately it includes the stop word
+            # We want to exclude it to be consistent with other invocation layers
+            for idx, _ in enumerate(generated_texts):
+                for stop_word in stop_words:
+                    generated_texts[idx] = generated_texts[idx].replace(stop_word, "").strip()
+        return generated_texts
 
     @classmethod
     def supports(cls, model_name_or_path: str) -> bool:
@@ -655,6 +681,7 @@ def __init__(
         use_auth_token: Optional[Union[str, bool]] = None,
         use_gpu: Optional[bool] = None,
         devices: Optional[List[Union[str, torch.device]]] = None,
+        stop_words: Optional[List[str]] = None,
     ):
         """
         Creates a PromptNode instance.
@@ -674,6 +701,7 @@ def __init__(
         self.output_variable: Optional[str] = output_variable
         self.model_name_or_path: Union[str, PromptModel] = model_name_or_path
         self.prompt_model: PromptModel
+        self.stop_words: Optional[List[str]] = stop_words
         if isinstance(self.default_prompt_template, str) and not self.is_supported_template(
             self.default_prompt_template
         ):
@@ -743,7 +771,9 @@ def prompt(self, prompt_template: Optional[Union[str, PromptTemplate]], *args, *
 
             # prompt template used, yield prompts from inputs args
             for prompt in template_to_fill.fill(*args, **kwargs):
-                # and pass the prepared prompt to the model
+                # merge any additional model kwargs
+                kwargs = {**kwargs, **self._prepare_model_kwargs()}
+                # and pass the prepared prompt and kwargs to the model
                 output = self.prompt_model.invoke(prompt, **kwargs)
                 results.extend(output)
         else:
@@ -894,3 +924,6 @@ def run_batch(
         debug: Optional[bool] = None,
     ):
         pass
+
+    def _prepare_model_kwargs(self):
+        return {"stop": self.stop_words}


Test patch:
diff --git a/test/nodes/test_prompt_node.py b/test/nodes/test_prompt_node.py
--- a/test/nodes/test_prompt_node.py
+++ b/test/nodes/test_prompt_node.py
@@ -240,6 +240,16 @@ def test_open_ai_prompt_with_params():
     assert len(r) == 1 and len(r[0]) > 0
 
 
+@pytest.mark.parametrize("prompt_model", ["hf", "openai"], indirect=True)
+def test_stop_words(prompt_model):
+    if prompt_model.api_key is not None and not is_openai_api_key_set(prompt_model.api_key):
+        pytest.skip("No API key found for OpenAI, skipping test")
+
+    node = PromptNode(prompt_model, stop_words=["capital", "Germany"])
+    r = node.prompt("question-generation", documents=["Berlin is the capital of Germany."])
+    assert r[0] == "What is the"
+
+
 @pytest.mark.parametrize("prompt_model", ["hf", "openai"], indirect=True)
 def test_simple_pipeline(prompt_model):
     if prompt_model.api_key is not None and not is_openai_api_key_set(prompt_model.api_key):