Repository: deepset-ai/haystack
Pull request: feat: add `top_k` to `PromptNode`
### Related Issues
- fixes https://github.com/deepset-ai/haystack/issues/4158

### Proposed Changes:
- add `top_k` parameter to `PromptNode`
- push down `top_k` to invocation layers and map them to their invocation interface

### How did you test it?
- added tests

### Notes for the reviewer
- I don't know why `test/nodes/test_audio.py::TestTextToSpeech::test_text_to_speech_compress_audio` fails. The code didn't touch any of TestToSpeech. Also other PR CIs are failing because of this. Maybe a dependency update?...

### Checklist
- [x] I have read the [contributors guidelines](https://github.com/deepset-ai/haystack/blob/main/CONTRIBUTING.md) and the [code of conduct](https://github.com/deepset-ai/haystack/blob/main/code_of_conduct.txt)
- [x] I have updated the related issue with new insights and changes
- [x] I added tests that demonstrate the correct behavior of the change
- [x] I've used one of the [conventional commit types](https://www.conventionalcommits.org/en/v1.0.0/) for my PR title: `fix:`, `feat:`, `build:`, `chore:`, `ci:`, `docs:`, `style:`, `refactor:`, `perf:`, `test:`.
- [x] I documented my code
- [x] I ran [pre-commit hooks](https://github.com/deepset-ai/haystack/blob/main/CONTRIBUTING.md#installation) and fixed any issue


Problem statement:
Add `top_k` to `PromptNode`
**Is your feature request related to a problem? Please describe.**
Setting top_k on `PromptNode` is quite cumbersome. Depending on which invocation context you use, you have to set different params on `PromptModel`:
***For hf models:***
```yaml
....
  - name: my_prompt_model
    type: PromptModel
    params:
      model_name_or_path: google/flan-t5-large
      model_kwargs:
        model_kwargs:
          num_return_sequences: 3
          num_beams: 3
....
```
***For OpenAI models:***
```yaml
...
  - name: my_prompt_model
    type: PromptModel
    params:
      model_name_or_path: text-davinci-003
      model_kwargs:
        n: 3
        best_of: 3
...
```

`PromptNode` misses the functionality to set `top_k` completely. This is a problem if you want to use the same `PromptModel` with different `PromptNodes` (i.e. for different use-cases), that require differnt `top_k`.

**Describe the solution you'd like**
- Add `top_k` param to `PromptNode`
- push down `top_k` to the invocation layers

**Describe alternatives you've considered**
- make `top_k` part of `PromptTemplate`: it's not the right place, as with `PromptTemplate you specify the use-case in general, but not the instance of a use-case. `top_k` should be part of the instance of a use-case and thus part of `PromptNode`

Hints from issue discussion:

Commit message: add top_k to PromptNode

Patch to review:
diff --git a/haystack/nodes/prompt/prompt_node.py b/haystack/nodes/prompt/prompt_node.py
index 520c755912..a346e73467 100644
--- a/haystack/nodes/prompt/prompt_node.py
+++ b/haystack/nodes/prompt/prompt_node.py
@@ -333,6 +333,7 @@ def invoke(self, *args, **kwargs):
         """
         output: List[Dict[str, str]] = []
         stop_words = kwargs.pop("stop_words", None)
+        top_k = kwargs.pop("top_k", None)
         if kwargs and "prompt" in kwargs:
             prompt = kwargs.pop("prompt")
 
@@ -347,6 +348,9 @@ def invoke(self, *args, **kwargs):
             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)
         generated_texts = [o["generated_text"] for o in output if "generated_text" in o]
 
@@ -459,6 +463,8 @@ def invoke(self, *args, **kwargs):
             # we use keyword stop_words but OpenAI uses stop
             if "stop_words" in kwargs:
                 kwargs["stop"] = kwargs.pop("stop_words")
+            if "top_k" in kwargs:
+                kwargs["n"] = kwargs.pop("top_k")
             kwargs_with_defaults.update(kwargs)
         payload = {
             "model": self.model_name_or_path,
@@ -711,6 +717,7 @@ def __init__(
         use_gpu: Optional[bool] = None,
         devices: Optional[List[Union[str, torch.device]]] = None,
         stop_words: Optional[List[str]] = None,
+        top_k: int = 1,
     ):
         """
         Creates a PromptNode instance.
@@ -723,6 +730,7 @@ def __init__(
         :param use_auth_token: The authentication token to use for the model.
         :param use_gpu: Whether to use GPU or not.
         :param devices: The devices to use for the model.
+        :param top_k: Number of independently generated text to return per prompt.
         """
         super().__init__()
         self.prompt_templates: Dict[str, PromptTemplate] = {pt.name: pt for pt in get_predefined_prompt_templates()}  # type: ignore
@@ -731,6 +739,7 @@ def __init__(
         self.model_name_or_path: Union[str, PromptModel] = model_name_or_path
         self.prompt_model: PromptModel
         self.stop_words: Optional[List[str]] = stop_words
+        self.top_k: int = top_k
         if isinstance(self.default_prompt_template, str) and not self.is_supported_template(
             self.default_prompt_template
         ):
@@ -979,4 +988,4 @@ def run_batch(
     def _prepare_model_kwargs(self):
         # these are the parameters from PromptNode level
         # that are passed to the prompt model invocation layer
-        return {"stop_words": self.stop_words}
+        return {"stop_words": self.stop_words, "top_k": self.top_k}
diff --git a/test/nodes/test_prompt_node.py b/test/nodes/test_prompt_node.py
index f53b1a5bdb..5da8f6ff98 100644
--- a/test/nodes/test_prompt_node.py
+++ b/test/nodes/test_prompt_node.py
@@ -7,6 +7,7 @@
 
 from haystack import Document, Pipeline, BaseComponent, MultiLabel
 from haystack.errors import OpenAIError
+from haystack.nodes.other.shaper import Shaper
 from haystack.nodes.prompt import PromptTemplate, PromptNode, PromptModel
 from haystack.nodes.prompt.prompt_node import HFLocalInvocationLayer
 
@@ -329,6 +330,22 @@ def test_complex_pipeline(prompt_model):
     assert "berlin" in result["results"][0].casefold()
 
 
+@pytest.mark.integration
+@pytest.mark.parametrize("prompt_model", ["hf", "openai"], indirect=True)
+def test_simple_pipeline_with_topk(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, default_prompt_template="question-generation", top_k=2)
+
+    pipe = Pipeline()
+    pipe.add_node(component=node, name="prompt_node", inputs=["Query"])
+    result = pipe.run(query="not relevant", documents=[Document("Berlin is the capital of Germany")])
+
+    assert len(result["results"]) == 2
+    assert result["results"][0] != result["results"][1]
+
+
 @pytest.mark.integration
 @pytest.mark.parametrize("prompt_model", ["hf", "openai"], indirect=True)
 def test_complex_pipeline_with_qa(prompt_model):


Code review comments:
diff_hunk: @@ -459,6 +463,8 @@ def invoke(self, *args, **kwargs):
             # we use keyword stop_words but OpenAI uses stop
             if "stop_words" in kwargs:
                 kwargs["stop"] = kwargs.pop("stop_words")
+            if "top_k" in kwargs:
+                kwargs["n"] = kwargs.pop("top_k")
path: haystack/nodes/prompt/prompt_node.py
text: @author:
Handling in `OpenAIInvocationLayer`