/pyterrier_doc2query

Primary LanguageJupyter Notebook

PyTerrier_doc2query

New: Check out our interactive demo on 🤗 HuggingFace Spaces

New: Improve effectiveness and efficiency using Doc2Query−−

This is the PyTerrier plugin for the docTTTTTquery [Nogueira20] and Doc2Query−− [Gospodinov23] approaches for document expansion by query prediction.

Installation

This repostory can be installed using Pip.

pip install --upgrade git+https://github.com/terrierteam/pyterrier_doc2query.git

What does it do?

A Doc2Query object has a transform() function, which takes the text of each document, and suggests questions for that text.

sample_doc = "The presence of communication amid scientific minds was equally important to the success of the Manhattan Project as scientific intellect was. The only cloud hanging over the impressive achievement of the atomic researchers and engineers is what their success truly meant; hundreds of thousands of innocent lives obliterated"

import pyterrier_doc2query
doc2query = pyterrier_doc2query.Doc2Query()
doc2query([{"docno" : "d1", "text" : sample_doc}])

The resulting dataframe will have an additional "querygen" column, which contains the generated queries, such as:

docno querygen
"d1" 'what was the importance of the manhattan project to the united states atom project? what influenced the success of the united states why was the manhattan project a success? why was it important'

As a PyTerrier transformer, there are lots of ways to introduce Doc2query into a PyTerrier retrieval process.

By default, the plugin loads macavaney/doc2query-t5-base-msmarco, which is a a version of the checkpoint released by the original authors, converted to pytorch format. You can load another T5 model by passing another huggingface model name (or path to model on the file system) by passing it as the first argument:

doc2query = pyterrier_doc2query.Doc2Query('some/other/model')

Using Doc2Query for Indexing

Then, indexing is as easy as instantiating the Doc2Query object and a PyTerrier indexer:

import pyterrier as pt
pt.init()
dataset = pt.get_dataset("irds:vaswani")
import pyterrier_doc2query
doc2query = pyterrier_doc2query.Doc2Query(append=True) # append generated queries to the orignal document text
indexer = doc2query >> pt.IterDictIndexer(index_loc)
indexer.index(dataset.get_corpus_iter())

Doc2Query−−: When Less is More

The performance of Doc2Query can be significantly improved by removing queries that are not relevant to the documents that generated them. This involves first scoring the generated queries (using QueryScorer) and then filtering out the least relevant ones (using QueryFilter).

from pyterrier_doc2query import Doc2Query, QueryScorer, QueryFilter
from pyterrier_dr import ElectraScorer

doc2query = Doc2Query(append=False, num_samples=5)
scorer = ElectraScorer()
indexer = pt.IterDictIndexer('./index')
pipeline = doc2query >> QueryScorer(scorer) >> QueryFilter(t=3.21484375) >> indexer # t=3.21484375 is the 70th percentile for generated queries on MS MARCO

pipeline.index(dataset.get_corpus_iter())

We've also released pre-computed filter scores for various models on HuggingFace datasets:

Using Doc2Query for Retrieval

Doc2query can also be used at retrieval time (i.e. on retrieved documents) rather than at indexing time.

import pyterrier_doc2query
doc2query = pyterrier_doc2query.Doc2Query()

dataset = pt.get_dataset("irds:vaswani")
bm25 = pt.terrier.Retriever.from_dataset("vaswani", "terrier_stemmed", wmodel="BM25")
bm25 >> pt.get_text(dataset) >> doc2query >> pt.text.scorer(body_attr="querygen", wmodel="BM25")

Examples

Check out out the notebooks, even on Colab:

Implementation Details

We use a PyTerrier transformer to rewrite documents by doc2query.

References

Credits

  • Craig Macdonald, University of Glasgow
  • Sean MacAvaney, University of Glasgow
  • Mitko Gospodinov, University of Glasgow