diff --git a/pyterrier_dr/__init__.py b/pyterrier_dr/__init__.py index 8291daf..b37eeec 100644 --- a/pyterrier_dr/__init__.py +++ b/pyterrier_dr/__init__.py @@ -12,8 +12,9 @@ from pyterrier_dr.cde import CDE, CDECache from pyterrier_dr.prf import AveragePrf, VectorPrf from pyterrier_dr._ils import ILS, ils +from pyterrier_dr._mmr import MmrScorer __all__ = ["FlexIndex", "DocnoFile", "NilIndex", "NumpyIndex", "RankedLists", "FaissFlat", "FaissHnsw", "MemIndex", "TorchIndex", "BiEncoder", "BiQueryEncoder", "BiDocEncoder", "BiScorer", "HgfBiEncoder", "TasB", "RetroMAE", "SBertBiEncoder", "Ance", "Query2Query", "GTR", "E5", "TctColBert", "ElectraScorer", "BGEM3", "BGEM3QueryEncoder", "BGEM3DocEncoder", "CDE", "CDECache", - "SimFn", "infer_device", "AveragePrf", "VectorPrf", "ILS", "ils"] + "SimFn", "infer_device", "AveragePrf", "VectorPrf", "ILS", "ils", "MmrScorer"] diff --git a/pyterrier_dr/_mmr.py b/pyterrier_dr/_mmr.py new file mode 100644 index 0000000..b9526cb --- /dev/null +++ b/pyterrier_dr/_mmr.py @@ -0,0 +1,66 @@ +import numpy as np +import pandas as pd +import pyterrier as pt +import pyterrier_alpha as pta + + +class MmrScorer(pt.Transformer): + """An MMR (Maximal Marginal Relevance) scorer (i.e., re-ranker). + + The MMR scorer re-orders documents by balancing relevance (from the initial scores) and diversity (based on the + similarity of the document vectors). + + .. cite.dblp:: conf/sigir/CarbonellG98 + """ + def __init__(self, *, Lambda: float = 0.5, norm_rel: bool = False, norm_sim: bool = False, drop_doc_vec: bool = True, verbose: bool = False): + """ + Args: + Lambda: The balance parameter between relevance and diversity (default: 0.5) + norm_rel: Whether to normalize relevance scores to [0, 1] (default: False) + norm_sim: Whether to normalize similarity scores to [0, 1] (default: False) + drop_doc_vec: Whether to drop the 'doc_vec' column after re-ranking (default: True) + verbose: Whether to display verbose output (e.g., progress bars) (default: False) + """ + self.Lambda = Lambda + self.norm_rel = norm_rel + self.norm_sim = norm_sim + self.drop_doc_vec = drop_doc_vec + self.verbose = verbose + + def transform(self, inp: pd.DataFrame) -> pd.DataFrame: + pta.validate.result_frame(inp, extra_columns=['doc_vec']) + out = [] + + it = inp.groupby('qid') + if self.verbose: + it = pt.tqdm(it, unit='q', desc=repr(self)) + + for qid, frame in it: + scores = frame['score'].values + dvec_matrix = np.stack(frame['doc_vec']) + dvec_matrix = dvec_matrix / np.linalg.norm(dvec_matrix, axis=1)[:, None] + dvec_sims = dvec_matrix @ dvec_matrix.T + if self.norm_rel: + scores = (scores - scores.min()) / (scores.max() - scores.min()) + if self.norm_sim: + dvec_sims = (dvec_sims - dvec_sims.min()) / (dvec_sims.max() - dvec_sims.min()) + marg_rels = np.zeros_like(scores) + new_idxs = [] + for _ in range(scores.shape[0]): + mmr_scores = (self.Lambda * scores) - ((1 - self.Lambda) * marg_rels) + idx = mmr_scores.argmax() + new_idxs.append(idx) + if marg_rels.shape[0] > 1: + marg_rels = np.max(np.stack([marg_rels, dvec_sims[idx]]), axis=0) + marg_rels[idx] = float('inf') # ignore this document from now on + new_frame = frame.iloc[new_idxs].reset_index(drop=True).assign( + score=-np.arange(len(new_idxs), dtype=float), + rank=np.arange(len(new_idxs)) + ) + if self.drop_doc_vec: + new_frame = new_frame.drop(columns='doc_vec') + out.append(new_frame) + + return pd.concat(out, ignore_index=True) + + __repr__ = pta.transformer_repr diff --git a/pyterrier_dr/flex/__init__.py b/pyterrier_dr/flex/__init__.py index b736036..1f90d0e 100644 --- a/pyterrier_dr/flex/__init__.py +++ b/pyterrier_dr/flex/__init__.py @@ -8,5 +8,6 @@ from pyterrier_dr.flex import ladr from pyterrier_dr.flex import gar from pyterrier_dr.flex import voyager_retr +from pyterrier_dr.flex import diversity -__all__ = ["FlexIndex", "IndexingMode", "np_retr", "torch_retr", "corpus_graph", "faiss_retr", "flatnav_retr", "scann_retr", "ladr", "gar", "voyager_retr"] +__all__ = ["FlexIndex", "IndexingMode", "np_retr", "torch_retr", "corpus_graph", "faiss_retr", "flatnav_retr", "scann_retr", "ladr", "gar", "voyager_retr", "diversity"] diff --git a/pyterrier_dr/flex/core.py b/pyterrier_dr/flex/core.py index e9db19f..04cfe5e 100644 --- a/pyterrier_dr/flex/core.py +++ b/pyterrier_dr/flex/core.py @@ -196,9 +196,13 @@ def docnos(self) -> Lookup: @property def ILS(self) -> ir_measures.Measure: - """Return an ILS (Intra-List Similarity) measure for this index. See: :func:`pyterrier_dr.ILS` for more details.""" + """Return an ILS (Intra-List Similarity) measure for this index. See :func:`pyterrier_dr.ILS` for more details.""" return pyterrier_dr.ILS(self) + def __repr__(self): + return f'FlexIndex({str(self.index_path)!r})' + + class FlexIndexer(pt.Indexer): def __init__(self, index: FlexIndex, mode: Union[IndexingMode, str] = IndexingMode.create): self._index = index diff --git a/pyterrier_dr/flex/diversity.py b/pyterrier_dr/flex/diversity.py new file mode 100644 index 0000000..66cc287 --- /dev/null +++ b/pyterrier_dr/flex/diversity.py @@ -0,0 +1,22 @@ +import pyterrier as pt +import pyterrier_dr +from . import FlexIndex + + +def _mmr(self, *, Lambda: float = 0.5, norm_rel: bool = False, norm_sim: bool = False, drop_doc_vec: bool = True, verbose: bool = False) -> pt.Transformer: + """Returns an MMR (Maximal Marginal Relevance) scorer (i.e., re-ranker) over this index. + + The method first loads vectors from the index and then applies :class:`MmrScorer` to re-rank the results. See + :class:`MmrScorer` for more details on MMR. + + Args: + Lambda: The balance parameter between relevance and diversity (default: 0.5) + norm_rel: Whether to normalize relevance scores to [0, 1] (default: False) + norm_sim: Whether to normalize similarity scores to [0, 1] (default: False) + drop_doc_vec: Whether to drop the 'doc_vec' column after re-ranking (default: True) + verbose: Whether to display verbose output (e.g., progress bars) (default: False) + + .. cite.dblp:: conf/sigir/CarbonellG98 + """ + return self.vec_loader() >> pyterrier_dr.MmrScorer(Lambda=Lambda, norm_rel=norm_rel, norm_sim=norm_sim, drop_doc_vec=drop_doc_vec, verbose=verbose) +FlexIndex.mmr = _mmr diff --git a/pyterrier_dr/pt_docs/diversity.rst b/pyterrier_dr/pt_docs/diversity.rst index 97b2c51..277d71a 100644 --- a/pyterrier_dr/pt_docs/diversity.rst +++ b/pyterrier_dr/pt_docs/diversity.rst @@ -1,7 +1,19 @@ Diversity ======================================================= -``pyterrier-dr`` provides a diversity evaluation measure, :func:`~pyterrier_dr.ILS` (Intra-List Similarity), +Search Result Diversification +------------------------------------------------------- + +``pyterrier-dr`` provides one diversification algorithm, :class:`~pyterrier_dr.MmrScorer` (Maximal Marginal Relevance). +The transformer works over input dataframes that contain the dense vectors of the documents and the query. You can also +use :meth:`~pyterrier_dr.FlexIndex.mmr` to first load vectors from an index and then apply MMR. + +.. autoclass:: pyterrier_dr.MmrScorer + +Diversity Evaluation +------------------------------------------------------- + +``pyterrier-dr`` provides one diversity evaluation measure, :func:`~pyterrier_dr.ILS` (Intra-List Similarity), which can be used to evaluate the diversity of search results based on the dense vectors of a :class:`~pyterrier_dr.FlexIndex`. This measure can be used alongside PyTerrier's built-in evaluation measures in a :func:`pyterrier.Experiment`. @@ -22,15 +34,17 @@ This measure can be used alongside PyTerrier's built-in evaluation measures in a pt.Experiment( [ bm25, - model >> index, + model >> index.retriever(), + model >> index.retriever() >> index.mmr(), ], dataset.get_topics(), dataset.get_qrels(), [nDCG@10, R(rel=2)@1000, index.ILS@10, index.ILS@1000] ) - # name nDCG@10 R(rel=2)@1000 ILS@10 ILS@1000 - # BM25 0.498902 0.755495 0.852248 0.754691 - # TAS-B 0.716068 0.841756 0.889112 0.775415 + # name nDCG@10 R(rel=2)@1000 ILS@10 ILS@1000 + # BM25 0.498 0.755 0.852 0.754 + # TasB 0.716 0.841 0.889 0.775 + # TasB w/ MMR 0.714 0.841 0.888 0.775 .. autofunction:: pyterrier_dr.ILS .. autofunction:: pyterrier_dr.ils diff --git a/pyterrier_dr/pt_docs/indexing-retrieval.rst b/pyterrier_dr/pt_docs/indexing-retrieval.rst index 13f2286..360e809 100644 --- a/pyterrier_dr/pt_docs/indexing-retrieval.rst +++ b/pyterrier_dr/pt_docs/indexing-retrieval.rst @@ -59,6 +59,12 @@ API Documentation .. automethod:: gar .. automethod:: ladr_proactive .. automethod:: ladr_adaptive + .. automethod:: mmr + + Evaluation + ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + + .. autoproperty:: ILS Index Data Access ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ diff --git a/tests/test_mmr.py b/tests/test_mmr.py new file mode 100644 index 0000000..2d3cd9b --- /dev/null +++ b/tests/test_mmr.py @@ -0,0 +1,31 @@ +import unittest +import numpy as np +import pandas as pd +from pyterrier_dr import MmrScorer + + +class TestMmr(unittest.TestCase): + def test_mmr(self): + mmr = MmrScorer() + results = mmr(pd.DataFrame([ + ['q0', 'd0', 1.0, np.array([0, 1, 0])], + ['q0', 'd1', 0.5, np.array([0, 1, 1])], + ['q0', 'd2', 0.5, np.array([1, 1, 1])], + ['q0', 'd3', 0.1, np.array([1, 1, 0])], + ['q1', 'd0', 0.6, np.array([0, 1, 0])], + ['q2', 'd0', 0.4, np.array([0, 1, 0])], + ['q2', 'd1', 0.3, np.array([0, 1, 1])], + ], columns=['qid', 'docno', 'score', 'doc_vec'])) + pd.testing.assert_frame_equal(results, pd.DataFrame([ + ['q0', 'd0', 0.0, 0], + ['q0', 'd2', -1.0, 1], + ['q0', 'd1', -2.0, 2], + ['q0', 'd3', -3.0, 3], + ['q1', 'd0', 0.0, 0], + ['q2', 'd0', 0.0, 0], + ['q2', 'd1', -1.0, 1], + ], columns=['qid', 'docno', 'score', 'rank'])) + + +if __name__ == '__main__': + unittest.main()