# SPDX-FileCopyrightText: 2022-present deepset GmbH # # SPDX-License-Identifier: Apache-2.0 from typing import Any from haystack import Document, component, logging from haystack.lazy_imports import LazyImport logger = logging.getLogger(__name__) with LazyImport(message="Run 'pip install \"torch>=1.13\"'") as torch_import: import torch @component class TopPSampler: """ Implements top-p (nucleus) sampling for document filtering based on cumulative probability scores. This component provides functionality to filter a list of documents by selecting those whose scores fall within the top 'p' percent of the cumulative distribution. It is useful for focusing on high-probability documents while filtering out less relevant ones based on their assigned scores. Usage example: ```python from haystack import Document from haystack.components.samplers import TopPSampler sampler = TopPSampler(top_p=0.95, score_field="similarity_score") docs = [ Document(content="Berlin", meta={"similarity_score": -10.6}), Document(content="Belgrade", meta={"similarity_score": -8.9}), Document(content="Sarajevo", meta={"similarity_score": -4.6}), ] output = sampler.run(documents=docs) docs = output["documents"] assert len(docs) == 1 assert docs[0].content == "Sarajevo" ``` """ def __init__(self, top_p: float = 1.0, score_field: str | None = None, min_top_k: int | None = None) -> None: """ Creates an instance of TopPSampler. :param top_p: Float between 0 and 1 representing the cumulative probability threshold for document selection. A value of 1.0 indicates no filtering (all documents are retained). :param score_field: Name of the field in each document's metadata that contains the score. If None, the default document score field is used. :param min_top_k: If specified, the minimum number of documents to return. If the top_p selects fewer documents, additional ones with the next highest scores are added to the selection. """ torch_import.check() self.top_p = top_p if not 0 <= top_p <= 1: raise ValueError(f"top_p must be between 0 and 1. Got {top_p}.") self.score_field = score_field self.min_top_k = min_top_k @component.output_types(documents=list[Document]) def run(self, documents: list[Document], top_p: float | None = None) -> dict[str, Any]: """ Filters documents using top-p sampling based on their scores. If the specified top_p results in no documents being selected (especially in cases of a low top_p value), the method returns the document with the highest score. :param documents: List of Document objects to be filtered. :param top_p: If specified, a float to override the cumulative probability threshold set during initialization. :returns: A dictionary with the following key: - `documents`: List of Document objects that have been selected based on the top-p sampling. :raises ValueError: If the top_p value is not within the range [0, 1]. """ if not documents: return {"documents": []} top_p = top_p if top_p is not None else self.top_p if not 0 <= top_p <= 1: raise ValueError(f"top_p must be between 0 and 1. Got {top_p}.") documents_with_scores, scores = self._get_documents_and_scores(documents) if len(documents_with_scores) == 0: logger.warning("No documents with scores found. Returning the original documents.") return {"documents": documents} sorted_docs_with_scores = sorted( zip(documents_with_scores, scores, strict=True), key=lambda x: x[1], reverse=True ) sorted_documents, sorted_scores = [list(t) for t in zip(*sorted_docs_with_scores, strict=True)] tensor_scores = torch.tensor(sorted_scores, dtype=torch.float32) probs = torch.nn.functional.softmax(tensor_scores, dim=-1) cumulative_probs = torch.cumsum(probs, dim=-1) # Check if the cumulative probabilities are close to top_p with a 1e-6 tolerance close_to_top_p = torch.isclose(cumulative_probs, torch.tensor(top_p, device=cumulative_probs.device), atol=1e-6) # Combine the close_to_top_p with original condition using logical OR condition = (cumulative_probs <= top_p) | close_to_top_p # Find the indices with cumulative probabilities that exceed top_p top_p_indices = torch.where(torch.BoolTensor(condition))[0] # Map the selected indices back to their original indices selected_docs = [sorted_documents[i.item()] for i in top_p_indices] if self.min_top_k and len(selected_docs) < self.min_top_k: selected_docs = sorted_documents[: self.min_top_k] # If low p resulted in no documents being selected, then return at least one document if len(selected_docs) == 0: logger.warning( "Top-p sampling with p={top_p} resulted in no documents being selected. " "Returning the document with the highest score.", top_p=top_p, ) selected_docs = [sorted_documents[0]] return {"documents": selected_docs} @staticmethod def _get_doc_score(doc: Document, score_field: str | None = None) -> float | None: """ Get the score of a document. :param doc: Document object. :param score_field: Name of the field in the document's metadata that contains the score. If None, the document score field is used. :return: Score of the document. """ if score_field: score = doc.meta.get(score_field) else: score = doc.score # bool is a subclass of int but is not a valid score if isinstance(score, bool) or not isinstance(score, (int, float)): return None return float(score) def _get_documents_and_scores(self, documents: list[Document]) -> tuple[list[Document], list[float]]: """ Checks if documents have scores in their metadata or score field and returns the documents with scores. :param documents: List of Documents. :return: List of scores. """ docs_with_scores = [] scores = [] docs_missing_scores = [] for doc in documents: score = self._get_doc_score(doc=doc, score_field=self.score_field) if score is None: docs_missing_scores.append(doc) else: scores.append(score) docs_with_scores.append(doc) if len(docs_missing_scores) > 0: missing_scores_docs_ids = [d.id for d in docs_missing_scores if d.id] if self.score_field: logger.warning( "Score field '{score_field}' not found in metadata of documents with IDs: {doc_ids}." "Make sure that all documents have a score field '{score_field_2}' in their metadata.", score_field=self.score_field, doc_ids=",".join(missing_scores_docs_ids), score_field_2=self.score_field, ) else: logger.warning( "Ensure all documents have a valid score value. These documents {doc_ids} are missing scores.", doc_ids=",".join(missing_scores_docs_ids), ) return docs_with_scores, scores