feat: Added DiversificationResult, added docs (Pringled#5)

* Added docs * Added docs * Added docs * Added docs * Added docs * Made returning gains optional * Added return datamodel * Added return datamodel * Updated docs * Updated docs * Updated docs * Updated docs

Author: Pringled

README.md

@@ -1 +1,63 @@
 # Pyversity — Diversified Re‑Ranking for Retrieval
+
+Pyversity is a small, fast library for diversifying retrieval results.
+Retrieval systems often return highly similar items. Pyversity efficiently re-ranks these results to encourage diversity, surfacing items that remain relevant but less redundant.
+
+It implements several popular strategies such as MMR, MSD, DPP, and Cover with a clear, unified API. More information about the supported strategies can be found in the [supported strategies section](#supported-strategies).
+
+
+## Quickstart
+
+Install `pyversity` with:
+
+```bash
+pip install pyversity
+```
+
+Diversify retrieval results:
+```python
+import numpy as np
+from pyversity import diversify, Strategy
+
+# Define embeddings and scores
+embeddings = np.random.randn(100, 256).astype(np.float32)
+scores = np.random.rand(100).astype(np.float32)
+
+# Diversify with with a chosen strategy (in this case MMR)
+diversified_result = diversify(
+ embeddings=embeddings,
+ scores=scores,
+ k=10,
+ strategy=Strategy.MMR,
+)
+# Get the indicices of the diversified result
+diversified_indices = diversified_result.indices
+```
+
+
+
+## Supported Strategies
+
+The following table describes the supported strategies, how they work, their time complexity, and when to use them.
+
+| Strategy | What It Does | Time Complexity | When to Use |
+| ------------------------------------- | ---------------------------------------------------------------------------------------------- | ------------------------- | ---------------------------------------------------------------------------------------------- |
+| **MMR** (Maximum Marginal Relevance) | Keeps the most relevant items while down-weighting those too similar to what’s already picked. | **O(k · n · d)** | Best **default**. Fast, simple, and works well when you just want to avoid near-duplicates. |
+| **MSD** (Max Sum of Distances) | Prefers items that are both relevant and far from *all* previous selections. | **O(k · n · d)** | Use when you want stronger spread, i.e. results that cover a wider range of topics or styles. |
+| **DPP** (Determinantal Point Process) | Samples diverse yet relevant items using probabilistic “repulsion.” | **O(k · n · d + n · k²)** | Ideal when you want to eliminate redundancy or ensure diversity is built-in to selection. |
+| **COVER** (Facility-Location) | Ensures selected items collectively represent the full dataset’s structure. | **O(k · n²)** | Great for topic coverage or clustering scenarios, but slower for large `n`. |
+
+## References
+
+The implementations in this package are based on the following research papers:
+
+- **MMR**: Carbonell, J., & Goldstein, J. (1998). The use of MMR, diversity-based reranking for reordering documents and producing summaries. [Link](https://dl.acm.org/doi/pdf/10.1145/290941.291025)
+
+- **MSD**: Borodin, A., Lee, H. C., & Ye, Y. (2012). Max-sum diversification, monotone submodular functions and dynamic updates. [Link](https://arxiv.org/pdf/1203.6397)
+
+- **COVER**: Puthiya Parambath, S. A., Usunier, N., & Grandvalet, Y. (2016). A coverage-based approach to recommendation diversity on similarity graph. [Link](https://dl.acm.org/doi/10.1145/2959100.2959149)
+
+- **DPP**: Kulesza, A., & Taskar, B. (2012). Determinantal Point Processes for Machine Learning. [Link](https://arxiv.org/pdf/1207.6083)
+
+- **DPP (efficient greedy implementation)**: Chen, L., Zhang, G., & Zhou, H. (2018). Fast greedy MAP inference for determinantal point process to improve recommendation diversity.
+[Link](https://arxiv.org/pdf/1709.05135)

src/pyversity/__init__.py

@@ -1,4 +1,6 @@
 from pyversity.core import diversify
+from pyversity.datatypes import DiversificationResult, Metric, Strategy
 from pyversity.strategies import cover, dpp, mmr, msd
+from pyversity.version import __version__
 
-__all__ = ["diversify", "mmr", "msd", "cover", "dpp", "__version__"]
+__all__ = ["diversify", "Strategy", "Metric", "DiversificationResult", "mmr", "msd", "cover", "dpp", "__version__"]

src/pyversity/core.py

@@ -2,27 +2,28 @@
 
 import numpy as np
 
-from pyversity.datatypes import Strategy
+from pyversity.datatypes import DiversificationResult, Strategy
 from pyversity.strategies import cover, dpp, mmr, msd
 
 
 def diversify(
- strategy: Strategy,
  embeddings: np.ndarray,
  scores: np.ndarray,
  k: int,
+ strategy: Strategy = Strategy.MMR,
  **kwargs: Any,
-) -> tuple[np.ndarray, np.ndarray]:
+) -> DiversificationResult:
  """
  Diversify a retrieval result using a selected strategy.
 
- :param strategy: The diversification strategy to apply. Supported strategies are: MMR, MSD, COVER, and DPP.
- :param embeddings: Array of embeddings for the items.
- :param scores: Array of relevance scores for the items.
- :param k: The number of items to select in the diversified result.
+ :param embeddings: Embeddings of the items to be diversified.
+ :param scores: Scores (relevances) of the items to be diversified.
+ :param k: The number of items to select for the diversified result.
+ :param strategy: The diversification strategy to apply.
+ Supported strategies are: 'mmr' (default), 'msd', 'cover', and 'dpp'.
  :param **kwargs: Additional keyword arguments passed to the specific strategy function.
- :return: A tuple containing an array of indices of the selected items
- and an array of corresponding relevance scores for the selected items.
+ :return: A DiversificationResult containing the selected item indices,
+ their marginal gains, the strategy used, and the parameters.
  :raises ValueError: If the provided strategy is not recognized.
  """
  if strategy == Strategy.MMR:

src/pyversity/datatypes.py

@@ -1,5 +1,8 @@
+from dataclasses import dataclass
 from enum import Enum
 
+import numpy as np
+
 
 class Strategy(str, Enum):
  """Supported diversification strategies."""
@@ -15,3 +18,23 @@ class Metric(str, Enum):
 
  COSINE = "cosine"
  DOT = "dot"
+
+
+@dataclass
+class DiversificationResult:
+ """
+ Result of a diversification operation.
+
+ Attributes
+ ----------
+ indices: Diversified item indices.
+ marginal_gains: Marginal gains/relevance scores for the diversified items.
+ strategy: Diversification strategy used.
+ parameters: Additional parameters used in the strategy.
+
+ """
+
+ indices: np.ndarray
+ marginal_gains: np.ndarray
+ strategy: Strategy
+ parameters: dict

src/pyversity/strategies/cover.py

@@ -1,6 +1,6 @@
 import numpy as np
 
-from pyversity.datatypes import Metric
+from pyversity.datatypes import DiversificationResult, Metric, Strategy
 from pyversity.utils import normalize_rows, pairwise_similarity, prepare_inputs
 
 
@@ -12,7 +12,7 @@ def cover(
  gamma: float = 0.5,
  metric: Metric = Metric.COSINE,
  normalize: bool = True,
-) -> tuple[np.ndarray, np.ndarray]:
+) -> DiversificationResult:
  """
  Select a subset of items that balances relevance and coverage.
 
@@ -27,7 +27,8 @@ def cover(
  :param gamma: Concavity parameter in (0, 1]; lower values emphasize diversity.
  :param metric: Similarity metric to use. Default is Metric.COSINE.
  :param normalize: Whether to normalize embeddings before computing similarity.
- :return: Tuple of selected indices and their marginal gains.
+ :return: A DiversificationResult containing the selected item indices,
+ their marginal gains, the strategy used, and the parameters.
  :raises ValueError: If theta is not in [0, 1].
  :raises ValueError: If gamma is not in (0, 1].
  """
@@ -37,11 +38,22 @@ def cover(
  if not (0.0 < float(gamma) <= 1.0):
  raise ValueError("gamma must be in (0, 1]")
 
+ params = {
+ "theta": theta,
+ "gamma": gamma,
+ "metric": metric,
+ }
+
  # Prepare inputs
  feature_matrix, relevance_scores, top_k, early_exit = prepare_inputs(embeddings, scores, k)
  if early_exit:
  # Nothing to select: return empty arrays
- return np.empty(0, np.int32), np.empty(0, np.float32)
+ return DiversificationResult(
+ indices=np.empty(0, np.int32),
+ marginal_gains=np.empty(0, np.float32),
+ strategy=Strategy.COVER,
+ parameters=params,
+ )
 
  if metric == Metric.COSINE and normalize:
  # Normalize feature vectors to unit length for cosine similarity
@@ -51,7 +63,12 @@ def cover(
  # Pure relevance: select top-k by relevance scores
  topk = np.argsort(-relevance_scores)[:top_k].astype(np.int32)
  gains = relevance_scores[topk].astype(np.float32, copy=False)
- return topk, gains
+ return DiversificationResult(
+ indices=topk,
+ marginal_gains=gains,
+ strategy=Strategy.COVER,
+ parameters=params,
+ )
 
  # Compute non-negative similarities for coverage to avoid concave-power NaNs
  similarity_matrix = pairwise_similarity(feature_matrix, metric)
@@ -82,4 +99,9 @@ def cover(
  # Update accumulated coverage
  accumulated_coverage += similarity_matrix[:, best_index]
 
- return selected_indices, marginal_gains
+ return DiversificationResult(
+ indices=selected_indices,
+ marginal_gains=marginal_gains,
+ strategy=Strategy.COVER,
+ parameters=params,
+ )

src/pyversity/strategies/dpp.py

@@ -1,5 +1,6 @@
 import numpy as np
 
+from pyversity.datatypes import DiversificationResult, Strategy
 from pyversity.utils import EPS32, normalize_rows, prepare_inputs
 
 
@@ -16,7 +17,7 @@ def dpp(
  scores: np.ndarray,
  k: int,
  beta: float = 1.0,
-) -> tuple[np.ndarray, np.ndarray]:
+) -> DiversificationResult:
  """
  Greedy determinantal point process (DPP) selection.
 
@@ -29,14 +30,19 @@ def dpp(
  :param k: Number of items to select.
  :param beta: Controls the influence of relevance scores in the DPP kernel.
  Higher values increase the emphasis on relevance.
- :return: Tuple of selected indices and their marginal gains.
+ :return: A DiversificationResult containing the selected item indices,
+ their marginal gains, the strategy used, and the parameters.
  """
  # Prepare inputs
  feature_matrix, relevance_scores, top_k, early_exit = prepare_inputs(embeddings, scores, k)
  if early_exit:
  # Nothing to select: return empty arrays
- return np.empty(0, np.int32), np.empty(0, np.float32)
-
+ return DiversificationResult(
+ indices=np.empty(0, np.int32),
+ marginal_gains=np.empty(0, np.float32),
+ strategy=Strategy.DPP,
+ parameters={"beta": beta},
+ )
  # Normalize feature vectors to unit length for cosine similarity
  feature_matrix = normalize_rows(feature_matrix)
 
@@ -87,4 +93,9 @@ def dpp(
  residual_variance -= update_component * update_component
  np.maximum(residual_variance, 0.0, out=residual_variance)
 
- return selected_indices[:step], marginal_gains[:step]
+ return DiversificationResult(
+ indices=selected_indices[:step],
+ marginal_gains=marginal_gains[:step],
+ strategy=Strategy.DPP,
+ parameters={"beta": beta},
+ )

src/pyversity/strategies/mmr.py

@@ -1,6 +1,6 @@
 import numpy as np
 
-from pyversity.datatypes import Metric
+from pyversity.datatypes import DiversificationResult, Metric
 from pyversity.strategies.utils import greedy_select
 
 
@@ -11,7 +11,7 @@ def mmr(
  lambda_param: float = 0.5,
  metric: Metric = Metric.COSINE,
  normalize: bool = True,
-) -> tuple[np.ndarray, np.ndarray]:
+) -> DiversificationResult:
  """
  Maximal Marginal Relevance (MMR) selection.
 
@@ -26,7 +26,8 @@ def mmr(
  1.0 = pure relevance, 0.0 = pure diversity.
  :param metric: Similarity metric to use. Default is Metric.COSINE.
  :param normalize: Whether to normalize embeddings before computing similarity.
- :return: Tuple of selected indices and their marginal gains.
+ :return: A DiversificationResult containing the selected item indices,
+ their marginal gains, the strategy used, and the parameters.
  """
  return greedy_select(
  "mmr",

src/pyversity/strategies/msd.py

@@ -1,6 +1,6 @@
 import numpy as np
 
-from pyversity.datatypes import Metric
+from pyversity.datatypes import DiversificationResult, Metric
 from pyversity.strategies.utils import greedy_select
 
 
@@ -11,7 +11,7 @@ def msd(
  lambda_param: float = 0.5,
  metric: Metric = Metric.COSINE,
  normalize: bool = True,
-) -> tuple[np.ndarray, np.ndarray]:
+) -> DiversificationResult:
  """
  Maximal Sum of Distances (MSD) selection.
 
@@ -27,7 +27,8 @@ def msd(
 
  :param metric: Similarity metric to use. Default is Metric.COSINE.
  :param normalize: Whether to normalize embeddings before computing similarity.
- :return: Tuple of selected indices and their marginal gains.
+ :return: A DiversificationResult containing the selected item indices,
+ their marginal gains, the strategy used, and the parameters.
  """
  return greedy_select(
  "msd",

src/pyversity/strategies/utils.py

@@ -2,7 +2,7 @@
 
 import numpy as np
 
-from pyversity.datatypes import Metric
+from pyversity.datatypes import DiversificationResult, Metric, Strategy
 from pyversity.utils import normalize_rows, prepare_inputs, vector_similarity
 
 
@@ -15,7 +15,7 @@ def greedy_select(
  metric: Metric,
  normalize: bool,
  lambda_param: float,
-) -> tuple[np.ndarray, np.ndarray]:
+) -> DiversificationResult:
  """
  Greedy selection for MMR/MSD strategies.
 
@@ -32,19 +32,30 @@ def greedy_select(
  :param normalize: Whether to normalize embeddings before computing similarity.
  :param lambda_param: Trade-off parameter in [0, 1].
  1.0 = pure relevance, 0.0 = pure diversity.
- :return: Tuple of selected indices and their marginal gains.
+ :return: A DiversificationResult containing the selected item indices,
+ their marginal gains, the strategy used, and the parameters.
  :raises ValueError: If lambda_param is not in [0, 1].
  :raises ValueError: If input shapes are inconsistent.
  """
  # Validate parameters
  if not (0.0 <= float(lambda_param) <= 1.0):
  raise ValueError("lambda_param must be in [0, 1]")
 
+ params = {
+ "lambda_param": lambda_param,
+ "metric": metric,
+ }
+
  # Prepare inputs
  feature_matrix, relevance_scores, top_k, early_exit = prepare_inputs(embeddings, scores, k)
  if early_exit:
  # Nothing to select: return empty arrays
- return np.empty(0, np.int32), np.empty(0, np.float32)
+ return DiversificationResult(
+ indices=np.empty(0, np.int32),
+ marginal_gains=np.empty(0, np.float32),
+ strategy=Strategy.MMR if strategy == "mmr" else Strategy.MSD,
+ parameters=params,
+ )
 
  if metric == Metric.COSINE and normalize:
  # Normalize feature vectors to unit length for cosine similarity
@@ -93,4 +104,9 @@ def greedy_select(
  marginal_gains[step] = float(candidate_scores[best_index])
  selected_mask[best_index] = True
 
- return selected_indices, marginal_gains
+ return DiversificationResult(
+ indices=selected_indices,
+ marginal_gains=marginal_gains,
+ strategy=Strategy.MMR if strategy == "mmr" else Strategy.MSD,
+ parameters=params,
+ )