We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
Code snippet wrongly formatted in docs
**Describe the bug**
The code snippet at:

https://github.com/deepset-ai/haystack/blob/master/haystack/nodes/ranker/sentence_transformers.py#L32

is rendered as plain text in the corresponding documentation page

https://haystack.deepset.ai/reference/ranker


**Expected behavior**
The snipped should be formatted as Python code


--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `haystack/nodes/ranker/sentence_transformers.py`
Content:
```
1 from typing import List, Optional, Union, Tuple, Iterator, Any
2 import logging
3 from pathlib import Path
4 
5 import torch
6 from torch.nn import DataParallel
7 from transformers import AutoModelForSequenceClassification, AutoTokenizer
8 
9 from haystack.errors import HaystackError
10 from haystack.schema import Document
11 from haystack.nodes.ranker.base import BaseRanker
12 from haystack.modeling.utils import initialize_device_settings
13 
14 logger = logging.getLogger(__name__)
15 
16 
17 class SentenceTransformersRanker(BaseRanker):
18     """
19     Sentence Transformer based pre-trained Cross-Encoder model for Document Re-ranking (https://huggingface.co/cross-encoder).
20     Re-Ranking can be used on top of a retriever to boost the performance for document search. This is particularly useful if the retriever has a high recall but is bad in sorting the documents by relevance.
21 
22     SentenceTransformerRanker handles Cross-Encoder models
23         - use a single logit as similarity score e.g.  cross-encoder/ms-marco-MiniLM-L-12-v2
24         - use two output logits (no_answer, has_answer) e.g. deepset/gbert-base-germandpr-reranking
25     https://www.sbert.net/docs/pretrained-models/ce-msmarco.html#usage-with-transformers
26 
27     |  With a SentenceTransformersRanker, you can:
28      - directly get predictions via predict()
29 
30     Usage example:
31     ...
32     retriever = BM25Retriever(document_store=document_store)
33     ranker = SentenceTransformersRanker(model_name_or_path="cross-encoder/ms-marco-MiniLM-L-12-v2")
34     p = Pipeline()
35     p.add_node(component=retriever, name="ESRetriever", inputs=["Query"])
36     p.add_node(component=ranker, name="Ranker", inputs=["ESRetriever"])
37     """
38 
39     def __init__(
40         self,
41         model_name_or_path: Union[str, Path],
42         model_version: Optional[str] = None,
43         top_k: int = 10,
44         use_gpu: bool = True,
45         devices: Optional[List[Union[str, torch.device]]] = None,
46         batch_size: Optional[int] = None,
47         scale_score: bool = True,
48     ):
49         """
50         :param model_name_or_path: Directory of a saved model or the name of a public model e.g.
51         'cross-encoder/ms-marco-MiniLM-L-12-v2'.
52         See https://huggingface.co/cross-encoder for full list of available models
53         :param model_version: The version of model to use from the HuggingFace model hub. Can be tag name, branch name, or commit hash.
54         :param top_k: The maximum number of documents to return
55         :param use_gpu: Whether to use all available GPUs or the CPU. Falls back on CPU if no GPU is available.
56         :param devices: List of GPU (or CPU) devices, to limit inference to certain GPUs and not use all available ones
57                         The strings will be converted into pytorch devices, so use the string notation described here:
58                         https://pytorch.org/docs/stable/tensor_attributes.html?highlight=torch%20device#torch.torch.device
59                         (e.g. ["cuda:0"]).
60         :param batch_size: Number of documents to process at a time.
61         :param scale_score: The raw predictions will be transformed using a Sigmoid activation function in case the model
62                             only predicts a single label. For multi-label predictions, no scaling is applied. Set this
63                             to False if you do not want any scaling of the raw predictions.
64         """
65         super().__init__()
66 
67         self.top_k = top_k
68 
69         if devices is not None:
70             self.devices = [torch.device(device) for device in devices]
71         else:
72             self.devices, _ = initialize_device_settings(use_cuda=use_gpu, multi_gpu=True)
73 
74         self.transformer_model = AutoModelForSequenceClassification.from_pretrained(
75             pretrained_model_name_or_path=model_name_or_path, revision=model_version
76         )
77         self.transformer_model.to(str(self.devices[0]))
78         self.transformer_tokenizer = AutoTokenizer.from_pretrained(
79             pretrained_model_name_or_path=model_name_or_path, revision=model_version
80         )
81         self.transformer_model.eval()
82 
83         # we use sigmoid activation function to scale the score in case there is only a single label
84         # we do not apply any scaling when scale_score is set to False
85         num_labels = self.transformer_model.num_labels
86         self.activation_function: torch.nn.Module
87         if num_labels == 1 and scale_score:
88             self.activation_function = torch.nn.Sigmoid()
89         else:
90             self.activation_function = torch.nn.Identity()
91 
92         if len(self.devices) > 1:
93             self.model = DataParallel(self.transformer_model, device_ids=self.devices)
94 
95         self.batch_size = batch_size
96 
97     def predict(self, query: str, documents: List[Document], top_k: Optional[int] = None) -> List[Document]:
98         """
99         Use loaded ranker model to re-rank the supplied list of Document.
100 
101         Returns list of Document sorted by (desc.) similarity with the query.
102 
103         :param query: Query string
104         :param documents: List of Document to be re-ranked
105         :param top_k: The maximum number of documents to return
106         :return: List of Document
107         """
108         if top_k is None:
109             top_k = self.top_k
110 
111         features = self.transformer_tokenizer(
112             [query for doc in documents],
113             [doc.content for doc in documents],
114             padding=True,
115             truncation=True,
116             return_tensors="pt",
117         ).to(self.devices[0])
118 
119         # SentenceTransformerRanker uses:
120         # 1. the logit as similarity score/answerable classification
121         # 2. the logits as answerable classification  (no_answer / has_answer)
122         # https://www.sbert.net/docs/pretrained-models/ce-msmarco.html#usage-with-transformers
123         with torch.no_grad():
124             similarity_scores = self.transformer_model(**features).logits
125 
126         logits_dim = similarity_scores.shape[1]  # [batch_size, logits_dim]
127         sorted_scores_and_documents = sorted(
128             zip(similarity_scores, documents),
129             key=lambda similarity_document_tuple:
130             # assume the last element in logits represents the `has_answer` label
131             similarity_document_tuple[0][-1] if logits_dim >= 2 else similarity_document_tuple[0],
132             reverse=True,
133         )
134 
135         # add normalized scores to documents
136         sorted_documents = self._add_scores_to_documents(sorted_scores_and_documents[:top_k], logits_dim)
137 
138         return sorted_documents
139 
140     def _add_scores_to_documents(
141         self, sorted_scores_and_documents: List[Tuple[Any, Document]], logits_dim: int
142     ) -> List[Document]:
143         """
144         Normalize and add scores to retrieved result documents.
145 
146         :param sorted_scores_and_documents: List of score, Document Tuples.
147         :param logits_dim: Dimensionality of the returned scores.
148         """
149         sorted_documents = []
150         for raw_score, doc in sorted_scores_and_documents:
151             if logits_dim >= 2:
152                 score = self.activation_function(raw_score)[-1]
153             else:
154                 score = self.activation_function(raw_score)[0]
155 
156             doc.score = score.detach().cpu().numpy().tolist()
157             sorted_documents.append(doc)
158 
159         return sorted_documents
160 
161     def predict_batch(
162         self,
163         queries: List[str],
164         documents: Union[List[Document], List[List[Document]]],
165         top_k: Optional[int] = None,
166         batch_size: Optional[int] = None,
167     ) -> Union[List[Document], List[List[Document]]]:
168         """
169         Use loaded ranker model to re-rank the supplied lists of Documents.
170 
171         Returns lists of Documents sorted by (desc.) similarity with the corresponding queries.
172 
173 
174         - If you provide a list containing a single query...
175 
176             - ... and a single list of Documents, the single list of Documents will be re-ranked based on the
177               supplied query.
178             - ... and a list of lists of Documents, each list of Documents will be re-ranked individually based on the
179               supplied query.
180 
181 
182         - If you provide a list of multiple queries...
183 
184             - ... you need to provide a list of lists of Documents. Each list of Documents will be re-ranked based on
185               its corresponding query.
186 
187         :param queries: Single query string or list of queries
188         :param documents: Single list of Documents or list of lists of Documents to be reranked.
189         :param top_k: The maximum number of documents to return per Document list.
190         :param batch_size: Number of Documents to process at a time.
191         """
192         if top_k is None:
193             top_k = self.top_k
194 
195         if batch_size is None:
196             batch_size = self.batch_size
197 
198         number_of_docs, all_queries, all_docs, single_list_of_docs = self._preprocess_batch_queries_and_docs(
199             queries=queries, documents=documents
200         )
201 
202         batches = self._get_batches(all_queries=all_queries, all_docs=all_docs, batch_size=batch_size)
203         preds = []
204         for cur_queries, cur_docs in batches:
205             features = self.transformer_tokenizer(
206                 cur_queries, [doc.content for doc in cur_docs], padding=True, truncation=True, return_tensors="pt"
207             ).to(self.devices[0])
208 
209             with torch.no_grad():
210                 similarity_scores = self.transformer_model(**features).logits
211                 preds.extend(similarity_scores)
212 
213         logits_dim = similarity_scores.shape[1]  # [batch_size, logits_dim]
214         if single_list_of_docs:
215             sorted_scores_and_documents = sorted(
216                 zip(similarity_scores, documents),
217                 key=lambda similarity_document_tuple:
218                 # assume the last element in logits represents the `has_answer` label
219                 similarity_document_tuple[0][-1] if logits_dim >= 2 else similarity_document_tuple[0],
220                 reverse=True,
221             )
222 
223             # is this step needed?
224             sorted_documents = [(score, doc) for score, doc in sorted_scores_and_documents if isinstance(doc, Document)]
225             sorted_documents_with_scores = self._add_scores_to_documents(sorted_documents[:top_k], logits_dim)
226 
227             return sorted_documents_with_scores
228         else:
229             # Group predictions together
230             grouped_predictions = []
231             left_idx = 0
232             right_idx = 0
233             for number in number_of_docs:
234                 right_idx = left_idx + number
235                 grouped_predictions.append(similarity_scores[left_idx:right_idx])
236                 left_idx = right_idx
237 
238             result = []
239             for pred_group, doc_group in zip(grouped_predictions, documents):
240                 sorted_scores_and_documents = sorted(
241                     zip(pred_group, doc_group),  # type: ignore
242                     key=lambda similarity_document_tuple:
243                     # assume the last element in logits represents the `has_answer` label
244                     similarity_document_tuple[0][-1] if logits_dim >= 2 else similarity_document_tuple[0],
245                     reverse=True,
246                 )
247 
248                 # rank documents according to scores
249                 sorted_documents = [
250                     (score, doc) for score, doc in sorted_scores_and_documents if isinstance(doc, Document)
251                 ]
252                 sorted_documents_with_scores = self._add_scores_to_documents(sorted_documents[:top_k], logits_dim)
253 
254                 result.append(sorted_documents_with_scores)
255 
256             return result
257 
258     def _preprocess_batch_queries_and_docs(
259         self, queries: List[str], documents: Union[List[Document], List[List[Document]]]
260     ) -> Tuple[List[int], List[str], List[Document], bool]:
261         number_of_docs = []
262         all_queries = []
263         all_docs: List[Document] = []
264         single_list_of_docs = False
265 
266         # Docs case 1: single list of Documents -> rerank single list of Documents based on single query
267         if len(documents) > 0 and isinstance(documents[0], Document):
268             if len(queries) != 1:
269                 raise HaystackError("Number of queries must be 1 if a single list of Documents is provided.")
270             query = queries[0]
271             number_of_docs = [len(documents)]
272             all_queries = [query] * len(documents)
273             all_docs = documents  # type: ignore
274             single_list_of_docs = True
275 
276         # Docs case 2: list of lists of Documents -> rerank each list of Documents based on corresponding query
277         # If queries contains a single query, apply it to each list of Documents
278         if len(documents) > 0 and isinstance(documents[0], list):
279             if len(queries) == 1:
280                 queries = queries * len(documents)
281             if len(queries) != len(documents):
282                 raise HaystackError("Number of queries must be equal to number of provided Document lists.")
283             for query, cur_docs in zip(queries, documents):
284                 if not isinstance(cur_docs, list):
285                     raise HaystackError(f"cur_docs was of type {type(cur_docs)}, but expected a list of Documents.")
286                 number_of_docs.append(len(cur_docs))
287                 all_queries.extend([query] * len(cur_docs))
288                 all_docs.extend(cur_docs)
289 
290         return number_of_docs, all_queries, all_docs, single_list_of_docs
291 
292     @staticmethod
293     def _get_batches(
294         all_queries: List[str], all_docs: List[Document], batch_size: Optional[int]
295     ) -> Iterator[Tuple[List[str], List[Document]]]:
296         if batch_size is None:
297             yield all_queries, all_docs
298             return
299         else:
300             for index in range(0, len(all_queries), batch_size):
301                 yield all_queries[index : index + batch_size], all_docs[index : index + batch_size]
302
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/haystack/nodes/ranker/sentence_transformers.py b/haystack/nodes/ranker/sentence_transformers.py
--- a/haystack/nodes/ranker/sentence_transformers.py
+++ b/haystack/nodes/ranker/sentence_transformers.py
@@ -28,12 +28,14 @@
      - directly get predictions via predict()
 
     Usage example:
-    ...
-    retriever = BM25Retriever(document_store=document_store)
-    ranker = SentenceTransformersRanker(model_name_or_path="cross-encoder/ms-marco-MiniLM-L-12-v2")
-    p = Pipeline()
-    p.add_node(component=retriever, name="ESRetriever", inputs=["Query"])
-    p.add_node(component=ranker, name="Ranker", inputs=["ESRetriever"])
+
+    ```python
+    |     retriever = BM25Retriever(document_store=document_store)
+    |     ranker = SentenceTransformersRanker(model_name_or_path="cross-encoder/ms-marco-MiniLM-L-12-v2")
+    |     p = Pipeline()
+    |     p.add_node(component=retriever, name="ESRetriever", inputs=["Query"])
+    |     p.add_node(component=ranker, name="Ranker", inputs=["ESRetriever"])
+    ```
     """
 
     def __init__(

{"golden_diff": "diff --git a/haystack/nodes/ranker/sentence_transformers.py b/haystack/nodes/ranker/sentence_transformers.py\n--- a/haystack/nodes/ranker/sentence_transformers.py\n+++ b/haystack/nodes/ranker/sentence_transformers.py\n@@ -28,12 +28,14 @@\n      - directly get predictions via predict()\n \n     Usage example:\n-    ...\n-    retriever = BM25Retriever(document_store=document_store)\n-    ranker = SentenceTransformersRanker(model_name_or_path=\"cross-encoder/ms-marco-MiniLM-L-12-v2\")\n-    p = Pipeline()\n-    p.add_node(component=retriever, name=\"ESRetriever\", inputs=[\"Query\"])\n-    p.add_node(component=ranker, name=\"Ranker\", inputs=[\"ESRetriever\"])\n+\n+    ```python\n+    |     retriever = BM25Retriever(document_store=document_store)\n+    |     ranker = SentenceTransformersRanker(model_name_or_path=\"cross-encoder/ms-marco-MiniLM-L-12-v2\")\n+    |     p = Pipeline()\n+    |     p.add_node(component=retriever, name=\"ESRetriever\", inputs=[\"Query\"])\n+    |     p.add_node(component=ranker, name=\"Ranker\", inputs=[\"ESRetriever\"])\n+    ```\n     \"\"\"\n \n     def __init__(\n", "issue": "Code snippet wrongly formatted in docs\n**Describe the bug**\r\nThe code snippet at:\r\n\r\nhttps://github.com/deepset-ai/haystack/blob/master/haystack/nodes/ranker/sentence_transformers.py#L32\r\n\r\nis rendered as plain text in the corresponding documentation page\r\n\r\nhttps://haystack.deepset.ai/reference/ranker\r\n\r\n\r\n**Expected behavior**\r\nThe snipped should be formatted as Python code\r\n\n", "before_files": [{"content": "from typing import List, Optional, Union, Tuple, Iterator, Any\nimport logging\nfrom pathlib import Path\n\nimport torch\nfrom torch.nn import DataParallel\nfrom transformers import AutoModelForSequenceClassification, AutoTokenizer\n\nfrom haystack.errors import HaystackError\nfrom haystack.schema import Document\nfrom haystack.nodes.ranker.base import BaseRanker\nfrom haystack.modeling.utils import initialize_device_settings\n\nlogger = logging.getLogger(__name__)\n\n\nclass SentenceTransformersRanker(BaseRanker):\n    \"\"\"\n    Sentence Transformer based pre-trained Cross-Encoder model for Document Re-ranking (https://huggingface.co/cross-encoder).\n    Re-Ranking can be used on top of a retriever to boost the performance for document search. This is particularly useful if the retriever has a high recall but is bad in sorting the documents by relevance.\n\n    SentenceTransformerRanker handles Cross-Encoder models\n        - use a single logit as similarity score e.g.  cross-encoder/ms-marco-MiniLM-L-12-v2\n        - use two output logits (no_answer, has_answer) e.g. deepset/gbert-base-germandpr-reranking\n    https://www.sbert.net/docs/pretrained-models/ce-msmarco.html#usage-with-transformers\n\n    |  With a SentenceTransformersRanker, you can:\n     - directly get predictions via predict()\n\n    Usage example:\n    ...\n    retriever = BM25Retriever(document_store=document_store)\n    ranker = SentenceTransformersRanker(model_name_or_path=\"cross-encoder/ms-marco-MiniLM-L-12-v2\")\n    p = Pipeline()\n    p.add_node(component=retriever, name=\"ESRetriever\", inputs=[\"Query\"])\n    p.add_node(component=ranker, name=\"Ranker\", inputs=[\"ESRetriever\"])\n    \"\"\"\n\n    def __init__(\n        self,\n        model_name_or_path: Union[str, Path],\n        model_version: Optional[str] = None,\n        top_k: int = 10,\n        use_gpu: bool = True,\n        devices: Optional[List[Union[str, torch.device]]] = None,\n        batch_size: Optional[int] = None,\n        scale_score: bool = True,\n    ):\n        \"\"\"\n        :param model_name_or_path: Directory of a saved model or the name of a public model e.g.\n        'cross-encoder/ms-marco-MiniLM-L-12-v2'.\n        See https://huggingface.co/cross-encoder for full list of available models\n        :param model_version: The version of model to use from the HuggingFace model hub. Can be tag name, branch name, or commit hash.\n        :param top_k: The maximum number of documents to return\n        :param use_gpu: Whether to use all available GPUs or the CPU. Falls back on CPU if no GPU is available.\n        :param devices: List of GPU (or CPU) devices, to limit inference to certain GPUs and not use all available ones\n                        The strings will be converted into pytorch devices, so use the string notation described here:\n                        https://pytorch.org/docs/stable/tensor_attributes.html?highlight=torch%20device#torch.torch.device\n                        (e.g. [\"cuda:0\"]).\n        :param batch_size: Number of documents to process at a time.\n        :param scale_score: The raw predictions will be transformed using a Sigmoid activation function in case the model\n                            only predicts a single label. For multi-label predictions, no scaling is applied. Set this\n                            to False if you do not want any scaling of the raw predictions.\n        \"\"\"\n        super().__init__()\n\n        self.top_k = top_k\n\n        if devices is not None:\n            self.devices = [torch.device(device) for device in devices]\n        else:\n            self.devices, _ = initialize_device_settings(use_cuda=use_gpu, multi_gpu=True)\n\n        self.transformer_model = AutoModelForSequenceClassification.from_pretrained(\n            pretrained_model_name_or_path=model_name_or_path, revision=model_version\n        )\n        self.transformer_model.to(str(self.devices[0]))\n        self.transformer_tokenizer = AutoTokenizer.from_pretrained(\n            pretrained_model_name_or_path=model_name_or_path, revision=model_version\n        )\n        self.transformer_model.eval()\n\n        # we use sigmoid activation function to scale the score in case there is only a single label\n        # we do not apply any scaling when scale_score is set to False\n        num_labels = self.transformer_model.num_labels\n        self.activation_function: torch.nn.Module\n        if num_labels == 1 and scale_score:\n            self.activation_function = torch.nn.Sigmoid()\n        else:\n            self.activation_function = torch.nn.Identity()\n\n        if len(self.devices) > 1:\n            self.model = DataParallel(self.transformer_model, device_ids=self.devices)\n\n        self.batch_size = batch_size\n\n    def predict(self, query: str, documents: List[Document], top_k: Optional[int] = None) -> List[Document]:\n        \"\"\"\n        Use loaded ranker model to re-rank the supplied list of Document.\n\n        Returns list of Document sorted by (desc.) similarity with the query.\n\n        :param query: Query string\n        :param documents: List of Document to be re-ranked\n        :param top_k: The maximum number of documents to return\n        :return: List of Document\n        \"\"\"\n        if top_k is None:\n            top_k = self.top_k\n\n        features = self.transformer_tokenizer(\n            [query for doc in documents],\n            [doc.content for doc in documents],\n            padding=True,\n            truncation=True,\n            return_tensors=\"pt\",\n        ).to(self.devices[0])\n\n        # SentenceTransformerRanker uses:\n        # 1. the logit as similarity score/answerable classification\n        # 2. the logits as answerable classification  (no_answer / has_answer)\n        # https://www.sbert.net/docs/pretrained-models/ce-msmarco.html#usage-with-transformers\n        with torch.no_grad():\n            similarity_scores = self.transformer_model(**features).logits\n\n        logits_dim = similarity_scores.shape[1]  # [batch_size, logits_dim]\n        sorted_scores_and_documents = sorted(\n            zip(similarity_scores, documents),\n            key=lambda similarity_document_tuple:\n            # assume the last element in logits represents the `has_answer` label\n            similarity_document_tuple[0][-1] if logits_dim >= 2 else similarity_document_tuple[0],\n            reverse=True,\n        )\n\n        # add normalized scores to documents\n        sorted_documents = self._add_scores_to_documents(sorted_scores_and_documents[:top_k], logits_dim)\n\n        return sorted_documents\n\n    def _add_scores_to_documents(\n        self, sorted_scores_and_documents: List[Tuple[Any, Document]], logits_dim: int\n    ) -> List[Document]:\n        \"\"\"\n        Normalize and add scores to retrieved result documents.\n\n        :param sorted_scores_and_documents: List of score, Document Tuples.\n        :param logits_dim: Dimensionality of the returned scores.\n        \"\"\"\n        sorted_documents = []\n        for raw_score, doc in sorted_scores_and_documents:\n            if logits_dim >= 2:\n                score = self.activation_function(raw_score)[-1]\n            else:\n                score = self.activation_function(raw_score)[0]\n\n            doc.score = score.detach().cpu().numpy().tolist()\n            sorted_documents.append(doc)\n\n        return sorted_documents\n\n    def predict_batch(\n        self,\n        queries: List[str],\n        documents: Union[List[Document], List[List[Document]]],\n        top_k: Optional[int] = None,\n        batch_size: Optional[int] = None,\n    ) -> Union[List[Document], List[List[Document]]]:\n        \"\"\"\n        Use loaded ranker model to re-rank the supplied lists of Documents.\n\n        Returns lists of Documents sorted by (desc.) similarity with the corresponding queries.\n\n\n        - If you provide a list containing a single query...\n\n            - ... and a single list of Documents, the single list of Documents will be re-ranked based on the\n              supplied query.\n            - ... and a list of lists of Documents, each list of Documents will be re-ranked individually based on the\n              supplied query.\n\n\n        - If you provide a list of multiple queries...\n\n            - ... you need to provide a list of lists of Documents. Each list of Documents will be re-ranked based on\n              its corresponding query.\n\n        :param queries: Single query string or list of queries\n        :param documents: Single list of Documents or list of lists of Documents to be reranked.\n        :param top_k: The maximum number of documents to return per Document list.\n        :param batch_size: Number of Documents to process at a time.\n        \"\"\"\n        if top_k is None:\n            top_k = self.top_k\n\n        if batch_size is None:\n            batch_size = self.batch_size\n\n        number_of_docs, all_queries, all_docs, single_list_of_docs = self._preprocess_batch_queries_and_docs(\n            queries=queries, documents=documents\n        )\n\n        batches = self._get_batches(all_queries=all_queries, all_docs=all_docs, batch_size=batch_size)\n        preds = []\n        for cur_queries, cur_docs in batches:\n            features = self.transformer_tokenizer(\n                cur_queries, [doc.content for doc in cur_docs], padding=True, truncation=True, return_tensors=\"pt\"\n            ).to(self.devices[0])\n\n            with torch.no_grad():\n                similarity_scores = self.transformer_model(**features).logits\n                preds.extend(similarity_scores)\n\n        logits_dim = similarity_scores.shape[1]  # [batch_size, logits_dim]\n        if single_list_of_docs:\n            sorted_scores_and_documents = sorted(\n                zip(similarity_scores, documents),\n                key=lambda similarity_document_tuple:\n                # assume the last element in logits represents the `has_answer` label\n                similarity_document_tuple[0][-1] if logits_dim >= 2 else similarity_document_tuple[0],\n                reverse=True,\n            )\n\n            # is this step needed?\n            sorted_documents = [(score, doc) for score, doc in sorted_scores_and_documents if isinstance(doc, Document)]\n            sorted_documents_with_scores = self._add_scores_to_documents(sorted_documents[:top_k], logits_dim)\n\n            return sorted_documents_with_scores\n        else:\n            # Group predictions together\n            grouped_predictions = []\n            left_idx = 0\n            right_idx = 0\n            for number in number_of_docs:\n                right_idx = left_idx + number\n                grouped_predictions.append(similarity_scores[left_idx:right_idx])\n                left_idx = right_idx\n\n            result = []\n            for pred_group, doc_group in zip(grouped_predictions, documents):\n                sorted_scores_and_documents = sorted(\n                    zip(pred_group, doc_group),  # type: ignore\n                    key=lambda similarity_document_tuple:\n                    # assume the last element in logits represents the `has_answer` label\n                    similarity_document_tuple[0][-1] if logits_dim >= 2 else similarity_document_tuple[0],\n                    reverse=True,\n                )\n\n                # rank documents according to scores\n                sorted_documents = [\n                    (score, doc) for score, doc in sorted_scores_and_documents if isinstance(doc, Document)\n                ]\n                sorted_documents_with_scores = self._add_scores_to_documents(sorted_documents[:top_k], logits_dim)\n\n                result.append(sorted_documents_with_scores)\n\n            return result\n\n    def _preprocess_batch_queries_and_docs(\n        self, queries: List[str], documents: Union[List[Document], List[List[Document]]]\n    ) -> Tuple[List[int], List[str], List[Document], bool]:\n        number_of_docs = []\n        all_queries = []\n        all_docs: List[Document] = []\n        single_list_of_docs = False\n\n        # Docs case 1: single list of Documents -> rerank single list of Documents based on single query\n        if len(documents) > 0 and isinstance(documents[0], Document):\n            if len(queries) != 1:\n                raise HaystackError(\"Number of queries must be 1 if a single list of Documents is provided.\")\n            query = queries[0]\n            number_of_docs = [len(documents)]\n            all_queries = [query] * len(documents)\n            all_docs = documents  # type: ignore\n            single_list_of_docs = True\n\n        # Docs case 2: list of lists of Documents -> rerank each list of Documents based on corresponding query\n        # If queries contains a single query, apply it to each list of Documents\n        if len(documents) > 0 and isinstance(documents[0], list):\n            if len(queries) == 1:\n                queries = queries * len(documents)\n            if len(queries) != len(documents):\n                raise HaystackError(\"Number of queries must be equal to number of provided Document lists.\")\n            for query, cur_docs in zip(queries, documents):\n                if not isinstance(cur_docs, list):\n                    raise HaystackError(f\"cur_docs was of type {type(cur_docs)}, but expected a list of Documents.\")\n                number_of_docs.append(len(cur_docs))\n                all_queries.extend([query] * len(cur_docs))\n                all_docs.extend(cur_docs)\n\n        return number_of_docs, all_queries, all_docs, single_list_of_docs\n\n    @staticmethod\n    def _get_batches(\n        all_queries: List[str], all_docs: List[Document], batch_size: Optional[int]\n    ) -> Iterator[Tuple[List[str], List[Document]]]:\n        if batch_size is None:\n            yield all_queries, all_docs\n            return\n        else:\n            for index in range(0, len(all_queries), batch_size):\n                yield all_queries[index : index + batch_size], all_docs[index : index + batch_size]\n", "path": "haystack/nodes/ranker/sentence_transformers.py"}], "after_files": [{"content": "from typing import List, Optional, Union, Tuple, Iterator, Any\nimport logging\nfrom pathlib import Path\n\nimport torch\nfrom torch.nn import DataParallel\nfrom transformers import AutoModelForSequenceClassification, AutoTokenizer\n\nfrom haystack.errors import HaystackError\nfrom haystack.schema import Document\nfrom haystack.nodes.ranker.base import BaseRanker\nfrom haystack.modeling.utils import initialize_device_settings\n\nlogger = logging.getLogger(__name__)\n\n\nclass SentenceTransformersRanker(BaseRanker):\n    \"\"\"\n    Sentence Transformer based pre-trained Cross-Encoder model for Document Re-ranking (https://huggingface.co/cross-encoder).\n    Re-Ranking can be used on top of a retriever to boost the performance for document search. This is particularly useful if the retriever has a high recall but is bad in sorting the documents by relevance.\n\n    SentenceTransformerRanker handles Cross-Encoder models\n        - use a single logit as similarity score e.g.  cross-encoder/ms-marco-MiniLM-L-12-v2\n        - use two output logits (no_answer, has_answer) e.g. deepset/gbert-base-germandpr-reranking\n    https://www.sbert.net/docs/pretrained-models/ce-msmarco.html#usage-with-transformers\n\n    |  With a SentenceTransformersRanker, you can:\n     - directly get predictions via predict()\n\n    Usage example:\n\n    ```python\n    |     retriever = BM25Retriever(document_store=document_store)\n    |     ranker = SentenceTransformersRanker(model_name_or_path=\"cross-encoder/ms-marco-MiniLM-L-12-v2\")\n    |     p = Pipeline()\n    |     p.add_node(component=retriever, name=\"ESRetriever\", inputs=[\"Query\"])\n    |     p.add_node(component=ranker, name=\"Ranker\", inputs=[\"ESRetriever\"])\n    ```\n    \"\"\"\n\n    def __init__(\n        self,\n        model_name_or_path: Union[str, Path],\n        model_version: Optional[str] = None,\n        top_k: int = 10,\n        use_gpu: bool = True,\n        devices: Optional[List[Union[str, torch.device]]] = None,\n        batch_size: Optional[int] = None,\n        scale_score: bool = True,\n    ):\n        \"\"\"\n        :param model_name_or_path: Directory of a saved model or the name of a public model e.g.\n        'cross-encoder/ms-marco-MiniLM-L-12-v2'.\n        See https://huggingface.co/cross-encoder for full list of available models\n        :param model_version: The version of model to use from the HuggingFace model hub. Can be tag name, branch name, or commit hash.\n        :param top_k: The maximum number of documents to return\n        :param use_gpu: Whether to use all available GPUs or the CPU. Falls back on CPU if no GPU is available.\n        :param devices: List of GPU (or CPU) devices, to limit inference to certain GPUs and not use all available ones\n                        The strings will be converted into pytorch devices, so use the string notation described here:\n                        https://pytorch.org/docs/stable/tensor_attributes.html?highlight=torch%20device#torch.torch.device\n                        (e.g. [\"cuda:0\"]).\n        :param batch_size: Number of documents to process at a time.\n        :param scale_score: The raw predictions will be transformed using a Sigmoid activation function in case the model\n                            only predicts a single label. For multi-label predictions, no scaling is applied. Set this\n                            to False if you do not want any scaling of the raw predictions.\n        \"\"\"\n        super().__init__()\n\n        self.top_k = top_k\n\n        if devices is not None:\n            self.devices = [torch.device(device) for device in devices]\n        else:\n            self.devices, _ = initialize_device_settings(use_cuda=use_gpu, multi_gpu=True)\n\n        self.transformer_model = AutoModelForSequenceClassification.from_pretrained(\n            pretrained_model_name_or_path=model_name_or_path, revision=model_version\n        )\n        self.transformer_model.to(str(self.devices[0]))\n        self.transformer_tokenizer = AutoTokenizer.from_pretrained(\n            pretrained_model_name_or_path=model_name_or_path, revision=model_version\n        )\n        self.transformer_model.eval()\n\n        # we use sigmoid activation function to scale the score in case there is only a single label\n        # we do not apply any scaling when scale_score is set to False\n        num_labels = self.transformer_model.num_labels\n        self.activation_function: torch.nn.Module\n        if num_labels == 1 and scale_score:\n            self.activation_function = torch.nn.Sigmoid()\n        else:\n            self.activation_function = torch.nn.Identity()\n\n        if len(self.devices) > 1:\n            self.model = DataParallel(self.transformer_model, device_ids=self.devices)\n\n        self.batch_size = batch_size\n\n    def predict(self, query: str, documents: List[Document], top_k: Optional[int] = None) -> List[Document]:\n        \"\"\"\n        Use loaded ranker model to re-rank the supplied list of Document.\n\n        Returns list of Document sorted by (desc.) similarity with the query.\n\n        :param query: Query string\n        :param documents: List of Document to be re-ranked\n        :param top_k: The maximum number of documents to return\n        :return: List of Document\n        \"\"\"\n        if top_k is None:\n            top_k = self.top_k\n\n        features = self.transformer_tokenizer(\n            [query for doc in documents],\n            [doc.content for doc in documents],\n            padding=True,\n            truncation=True,\n            return_tensors=\"pt\",\n        ).to(self.devices[0])\n\n        # SentenceTransformerRanker uses:\n        # 1. the logit as similarity score/answerable classification\n        # 2. the logits as answerable classification  (no_answer / has_answer)\n        # https://www.sbert.net/docs/pretrained-models/ce-msmarco.html#usage-with-transformers\n        with torch.no_grad():\n            similarity_scores = self.transformer_model(**features).logits\n\n        logits_dim = similarity_scores.shape[1]  # [batch_size, logits_dim]\n        sorted_scores_and_documents = sorted(\n            zip(similarity_scores, documents),\n            key=lambda similarity_document_tuple:\n            # assume the last element in logits represents the `has_answer` label\n            similarity_document_tuple[0][-1] if logits_dim >= 2 else similarity_document_tuple[0],\n            reverse=True,\n        )\n\n        # add normalized scores to documents\n        sorted_documents = self._add_scores_to_documents(sorted_scores_and_documents[:top_k], logits_dim)\n\n        return sorted_documents\n\n    def _add_scores_to_documents(\n        self, sorted_scores_and_documents: List[Tuple[Any, Document]], logits_dim: int\n    ) -> List[Document]:\n        \"\"\"\n        Normalize and add scores to retrieved result documents.\n\n        :param sorted_scores_and_documents: List of score, Document Tuples.\n        :param logits_dim: Dimensionality of the returned scores.\n        \"\"\"\n        sorted_documents = []\n        for raw_score, doc in sorted_scores_and_documents:\n            if logits_dim >= 2:\n                score = self.activation_function(raw_score)[-1]\n            else:\n                score = self.activation_function(raw_score)[0]\n\n            doc.score = score.detach().cpu().numpy().tolist()\n            sorted_documents.append(doc)\n\n        return sorted_documents\n\n    def predict_batch(\n        self,\n        queries: List[str],\n        documents: Union[List[Document], List[List[Document]]],\n        top_k: Optional[int] = None,\n        batch_size: Optional[int] = None,\n    ) -> Union[List[Document], List[List[Document]]]:\n        \"\"\"\n        Use loaded ranker model to re-rank the supplied lists of Documents.\n\n        Returns lists of Documents sorted by (desc.) similarity with the corresponding queries.\n\n\n        - If you provide a list containing a single query...\n\n            - ... and a single list of Documents, the single list of Documents will be re-ranked based on the\n              supplied query.\n            - ... and a list of lists of Documents, each list of Documents will be re-ranked individually based on the\n              supplied query.\n\n\n        - If you provide a list of multiple queries...\n\n            - ... you need to provide a list of lists of Documents. Each list of Documents will be re-ranked based on\n              its corresponding query.\n\n        :param queries: Single query string or list of queries\n        :param documents: Single list of Documents or list of lists of Documents to be reranked.\n        :param top_k: The maximum number of documents to return per Document list.\n        :param batch_size: Number of Documents to process at a time.\n        \"\"\"\n        if top_k is None:\n            top_k = self.top_k\n\n        if batch_size is None:\n            batch_size = self.batch_size\n\n        number_of_docs, all_queries, all_docs, single_list_of_docs = self._preprocess_batch_queries_and_docs(\n            queries=queries, documents=documents\n        )\n\n        batches = self._get_batches(all_queries=all_queries, all_docs=all_docs, batch_size=batch_size)\n        preds = []\n        for cur_queries, cur_docs in batches:\n            features = self.transformer_tokenizer(\n                cur_queries, [doc.content for doc in cur_docs], padding=True, truncation=True, return_tensors=\"pt\"\n            ).to(self.devices[0])\n\n            with torch.no_grad():\n                similarity_scores = self.transformer_model(**features).logits\n                preds.extend(similarity_scores)\n\n        logits_dim = similarity_scores.shape[1]  # [batch_size, logits_dim]\n        if single_list_of_docs:\n            sorted_scores_and_documents = sorted(\n                zip(similarity_scores, documents),\n                key=lambda similarity_document_tuple:\n                # assume the last element in logits represents the `has_answer` label\n                similarity_document_tuple[0][-1] if logits_dim >= 2 else similarity_document_tuple[0],\n                reverse=True,\n            )\n\n            # is this step needed?\n            sorted_documents = [(score, doc) for score, doc in sorted_scores_and_documents if isinstance(doc, Document)]\n            sorted_documents_with_scores = self._add_scores_to_documents(sorted_documents[:top_k], logits_dim)\n\n            return sorted_documents_with_scores\n        else:\n            # Group predictions together\n            grouped_predictions = []\n            left_idx = 0\n            right_idx = 0\n            for number in number_of_docs:\n                right_idx = left_idx + number\n                grouped_predictions.append(similarity_scores[left_idx:right_idx])\n                left_idx = right_idx\n\n            result = []\n            for pred_group, doc_group in zip(grouped_predictions, documents):\n                sorted_scores_and_documents = sorted(\n                    zip(pred_group, doc_group),  # type: ignore\n                    key=lambda similarity_document_tuple:\n                    # assume the last element in logits represents the `has_answer` label\n                    similarity_document_tuple[0][-1] if logits_dim >= 2 else similarity_document_tuple[0],\n                    reverse=True,\n                )\n\n                # rank documents according to scores\n                sorted_documents = [\n                    (score, doc) for score, doc in sorted_scores_and_documents if isinstance(doc, Document)\n                ]\n                sorted_documents_with_scores = self._add_scores_to_documents(sorted_documents[:top_k], logits_dim)\n\n                result.append(sorted_documents_with_scores)\n\n            return result\n\n    def _preprocess_batch_queries_and_docs(\n        self, queries: List[str], documents: Union[List[Document], List[List[Document]]]\n    ) -> Tuple[List[int], List[str], List[Document], bool]:\n        number_of_docs = []\n        all_queries = []\n        all_docs: List[Document] = []\n        single_list_of_docs = False\n\n        # Docs case 1: single list of Documents -> rerank single list of Documents based on single query\n        if len(documents) > 0 and isinstance(documents[0], Document):\n            if len(queries) != 1:\n                raise HaystackError(\"Number of queries must be 1 if a single list of Documents is provided.\")\n            query = queries[0]\n            number_of_docs = [len(documents)]\n            all_queries = [query] * len(documents)\n            all_docs = documents  # type: ignore\n            single_list_of_docs = True\n\n        # Docs case 2: list of lists of Documents -> rerank each list of Documents based on corresponding query\n        # If queries contains a single query, apply it to each list of Documents\n        if len(documents) > 0 and isinstance(documents[0], list):\n            if len(queries) == 1:\n                queries = queries * len(documents)\n            if len(queries) != len(documents):\n                raise HaystackError(\"Number of queries must be equal to number of provided Document lists.\")\n            for query, cur_docs in zip(queries, documents):\n                if not isinstance(cur_docs, list):\n                    raise HaystackError(f\"cur_docs was of type {type(cur_docs)}, but expected a list of Documents.\")\n                number_of_docs.append(len(cur_docs))\n                all_queries.extend([query] * len(cur_docs))\n                all_docs.extend(cur_docs)\n\n        return number_of_docs, all_queries, all_docs, single_list_of_docs\n\n    @staticmethod\n    def _get_batches(\n        all_queries: List[str], all_docs: List[Document], batch_size: Optional[int]\n    ) -> Iterator[Tuple[List[str], List[Document]]]:\n        if batch_size is None:\n            yield all_queries, all_docs\n            return\n        else:\n            for index in range(0, len(all_queries), batch_size):\n                yield all_queries[index : index + batch_size], all_docs[index : index + batch_size]\n", "path": "haystack/nodes/ranker/sentence_transformers.py"}]}