integrations/dense/test_dpr.py

#
# Pyserini: Reproducible IR research with sparse and dense representations
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#

"""Integration tests for DPR model using pre-encoded queries."""

import json
import os
import socket
import unittest

from integrations.utils import clean_files, run_command, parse_score_qa
from pyserini.search import QueryEncoder
from pyserini.search import get_topics


class TestSearchIntegration(unittest.TestCase):
    def setUp(self):
        self.temp_files = []
        self.threads = 16
        self.batch_size = 256

        # Hard-code larger values for internal servers
        if socket.gethostname().startswith('damiano') or socket.gethostname().startswith('orca'):
            self.threads = 36
            self.batch_size = 144

    def test_dpr_nq_test_bf_otf(self):
        output_file = 'test_run.dpr.nq-test.multi.bf.otf.trec'
        retrieval_file = 'test_run.dpr.nq-test.multi.bf.otf.json'
        self.temp_files.extend([output_file, retrieval_file])
        cmd1 = f'python -m pyserini.search.faiss --topics dpr-nq-test \
                             --index wikipedia-dpr-multi-bf \
                             --encoder facebook/dpr-question_encoder-multiset-base \
                             --output {output_file} \
                             --batch-size {self.batch_size} --threads {self.threads}'
        cmd2 = f'python -m pyserini.eval.convert_trec_run_to_dpr_retrieval_run --topics dpr-nq-test \
                                                           --index wikipedia-dpr \
                                                           --input {output_file} \
                                                           --output {retrieval_file}'
        cmd3 = f'python -m pyserini.eval.evaluate_dpr_retrieval --retrieval {retrieval_file} --topk 20'
        status1 = os.system(cmd1)
        status2 = os.system(cmd2)
        stdout, stderr = run_command(cmd3)
        score = parse_score_qa(stdout, 'Top20')
        self.assertEqual(status1, 0)
        self.assertEqual(status2, 0)
        self.assertAlmostEqual(score, 0.7947, places=4)

    def test_dpr_nq_test_bf_bm25_hybrid_otf(self):
        output_file = 'test_run.dpr.nq-test.multi.bf.otf.bm25.trec'
        retrieval_file = 'test_run.dpr.nq-test.multi.bf.otf.bm25.json'
        self.temp_files.extend([output_file, retrieval_file])
        cmd1 = f'python -m pyserini.search.hybrid dense  --index wikipedia-dpr-multi-bf \
                                    --encoder facebook/dpr-question_encoder-multiset-base \
                             sparse --index wikipedia-dpr \
                             fusion --alpha 1.3 \
                             run    --topics dpr-nq-test \
                                    --batch-size {self.batch_size} --threads {self.threads} \
                                    --output {output_file} '
        cmd2 = f'python -m pyserini.eval.convert_trec_run_to_dpr_retrieval_run --topics dpr-nq-test \
                                                           --index wikipedia-dpr \
                                                           --input {output_file} \
                                                           --output {retrieval_file}'
        cmd3 = f'python -m pyserini.eval.evaluate_dpr_retrieval --retrieval {retrieval_file} --topk 20'
        status1 = os.system(cmd1)
        status2 = os.system(cmd2)
        stdout, stderr = run_command(cmd3)
        score = parse_score_qa(stdout, 'Top20')
        self.assertEqual(status1, 0)
        self.assertEqual(status2, 0)
        self.assertAlmostEqual(score, 0.8260, places=4)

    def test_dpr_nq_test_encoded_queries(self):
        encoder = QueryEncoder.load_encoded_queries('dpr_multi-nq-test')
        topics = get_topics('dpr-nq-test')
        for t in topics:
            self.assertTrue(topics[t]['title'] in encoder.embedding)

    def test_dpr_trivia_test_bf_otf(self):
        output_file = 'test_run.dpr.trivia-test.multi.bf.otf.trec'
        retrieval_file = 'test_run.dpr.trivia-test.multi.bf.otf.json'
        self.temp_files.extend([output_file, retrieval_file])
        cmd1 = f'python -m pyserini.search.faiss --topics dpr-trivia-test \
                             --encoder facebook/dpr-question_encoder-multiset-base \
                             --index wikipedia-dpr-multi-bf \
                             --output {output_file} \
                             --batch-size {self.batch_size} --threads {self.threads}'
        cmd2 = f'python -m pyserini.eval.convert_trec_run_to_dpr_retrieval_run --topics dpr-trivia-test \
                                                           --index wikipedia-dpr \
                                                           --input {output_file} \
                                                           --output {retrieval_file}'
        cmd3 = f'python -m pyserini.eval.evaluate_dpr_retrieval --retrieval {retrieval_file} --topk 20'
        status1 = os.system(cmd1)
        status2 = os.system(cmd2)
        stdout, stderr = run_command(cmd3)
        score = parse_score_qa(stdout, 'Top20')
        self.assertEqual(status1, 0)
        self.assertEqual(status2, 0)
        self.assertAlmostEqual(score, 0.7887, places=4)

    def test_dpr_trivia_test_bf_bm25_hybrid_otf(self):
        output_file = 'test_run.dpr.trivia-test.multi.bf.otf.bm25.trec'
        retrieval_file = 'test_run.dpr.trivia-test.multi.bf.otf.bm25.json'
        self.temp_files.extend([output_file, retrieval_file])
        cmd1 = f'python -m pyserini.search.hybrid dense  --index wikipedia-dpr-multi-bf \
                                    --encoder facebook/dpr-question_encoder-multiset-base \
                             sparse --index wikipedia-dpr \
                             fusion --alpha 0.95 \
                             run    --topics dpr-trivia-test \
                                    --batch-size {self.batch_size} --threads {self.threads} \
                                    --output {output_file} '
        cmd2 = f'python -m pyserini.eval.convert_trec_run_to_dpr_retrieval_run --topics dpr-trivia-test \
                                                           --index wikipedia-dpr \
                                                           --input {output_file} \
                                                           --output {retrieval_file}'
        cmd3 = f'python -m pyserini.eval.evaluate_dpr_retrieval --retrieval {retrieval_file} --topk 20'
        status1 = os.system(cmd1)
        status2 = os.system(cmd2)
        stdout, stderr = run_command(cmd3)
        score = parse_score_qa(stdout, 'Top20')
        self.assertEqual(status1, 0)
        self.assertEqual(status2, 0)
        self.assertAlmostEqual(score, 0.8264, places=4)

    def test_dpr_trivia_test_encoded_queries(self):
        encoder = QueryEncoder.load_encoded_queries('dpr_multi-trivia-test')
        topics = get_topics('dpr-trivia-test')
        for t in topics:
            self.assertTrue(topics[t]['title'] in encoder.embedding)

    def test_dpr_wq_test_bf_otf(self):
        output_file = 'test_run.dpr.wq-test.multi.bf.otf.trec'
        retrieval_file = 'test_run.dpr.wq-test.multi.bf.otf.json'
        self.temp_files.extend([output_file, retrieval_file])
        cmd1 = f'python -m pyserini.search.faiss --topics dpr-wq-test \
                             --index wikipedia-dpr-multi-bf \
                             --encoder facebook/dpr-question_encoder-multiset-base \
                             --output {output_file} \
                             --batch-size {self.batch_size} --threads {self.threads}'
        cmd2 = f'python -m pyserini.eval.convert_trec_run_to_dpr_retrieval_run --topics dpr-wq-test \
                                                           --index wikipedia-dpr \
                                                           --input {output_file} \
                                                           --output {retrieval_file}'
        cmd3 = f'python -m pyserini.eval.evaluate_dpr_retrieval --retrieval {retrieval_file} --topk 20'
        status1 = os.system(cmd1)
        status2 = os.system(cmd2)
        stdout, stderr = run_command(cmd3)
        score = parse_score_qa(stdout, 'Top20')
        self.assertEqual(status1, 0)
        self.assertEqual(status2, 0)
        self.assertAlmostEqual(score, 0.7505, places=4)

    def test_dpr_wq_test_bf_bm25_hybrid_otf(self):
        output_file = 'test_run.dpr.wq-test.multi.bf.otf.bm25.trec'
        retrieval_file = 'test_run.dpr.wq-test.multi.bf.otf.bm25.json'
        self.temp_files.extend([output_file, retrieval_file])
        cmd1 = f'python -m pyserini.search.hybrid dense  --index wikipedia-dpr-multi-bf \
                                    --encoder facebook/dpr-question_encoder-multiset-base \
                             sparse --index wikipedia-dpr \
                             fusion --alpha 0.95 \
                             run    --topics dpr-wq-test \
                                    --batch-size {self.batch_size} --threads {self.threads} \
                                    --output {output_file} '
        cmd2 = f'python -m pyserini.eval.convert_trec_run_to_dpr_retrieval_run --topics dpr-wq-test \
                                                           --index wikipedia-dpr \
                                                           --input {output_file} \
                                                           --output {retrieval_file}'
        cmd3 = f'python -m pyserini.eval.evaluate_dpr_retrieval --retrieval {retrieval_file} --topk 20'
        status1 = os.system(cmd1)
        status2 = os.system(cmd2)
        stdout, stderr = run_command(cmd3)
        score = parse_score_qa(stdout, 'Top20')
        self.assertEqual(status1, 0)
        self.assertEqual(status2, 0)
        self.assertAlmostEqual(score, 0.7712, places=4)

    def test_dpr_wq_test_encoded_queries(self):
        encoder = QueryEncoder.load_encoded_queries('dpr_multi-wq-test')
        topics = get_topics('dpr-wq-test')
        for t in topics:
            self.assertTrue(topics[t]['title'] in encoder.embedding)

    def test_dpr_curated_test_bf_otf(self):
        output_file = 'test_run.dpr.curated-test.multi.bf.otf.trec'
        retrieval_file = 'test_run.dpr.curated-test.multi.bf.otf.json'
        self.temp_files.extend([output_file, retrieval_file])
        cmd1 = f'python -m pyserini.search.faiss --topics dpr-curated-test \
                             --index wikipedia-dpr-multi-bf \
                             --encoder facebook/dpr-question_encoder-multiset-base \
                             --output {output_file} \
                             --batch-size {self.batch_size} --threads {self.threads}'
        cmd2 = f'python -m pyserini.eval.convert_trec_run_to_dpr_retrieval_run --topics dpr-curated-test \
                                                           --index wikipedia-dpr \
                                                           --input {output_file} \
                                                           --output {retrieval_file} \
                                                           --regex'
        cmd3 = f'python -m pyserini.eval.evaluate_dpr_retrieval --retrieval {retrieval_file} --topk 20 --regex'
        status1 = os.system(cmd1)
        status2 = os.system(cmd2)
        stdout, stderr = run_command(cmd3)
        score = parse_score_qa(stdout, 'Top20')
        self.assertEqual(status1, 0)
        self.assertEqual(status2, 0)
        self.assertAlmostEqual(score, 0.8876, places=4)

    def test_dpr_curated_test_bf_bm25_hybrid_otf(self):
        output_file = 'test_run.dpr.curated-test.multi.bf.otf.bm25.trec'
        retrieval_file = 'test_run.dpr.curated-test.multi.bf.otf.bm25.json'
        self.temp_files.extend([output_file, retrieval_file])
        cmd1 = f'python -m pyserini.search.hybrid dense  --index wikipedia-dpr-multi-bf \
                                    --encoder facebook/dpr-question_encoder-multiset-base \
                             sparse --index wikipedia-dpr \
                             fusion --alpha 1.05 \
                             run    --topics dpr-curated-test \
                                    --batch-size {self.batch_size} --threads {self.threads} \
                                    --output {output_file} '
        cmd2 = f'python -m pyserini.eval.convert_trec_run_to_dpr_retrieval_run --topics dpr-curated-test \
                                                           --index wikipedia-dpr \
                                                           --input {output_file} \
                                                           --output {retrieval_file} \
                                                           --regex'
        cmd3 = f'python -m pyserini.eval.evaluate_dpr_retrieval --retrieval {retrieval_file} --topk 20 --regex'
        status1 = os.system(cmd1)
        status2 = os.system(cmd2)
        stdout, stderr = run_command(cmd3)
        score = parse_score_qa(stdout, 'Top20')
        self.assertEqual(status1, 0)
        self.assertEqual(status2, 0)
        self.assertAlmostEqual(score, 0.9006, places=4)

    def test_dpr_curated_test_encoded_queries(self):
        encoder = QueryEncoder.load_encoded_queries('dpr_multi-curated-test')
        topics = get_topics('dpr-curated-test')
        for t in topics:
            self.assertTrue(topics[t]['title'] in encoder.embedding)

    def test_dpr_squad_test_bf_otf(self):
        output_file = 'test_run.dpr.squad-test.multi.bf.otf.trec'
        retrieval_file = 'test_run.dpr.squad-test.multi.bf.otf.json'
        self.temp_files.extend([output_file, retrieval_file])
        cmd1 = f'python -m pyserini.search.faiss --topics dpr-squad-test \
                             --index wikipedia-dpr-multi-bf \
                             --encoder facebook/dpr-question_encoder-multiset-base \
                             --output {output_file} \
                             --batch-size {self.batch_size} --threads {self.threads}'
        cmd2 = f'python -m pyserini.eval.convert_trec_run_to_dpr_retrieval_run --topics dpr-squad-test \
                                                           --index wikipedia-dpr \
                                                           --input {output_file} \
                                                           --output {retrieval_file}'
        cmd3 = f'python -m pyserini.eval.evaluate_dpr_retrieval --retrieval {retrieval_file} --topk 20'
        status1 = os.system(cmd1)
        status2 = os.system(cmd2)
        stdout, stderr = run_command(cmd3)
        score = parse_score_qa(stdout, 'Top20')
        self.assertEqual(status1, 0)
        self.assertEqual(status2, 0)
        self.assertAlmostEqual(score, 0.5199, places=4)

    def test_dpr_squad_test_bf_bm25_hybrid_otf(self):
        output_file = 'test_run.dpr.squad-test.multi.bf.otf.bm25.trec'
        retrieval_file = 'test_run.dpr.squad-test.multi.bf.otf.bm25.json'
        self.temp_files.extend([output_file, retrieval_file])
        cmd1 = f'python -m pyserini.search.hybrid dense  --index wikipedia-dpr-multi-bf \
                                    --encoder facebook/dpr-question_encoder-multiset-base \
                             sparse --index wikipedia-dpr \
                             fusion --alpha 2.0 \
                             run    --topics dpr-squad-test \
                                    --batch-size {self.batch_size} --threads {self.threads} \
                                    --output {output_file} '
        cmd2 = f'python -m pyserini.eval.convert_trec_run_to_dpr_retrieval_run --topics dpr-squad-test \
                                                           --index wikipedia-dpr \
                                                           --input {output_file} \
                                                           --output {retrieval_file}'
        cmd3 = f'python -m pyserini.eval.evaluate_dpr_retrieval --retrieval {retrieval_file} --topk 20'
        status1 = os.system(cmd1)
        status2 = os.system(cmd2)
        stdout, stderr = run_command(cmd3)
        score = parse_score_qa(stdout, 'Top20')
        self.assertEqual(status1, 0)
        self.assertEqual(status2, 0)
        # This appears to be a flaky test case; previously, we were getting a score of 0.7511, per
        # https://github.com/castorini/pyserini/pull/1273/files#diff-799c2c339e1d7defa31fa1e82f9b16886269b37805376ef93f7c8afedcee574e
        # Sometimes we get 0.7512. Fix is to reduce tolerance.
        self.assertAlmostEqual(score, 0.7514, places=3)

    def test_dpr_squad_test_encoded_queries(self):
        encoder = QueryEncoder.load_encoded_queries('dpr_multi-squad-test')
        topics = get_topics('dpr-squad-test')
        for t in topics:
            self.assertTrue(topics[t]['title'] in encoder.embedding)

    def test_convert_trec_run_to_dpr_retrieval_run(self):
        trec_run_file = 'tests/resources/simple_test_run_convert_trec_run_dpr.trec'
        topics_file = 'tests/resources/simple_topics_dpr.txt'
        dpr_run_file = 'test_run.convert.trec_run.dpr.json'
        collection_path = "tests/resources/sample_collection_dense"
        topic_reader = "io.anserini.search.topicreader.DprNqTopicReader"
        index_dir = 'temp_index'

        self.temp_files.extend([dpr_run_file, index_dir])
        cmd1 = f'python -m pyserini.index.lucene -collection JsonCollection ' + \
               f'-generator DefaultLuceneDocumentGenerator ' + \
               f'-threads 1 -input {collection_path} -index {index_dir} -storeRaw'

        cmd2 = f'python -m pyserini.eval.convert_trec_run_to_dpr_retrieval_run --topics-file {topics_file} \
                                                           --topics-reader {topic_reader} \
                                                           --index {index_dir} \
                                                           --input {trec_run_file} \
                                                           --output {dpr_run_file}'
        _ = os.system(cmd1)
        _ = os.system(cmd2)

        with open(dpr_run_file) as f:
            topic_data = json.load(f)

        self.assertEqual(topic_data["0"]["answers"], ['text'])
        self.assertEqual(topic_data["0"]["question"], "what is in document three")
        self.assertEqual(topic_data["1"]["answers"], ['contents'])
        self.assertEqual(topic_data["1"]["question"], "what is document two")

    def tearDown(self):
        clean_files(self.temp_files)


if __name__ == '__main__':
    unittest.main()