Package: referral-augment
Referral-augmented retrieval (RAR)InstallationInstall with pip:pip install referral-augmentAlternatively, install from source:git clone https://github.com/michaelwilliamtang/referral-augment
cd referral-augment
pip install -r requirements.txt
pip install -e .OverviewSimple, general implementations of referral-augmented retrieval are provided inrar.retrievers. We support three aggregation methods — concatenation, mean, and shortest path — as described in the paper, which can be specified via anAggregationTypeconstructor argument.Under our framework, retrieval with BM25 is as simple as:fromrar.retrieversimportBM25Retrieverretriever=BM25Retriever(docs,referrals)retriever.retrieve('paper that introduced the Flickr30k dataset',num_docs=10)Similarly, retrieval with any dense embedding model on HuggingFace:fromrar.retrieversimportDenseRetriever,AggregationTypefromrar.encodersimportHuggingFaceEncoderencoder=HuggingFaceEncoder('facebook/contriever')retriever=DenseRetriever(encoder,docs,referrals,aggregation=AggregationType.MEAN)retriever.retrieve('paper that introduced the Flickr30k dataset',num_docs=10)For convenience, we also include direct implementations of theSimCSEEncoderandSpecterEncoder.Example replications of paper results showing the advantage of referral augmentation and demonstrating a full, concise retrieval and evaluation pipeline can be found inexamples.ipynb.Optional: install with SimCSE supportNote that the only stable version of SimCSE is currently using itssourceas a module, which requires buildingrarfrom source. Thus, optionally to install with support forSimCSEEncoder:git clone https://github.com/michaelwilliamtang/referral-augment
cd referral-augment
pip install -r requirements.txt
cd src/rar/encoders
git clone https://github.com/princeton-nlp/SimCSE
cd SimCSE
pip install -r requirements.txt
cd ../../../..DataWe provide sample data in zipped formhere— to use, unzip and placedata/under the repository's root directory.Our sample data covers two domains, each with acorpusof documents and referrals and an evaluationdatasetof queries and ground truth documents. Under thepaper_retrievaldomain, we include theacl,acl_small, andarxivcorpuses and datasets, and under theentity_retrievaldomain, we include thedbpedia_smallcorpus and dataset.Construction details:Theacl_small,acl, andarxivcorpuses are constructed from the rich paper metadata parses provided by Allen AI'sS2ORCproject. Documents consist of concatenated paper titles and abstracts from up-to-2017 ACL and ArXiv papers, respectively, and referrals consist of in-text citations between up-to-2017 papers. The respective evaluation datasets are also from the parses, consisting of in-text citations from2018-and-onpapers citing the up-to-2017 papers in the corpus -- this time-based split prevents data leakage and mirrors deployment conditions.Thedbpedia_smallcorpus and dataset is sampled from the DBPedia task in theBEIRbenchmark. Referrals are mined from Wikipedia HTML usingWikiExtractor.Data can be loaded via our utility functions atrar.utils:fromrar.utilsimportload_corpus,load_eval_datasetdocs,referrals=load_corpus(domain='paper_retrieval',corpus='acl_small')queries,ground_truth=load_eval_dataset(domain='paper_retrieval',dataset='acl_small')Our data representations are simple and intuitive:
— Acorpusis a lists of document strings
— A set ofreferralsis a list of lists of document strings (one list of referrals per document)
Similarly:
— A set ofqueriesis a lists of query strings
— The correspondingground_truthiseithera list of document strings (one ground truth document per query, e.g. the cited paper in paper retrieval)ora list of lists of document strings (multiple relevant ground truth documents per query, e.g. all relevant Wikipedia pages for a givendbpedia_smallquery)Custom dataCreating a corpus is as simple as constructing these lists (referrals are optional). For example:docs=['Steve Jobs was a revolutionary technological thinker and designer',"Bill Gates founded the world's largest software company"]referrals=[['Apple CEO','Magic Leap founder'],['Microsoft CEO','The Giving Pledge co-founder']]retriever=BM25Retriever(docs,referrals)Creating an evaluation corpus is similarly easy:queries=['Who built the Apple Macintosh?']ground_truth=[docs[0]]EvaluationWe implement the Recall@k and MRR metrics underrar.metrics, which can be used standalone or with our utility functions atrar.utils:fromrar.utilsimportevaluate_retrieverevaluate_retriever(retriever,queries,ground_truth)By default,evaluate_retrievalattempts to compute MRR, Recall@1, and Recall@10 metrics. The keyword parametermultiple_correct=Falseremoves MRR, since it does not support multiple ground truth documents per query (e.g. fordbpedia_small). Seeexamples.ipynbfor example outputs.If you find this repository helpful, feel free to cite our publicationReferral Augmentation for Zero-Shot Information Retrieval:@misc{tang2023referral,
      title={Referral Augmentation for Zero-Shot Information Retrieval}, 
      author={Michael Tang and Shunyu Yao and John Yang and Karthik Narasimhan},
      year={2023},
      eprint={2305.15098},
      archivePrefix={arXiv},
      primaryClass={cs.CL}
}