Repository: deepset-ai/haystack

Issue:
Support for Bloom models
**Is your feature request related to a problem? Please describe.**
We currently support sequence-to-sequence encoder/decoder architectures that use text2text-generation pipelines. So, it's only possible to use google/flan t_5 models and OpenAI InstructGPT models at the moment. We can provide support for other LLMs. 

**Describe the solution you'd like**
Add support for textgeneration thus support for Bloom models (BloomZ)

**Describe alternatives you've considered**
N/A

**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
@@ -9,7 +9,8 @@
 
 from haystack import MultiLabel
 from haystack.nodes.base import BaseComponent
-from haystack.nodes.prompt.providers import PromptModelInvocationLayer
+from haystack.nodes.prompt.providers import PromptModelInvocationLayer, instruction_following_models
+
 from haystack.schema import Document
 from haystack.telemetry_2 import send_event
 
@@ -164,7 +165,7 @@ class PromptModel(BaseComponent):
     """
     The PromptModel class is a component that uses a pre-trained model to perform tasks based on a prompt. Out of
     the box, it supports model invocation layers for:
-    - Hugging Face transformers (all text2text-generation models)
+    - Hugging Face transformers (all text2text-generation and text-generation models)
     - OpenAI InstructGPT models
     - Azure OpenAI InstructGPT models
 
@@ -215,6 +216,14 @@ def __init__(
 
         self.model_kwargs = model_kwargs if model_kwargs else {}
         self.model_invocation_layer = self.create_invocation_layer(invocation_layer_class=invocation_layer_class)
+        is_instruction_following: bool = any(m in model_name_or_path for m in instruction_following_models())
+        if not is_instruction_following:
+            logger.warning(
+                "PromptNode has been potentially initialized with a language model not fine-tuned on instruction following tasks. "
+                "Many of the default prompts and PromptTemplates will likely not work as intended. "
+                "Use custom prompts and PromptTemplates specific to the %s model",
+                model_name_or_path,
+            )
 
     def create_invocation_layer(
         self, invocation_layer_class: Optional[Type[PromptModelInvocationLayer]]
@@ -391,7 +400,8 @@ class PromptNode(BaseComponent):
     additional custom model invocation layers.
 
     We recommend using LLMs fine-tuned on a collection of datasets phrased as instructions, otherwise we find that the
-    LLM does not "follow" prompt instructions well. This is why we recommend using T5 flan or OpenAI InstructGPT models.
+    LLM does not "follow" prompt instructions well. The list of instructions following models increases every month,
+    and the current list includes: Flan, OpenAI InstructGPT, opt-iml, bloomz, and mt0 models.
 
     For more details, see  [PromptNode](https://docs.haystack.deepset.ai/docs/prompt_node).
     """
diff --git a/haystack/nodes/prompt/providers.py b/haystack/nodes/prompt/providers.py
--- a/haystack/nodes/prompt/providers.py
+++ b/haystack/nodes/prompt/providers.py
@@ -6,15 +6,8 @@
 
 import sseclient
 import torch
-from transformers import (
-    pipeline,
-    AutoConfig,
-    StoppingCriteriaList,
-    StoppingCriteria,
-    PreTrainedTokenizer,
-    PreTrainedTokenizerFast,
-)
-from transformers.models.auto.modeling_auto import MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING_NAMES
+from transformers import pipeline, StoppingCriteriaList, StoppingCriteria, PreTrainedTokenizer, PreTrainedTokenizerFast
+from transformers.pipelines import get_task
 
 from haystack.errors import OpenAIError
 from haystack.modeling.utils import initialize_device_settings
@@ -123,6 +116,10 @@ def _ensure_token_limit(self, prompt: Union[str, List[Dict[str, str]]]) -> Union
         pass
 
 
+def instruction_following_models() -> List[str]:
+    return ["flan", "mt0", "bloomz", "davinci", "opt-iml"]
+
+
 class StopWordsCriteria(StoppingCriteria):
     """
     Stops text generation if any one of the stop words is generated.
@@ -223,9 +220,8 @@ def __init__(
 
         if len(model_input_kwargs) > 0:
             logger.info("Using model input kwargs %s in %s", model_input_kwargs, self.__class__.__name__)
-
+        self.task_name = get_task(model_name_or_path, use_auth_token=use_auth_token)
         self.pipe = pipeline(
-            "text2text-generation",
             model=model_name_or_path,
             device=self.devices[0] if "device_map" not in model_input_kwargs else None,
             use_auth_token=self.use_auth_token,
@@ -241,8 +237,8 @@ def invoke(self, *args, **kwargs):
         It takes a prompt and returns a list of generated text using the local Hugging Face transformers model
         :return: A list of generated text.
 
-        Note: Only kwargs relevant to Text2TextGenerationPipeline are passed to Hugging Face as model_input_kwargs.
-        Other kwargs are ignored.
+        Note: Only kwargs relevant to Text2TextGenerationPipeline and TextGenerationPipeline are passed to
+        Hugging Face as model_input_kwargs. Other kwargs are ignored.
         """
         output: List[Dict[str, str]] = []
         stop_words = kwargs.pop("stop_words", None)
