Coverage for haystack/components/samplers/top_p.py: 100%

1# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai> 

2# 

3# SPDX-License-Identifier: Apache-2.0 

4 

5 

6from typing import Any 

7 

8from haystack import Document, component, logging 

9from haystack.lazy_imports import LazyImport 

10 

11logger = logging.getLogger(__name__) 

12 

13 

14with LazyImport(message="Run 'pip install \"torch>=1.13\"'") as torch_import: 

15 import torch 

16 

17 

18@component 

19class TopPSampler: 

20 """ 

21 Implements top-p (nucleus) sampling for document filtering based on cumulative probability scores. 

22 

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. 

26 

27 Usage example: 

28 

29 ```python 

30 from haystack import Document 

31 from haystack.components.samplers import TopPSampler 

32 

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 """ 

45 

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. 

49 

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: If specified, the minimum number of documents to return. If the top_p selects 

55 fewer documents, additional ones with the next highest scores are added to the selection. 

56 """ 

57 torch_import.check() 

58 

59 self.top_p = top_p 

60 if not 0 <= top_p <= 1: 

61 raise ValueError(f"top_p must be between 0 and 1. Got {top_p}.") 

62 self.score_field = score_field 

63 self.min_top_k = min_top_k 

64 

65 @component.output_types(documents=list[Document]) 

66 def run(self, documents: list[Document], top_p: float | None = None) -> dict[str, Any]: 

67 """ 

68 Filters documents using top-p sampling based on their scores. 

69 

70 If the specified top_p results in no documents being selected (especially in cases of a low top_p value), the 

71 method returns the document with the highest score. 

72 

73 :param documents: List of Document objects to be filtered. 

74 :param top_p: If specified, a float to override the cumulative probability threshold set during initialization. 

75 

76 :returns: A dictionary with the following key: 

77 - `documents`: List of Document objects that have been selected based on the top-p sampling. 

78 :raises ValueError: If the top_p value is not within the range [0, 1]. 

79 """ 

80 if not documents: 

81 return {"documents": []} 

82 

83 top_p = top_p if top_p is not None else self.top_p 

84 if not 0 <= top_p <= 1: 

85 raise ValueError(f"top_p must be between 0 and 1. Got {top_p}.") 

86 

87 documents_with_scores, scores = self._get_documents_and_scores(documents) 

88 if len(documents_with_scores) == 0: 

89 logger.warning("No documents with scores found. Returning the original documents.") 

90 return {"documents": documents} 

91 

92 sorted_docs_with_scores = sorted( 

93 zip(documents_with_scores, scores, strict=True), key=lambda x: x[1], reverse=True 

94 ) 

95 sorted_documents, sorted_scores = [list(t) for t in zip(*sorted_docs_with_scores, strict=True)] 

96 

97 tensor_scores = torch.tensor(sorted_scores, dtype=torch.float32) 

98 probs = torch.nn.functional.softmax(tensor_scores, dim=-1) 

99 cumulative_probs = torch.cumsum(probs, dim=-1) 

100 

101 # Check if the cumulative probabilities are close to top_p with a 1e-6 tolerance 

102 close_to_top_p = torch.isclose(cumulative_probs, torch.tensor(top_p, device=cumulative_probs.device), atol=1e-6) 

103 

104 # Combine the close_to_top_p with original condition using logical OR 

105 condition = (cumulative_probs <= top_p) | close_to_top_p 

106 

107 # Find the indices with cumulative probabilities that exceed top_p 

108 top_p_indices = torch.where(torch.BoolTensor(condition))[0] 

109 

110 # Map the selected indices back to their original indices 

111 selected_docs = [sorted_documents[i.item()] for i in top_p_indices] 

112 

113 if self.min_top_k and len(selected_docs) < self.min_top_k: 

114 selected_docs = sorted_documents[: self.min_top_k] 

115 

116 # If low p resulted in no documents being selected, then return at least one document 

117 if len(selected_docs) == 0: 

118 logger.warning( 

119 "Top-p sampling with p={top_p} resulted in no documents being selected. " 

120 "Returning the document with the highest score.", 

121 top_p=top_p, 

122 ) 

123 selected_docs = [sorted_documents[0]] 

124 

125 return {"documents": selected_docs} 

126 

127 @staticmethod 

128 def _get_doc_score(doc: Document, score_field: str | None = None) -> float | None: 

129 """ 

130 Get the score of a document. 

131 

132 :param doc: Document object. 

133 :param score_field: Name of the field in the document's metadata that contains the score. 

134 If None, the document score field is used. 

135 

136 :return: Score of the document. 

137 """ 

138 if score_field: 

139 score = doc.meta.get(score_field) 

140 else: 

141 score = doc.score 

142 

143 # bool is a subclass of int but is not a valid score 

144 if isinstance(score, bool) or not isinstance(score, (int, float)): 

145 return None 

146 return float(score) 

147 

148 def _get_documents_and_scores(self, documents: list[Document]) -> tuple[list[Document], list[float]]: 

149 """ 

150 Checks if documents have scores in their metadata or score field and returns the documents with scores. 

151 

152 :param documents: List of Documents. 

153 :return: List of scores. 

154 """ 

155 docs_with_scores = [] 

156 scores = [] 

157 docs_missing_scores = [] 

158 for doc in documents: 

159 score = self._get_doc_score(doc=doc, score_field=self.score_field) 

160 if score is None: 

161 docs_missing_scores.append(doc) 

162 else: 

163 scores.append(score) 

164 docs_with_scores.append(doc) 

165 

166 if len(docs_missing_scores) > 0: 

167 missing_scores_docs_ids = [d.id for d in docs_missing_scores if d.id] 

168 if self.score_field: 

169 logger.warning( 

170 "Score field '{score_field}' not found in metadata of documents with IDs: {doc_ids}." 

171 "Make sure that all documents have a score field '{score_field_2}' in their metadata.", 

172 score_field=self.score_field, 

173 doc_ids=",".join(missing_scores_docs_ids), 

174 score_field_2=self.score_field, 

175 ) 

176 else: 

177 logger.warning( 

178 "Ensure all documents have a valid score value. These documents {doc_ids} are missing scores.", 

179 doc_ids=",".join(missing_scores_docs_ids), 

180 ) 

181 return docs_with_scores, scores