Coverage for haystack/components/samplers/top_p.py: 100%
67 statements
« prev ^ index » next coverage.py v7.16.2, created at 2026-10-12 09:42 +0000
1# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
2#
3# SPDX-License-Identifier: Apache-2.0
6from typing import Any
8from haystack import Document, component, logging
9from haystack.lazy_imports import LazyImport
11logger = logging.getLogger(__name__)
14with LazyImport(message="Run 'pip install \"torch>=1.13\"'") as torch_import:
15 import torch
18@component
19class TopPSampler:
20 """
21 Implements top-p (nucleus) sampling for document filtering based on cumulative probability scores.
23 This component provides functionality to filter a list of documents by selecting those whose scores fall
24 within the top 'p' percent of the cumulative distribution. It is useful for focusing on high-probability
25 documents while filtering out less relevant ones based on their assigned scores.
27 Usage example:
29 ```python
30 from haystack import Document
31 from haystack.components.samplers import TopPSampler
33 sampler = TopPSampler(top_p=0.95, score_field="similarity_score")
34 docs = [
35 Document(content="Berlin", meta={"similarity_score": -10.6}),
36 Document(content="Belgrade", meta={"similarity_score": -8.9}),
37 Document(content="Sarajevo", meta={"similarity_score": -4.6}),
38 ]
39 output = sampler.run(documents=docs)
40 docs = output["documents"]
41 assert len(docs) == 1
42 assert docs[0].content == "Sarajevo"
43 ```
44 """
46 def __init__(self, top_p: float = 1.0, score_field: str | None = None, min_top_k: int | None = None) -> None:
47 """
48 Creates an instance of TopPSampler.
50 :param top_p: Float between 0 and 1 representing the cumulative probability threshold for document selection.
51 A value of 1.0 indicates no filtering (all documents are retained).
52 :param score_field: Name of the field in each document's metadata that contains the score. If None, the default
53 document score field is used.
54 :param min_top_k: Minimum number of scored documents to return. If top-p sampling selects fewer documents,
55 documents with the next-highest scores are added. Must be a non-negative integer or None.
56 If greater than the number of scored documents, all scored documents are returned.
57 :raises ValueError: If top_p is not within [0, 1] or min_top_k is not a non-negative integer or None.
58 """
59 torch_import.check()
61 self.top_p = top_p
62 if not 0 <= top_p <= 1:
63 raise ValueError(f"top_p must be between 0 and 1. Got {top_p}.")
64 if min_top_k is not None and (isinstance(min_top_k, bool) or not isinstance(min_top_k, int) or min_top_k < 0):
65 raise ValueError(f"min_top_k must be a non-negative integer or None. Got {min_top_k!r}.")
66 self.score_field = score_field
67 self.min_top_k = min_top_k
69 @component.output_types(documents=list[Document])
70 def run(self, documents: list[Document], top_p: float | None = None) -> dict[str, Any]:
71 """
72 Filters documents using top-p sampling based on their scores.
74 If the specified top_p results in no documents being selected (especially in cases of a low top_p value), the
75 method returns the document with the highest score.
77 :param documents: List of Document objects to be filtered.
78 :param top_p: If specified, a float to override the cumulative probability threshold set during initialization.
80 :returns: A dictionary with the following key:
81 - `documents`: List of Document objects that have been selected based on the top-p sampling.
82 :raises ValueError: If the top_p value is not within the range [0, 1].
83 """
84 if not documents:
85 return {"documents": []}
87 top_p = top_p if top_p is not None else self.top_p
88 if not 0 <= top_p <= 1:
89 raise ValueError(f"top_p must be between 0 and 1. Got {top_p}.")
91 documents_with_scores, scores = self._get_documents_and_scores(documents)
92 if len(documents_with_scores) == 0:
93 logger.warning("No documents with scores found. Returning the original documents.")
94 return {"documents": documents}
96 sorted_docs_with_scores = sorted(
97 zip(documents_with_scores, scores, strict=True), key=lambda x: x[1], reverse=True
98 )
99 sorted_documents, sorted_scores = [list(t) for t in zip(*sorted_docs_with_scores, strict=True)]
101 tensor_scores = torch.tensor(sorted_scores, dtype=torch.float32)
102 probs = torch.nn.functional.softmax(tensor_scores, dim=-1)
103 cumulative_probs = torch.cumsum(probs, dim=-1)
105 # Check if the cumulative probabilities are close to top_p with a 1e-6 tolerance
106 close_to_top_p = torch.isclose(cumulative_probs, torch.tensor(top_p, device=cumulative_probs.device), atol=1e-6)
108 # Combine the close_to_top_p with original condition using logical OR
109 condition = (cumulative_probs <= top_p) | close_to_top_p
111 # Find the indices with cumulative probabilities that exceed top_p
112 top_p_indices = torch.where(torch.BoolTensor(condition))[0]
114 # Map the selected indices back to their original indices
115 selected_docs = [sorted_documents[i.item()] for i in top_p_indices]
117 if self.min_top_k and len(selected_docs) < self.min_top_k:
118 selected_docs = sorted_documents[: self.min_top_k]
120 # If low p resulted in no documents being selected, then return at least one document
121 if len(selected_docs) == 0:
122 logger.warning(
123 "Top-p sampling with p={top_p} resulted in no documents being selected. "
124 "Returning the document with the highest score.",
125 top_p=top_p,
126 )
127 selected_docs = [sorted_documents[0]]
129 return {"documents": selected_docs}
131 @staticmethod
132 def _get_doc_score(doc: Document, score_field: str | None = None) -> float | None:
133 """
134 Get the score of a document.
136 :param doc: Document object.
137 :param score_field: Name of the field in the document's metadata that contains the score.
138 If None, the document score field is used.
140 :return: Score of the document.
141 """
142 if score_field:
143 score = doc.meta.get(score_field)
144 else:
145 score = doc.score
147 # bool is a subclass of int but is not a valid score
148 if isinstance(score, bool) or not isinstance(score, (int, float)):
149 return None
150 return float(score)
152 def _get_documents_and_scores(self, documents: list[Document]) -> tuple[list[Document], list[float]]:
153 """
154 Checks if documents have scores in their metadata or score field and returns the documents with scores.
156 :param documents: List of Documents.
157 :return: List of scores.
158 """
159 docs_with_scores = []
160 scores = []
161 docs_missing_scores = []
162 for doc in documents:
163 score = self._get_doc_score(doc=doc, score_field=self.score_field)
164 if score is None:
165 docs_missing_scores.append(doc)
166 else:
167 scores.append(score)
168 docs_with_scores.append(doc)
170 if len(docs_missing_scores) > 0:
171 missing_scores_docs_ids = [d.id for d in docs_missing_scores if d.id]
172 if self.score_field:
173 logger.warning(
174 "Score field '{score_field}' not found in metadata of documents with IDs: {doc_ids}."
175 "Make sure that all documents have a score field '{score_field_2}' in their metadata.",
176 score_field=self.score_field,
177 doc_ids=",".join(missing_scores_docs_ids),
178 score_field_2=self.score_field,
179 )
180 else:
181 logger.warning(
182 "Ensure all documents have a valid score value. These documents {doc_ids} are missing scores.",
183 doc_ids=",".join(missing_scores_docs_ids),
184 )
185 return docs_with_scores, scores