@@ -250,21 +246,40 @@ def invoke(self, *args, **kwargs):
         if kwargs and "prompt" in kwargs:
             prompt = kwargs.pop("prompt")
 
-            # Consider only Text2TextGenerationPipeline relevant, ignore others
-            # For more details refer to Hugging Face Text2TextGenerationPipeline documentation
+            # Consider only Text2TextGenerationPipeline and TextGenerationPipeline relevant, ignore others
+            # For more details refer to Hugging Face Text2TextGenerationPipeline and TextGenerationPipeline
+            # documentation
             # TODO resolve these kwargs from the pipeline signature
             model_input_kwargs = {
                 key: kwargs[key]
-                for key in ["return_tensors", "return_text", "clean_up_tokenization_spaces", "truncation"]
+                for key in [
+                    "return_tensors",
+                    "return_text",
+                    "return_full_text",
+                    "clean_up_tokenization_spaces",
+                    "truncation",
+                ]
                 if key in kwargs
             }
+            is_text_generation = "text-generation" == self.task_name
+            # Prefer return_full_text is False for text-generation (unless explicitly set)
+            # Thus only generated text is returned (excluding prompt)
+            if is_text_generation and "return_full_text" not in model_input_kwargs:
+                model_input_kwargs["return_full_text"] = False
+                model_input_kwargs["max_new_tokens"] = self.max_length
             if stop_words:
                 sw = StopWordsCriteria(tokenizer=self.pipe.tokenizer, stop_words=stop_words)
                 model_input_kwargs["stopping_criteria"] = StoppingCriteriaList([sw])
             if top_k:
                 model_input_kwargs["num_return_sequences"] = top_k
                 model_input_kwargs["num_beams"] = top_k
-            output = self.pipe(prompt, max_length=self.max_length, **model_input_kwargs)
+            # max_new_tokens is used for text-generation and max_length for text2text-generation
+            if is_text_generation:
+                model_input_kwargs["max_new_tokens"] = self.max_length
+            else:
+                model_input_kwargs["max_length"] = self.max_length
+
+            output = self.pipe(prompt, **model_input_kwargs)
         generated_texts = [o["generated_text"] for o in output if "generated_text" in o]
 
         if stop_words:
@@ -304,22 +319,14 @@ def _ensure_token_limit(self, prompt: Union[str, List[Dict[str, str]]]) -> Union
 
     @classmethod
     def supports(cls, model_name_or_path: str, **kwargs) -> bool:
+        task_name: Optional[str] = None
         try:
-            config = AutoConfig.from_pretrained(model_name_or_path)
-        except OSError:
-            # This is needed so OpenAI models are skipped over
+            task_name = get_task(model_name_or_path, use_auth_token=kwargs.get("use_auth_token", None))
+        except RuntimeError:
+            # This will fail for all non-HF models
             return False
 
-        if not all(m in model_name_or_path for m in ["flan", "t5"]):
-            logger.warning(
-                "PromptNode has been potentially initialized with a language model not fine-tuned on instruction following tasks. "
-                "Many of the default prompts and PromptTemplates will likely not work as intended. "
-                "Use custom prompts and PromptTemplates specific to the %s model",
-                model_name_or_path,
-            )
-
-        supported_models = list(MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING_NAMES.values())
-        return config.architectures[0] in supported_models
+        return task_name in ["text2text-generation", "text-generation"]
 
 
 class OpenAIInvocationLayer(PromptModelInvocationLayer):


Test patch:
diff --git a/test/prompt/test_prompt_node.py b/test/prompt/test_prompt_node.py
--- a/test/prompt/test_prompt_node.py
+++ b/test/prompt/test_prompt_node.py
@@ -408,6 +408,21 @@ def test_streaming_prompt_node():
     assert ttsh.stream_handler_invoked, "Stream handler should have been invoked"
 
 
+def test_prompt_node_with_text_generation_model():
+    # test simple prompting with text generation model
+    # by default, we force the model not return prompt text
+    # Thus text-generation models can be used with PromptNode
+    # just like text2text-generation models
+    node = PromptNode("bigscience/bigscience-small-testing")
+    r = node("Hello big science!")
+    assert len(r[0]) > 0
+
+    # test prompting with parameter to return prompt text as well
+    # users can use this param to get the prompt text and the generated text
+    r = node("Hello big science!", return_full_text=True)
+    assert len(r[0]) > 0 and r[0].startswith("Hello big science!")
+
+
 @pytest.mark.integration
 @pytest.mark.parametrize("prompt_model", ["hf", "openai", "azure"], indirect=True)
 def test_simple_pipeline(prompt_model):