NetsPresso_QA / tests /test_search.py
geonmin-kim's picture
Upload folder using huggingface_hub
d6585f5
#
# 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.
#
import os
import shutil
import tarfile
import unittest
from random import randint
from typing import List, Dict
from urllib.request import urlretrieve
from pyserini.search.lucene import LuceneSearcher, JLuceneSearcherResult
from pyserini.index.lucene import Document
class TestSearch(unittest.TestCase):
@classmethod
def setUpClass(cls):
# Download pre-built CACM index built using Lucene 9; append a random value to avoid filename clashes.
r = randint(0, 10000000)
cls.collection_url = 'https://github.com/castorini/anserini-data/raw/master/CACM/lucene9-index.cacm.tar.gz'
cls.tarball_name = 'lucene-index.cacm-{}.tar.gz'.format(r)
cls.index_dir = 'index{}/'.format(r)
filename, headers = urlretrieve(cls.collection_url, cls.tarball_name)
tarball = tarfile.open(cls.tarball_name)
tarball.extractall(cls.index_dir)
tarball.close()
cls.searcher = LuceneSearcher(f'{cls.index_dir}lucene9-index.cacm')
# Create index without document vectors
# The current directory depends on if you're running inside an IDE or from command line.
curdir = os.getcwd()
if curdir.endswith('tests'):
corpus_path = '../tests/resources/sample_collection_json'
else:
corpus_path = 'tests/resources/sample_collection_json'
cls.no_vec_index_dir = 'no_vec_index'
cmd1 = f'python -m pyserini.index.lucene -collection JsonCollection ' + \
f'-generator DefaultLuceneDocumentGenerator ' + \
f'-threads 1 -input {corpus_path} -index {cls.no_vec_index_dir}'
os.system(cmd1)
cls.no_vec_searcher = LuceneSearcher(cls.no_vec_index_dir)
def test_basic(self):
self.assertTrue(self.searcher.get_similarity().toString().startswith('BM25'))
hits = self.searcher.search('information retrieval')
self.assertEqual(3204, self.searcher.num_docs)
self.assertTrue(isinstance(hits, List))
self.assertTrue(isinstance(hits[0], JLuceneSearcherResult))
self.assertEqual(hits[0].docid, 'CACM-3134')
self.assertEqual(hits[0].lucene_docid, 3133)
self.assertEqual(len(hits[0].contents), 1500)
self.assertEqual(len(hits[0].raw), 1532)
self.assertAlmostEqual(hits[0].score, 4.76550, places=5)
# Test accessing the raw Lucene document and fetching fields from it:
self.assertEqual(hits[0].lucene_document.getField('id').stringValue(), 'CACM-3134')
self.assertEqual(hits[0].lucene_document.get('id'), 'CACM-3134') # simpler call, same result as above
self.assertEqual(len(hits[0].lucene_document.getField('raw').stringValue()), 1532)
self.assertEqual(len(hits[0].lucene_document.get('raw')), 1532) # simpler call, same result as above
self.assertTrue(isinstance(hits[9], JLuceneSearcherResult))
self.assertEqual(hits[9].docid, 'CACM-2516')
self.assertAlmostEqual(hits[9].score, 4.21740, places=5)
hits = self.searcher.search('search')
self.assertTrue(isinstance(hits[0], JLuceneSearcherResult))
self.assertEqual(hits[0].docid, 'CACM-3058')
self.assertAlmostEqual(hits[0].score, 2.85760, places=5)
self.assertTrue(isinstance(hits[9], JLuceneSearcherResult))
self.assertEqual(hits[9].docid, 'CACM-3040')
self.assertAlmostEqual(hits[9].score, 2.68780, places=5)
def test_batch(self):
results = self.searcher.batch_search(['information retrieval', 'search'], ['q1', 'q2'], threads=2)
self.assertEqual(3204, self.searcher.num_docs)
self.assertTrue(isinstance(results, Dict))
self.assertTrue(isinstance(results['q1'], List))
self.assertTrue(isinstance(results['q1'][0], JLuceneSearcherResult))
self.assertEqual(results['q1'][0].docid, 'CACM-3134')
self.assertAlmostEqual(results['q1'][0].score, 4.76550, places=5)
self.assertTrue(isinstance(results['q1'][9], JLuceneSearcherResult))
self.assertEqual(results['q1'][9].docid, 'CACM-2516')
self.assertAlmostEqual(results['q1'][9].score, 4.21740, places=5)
self.assertTrue(isinstance(results['q2'], List))
self.assertTrue(isinstance(results['q2'][0], JLuceneSearcherResult))
self.assertEqual(results['q2'][0].docid, 'CACM-3058')
self.assertAlmostEqual(results['q2'][0].score, 2.85760, places=5)
self.assertTrue(isinstance(results['q2'][9], JLuceneSearcherResult))
self.assertEqual(results['q2'][9].docid, 'CACM-3040')
self.assertAlmostEqual(results['q2'][9].score, 2.68780, places=5)
def test_batch_rocchio(self):
self.searcher.set_rocchio()
results = self.searcher.batch_search(['information retrieval', 'search'], ['q1', 'q2'], threads=2)
self.searcher.unset_rocchio()
self.assertEqual(3204, self.searcher.num_docs)
self.assertTrue(isinstance(results, Dict))
self.assertTrue(isinstance(results['q1'], List))
self.assertTrue(isinstance(results['q1'][0], JLuceneSearcherResult))
self.assertEqual(results['q1'][0].docid, 'CACM-3134')
self.assertAlmostEqual(results['q1'][0].score, 7.18830, places=5)
self.assertTrue(isinstance(results['q1'][9], JLuceneSearcherResult))
self.assertEqual(results['q1'][9].docid, 'CACM-2140')
self.assertAlmostEqual(results['q1'][9].score, 5.57970, places=5)
self.assertTrue(isinstance(results['q2'], List))
self.assertTrue(isinstance(results['q2'][0], JLuceneSearcherResult))
self.assertEqual(results['q2'][0].docid, 'CACM-3041')
self.assertAlmostEqual(results['q2'][0].score, 6.14870, places=5)
self.assertTrue(isinstance(results['q2'][9], JLuceneSearcherResult))
self.assertEqual(results['q2'][9].docid, 'CACM-2251')
self.assertAlmostEqual(results['q2'][9].score, 4.97380, places=5)
def test_basic_k(self):
hits = self.searcher.search('information retrieval', k=100)
self.assertEqual(3204, self.searcher.num_docs)
self.assertTrue(isinstance(hits, List))
self.assertTrue(isinstance(hits[0], JLuceneSearcherResult))
self.assertEqual(len(hits), 100)
def test_batch_k(self):
results = self.searcher.batch_search(['information retrieval', 'search'], ['q1', 'q2'], k=100, threads=2)
self.assertEqual(3204, self.searcher.num_docs)
self.assertTrue(isinstance(results, Dict))
self.assertTrue(isinstance(results['q1'], List))
self.assertTrue(isinstance(results['q1'][0], JLuceneSearcherResult))
self.assertEqual(len(results['q1']), 100)
self.assertTrue(isinstance(results['q2'], List))
self.assertTrue(isinstance(results['q2'][0], JLuceneSearcherResult))
self.assertEqual(len(results['q2']), 100)
def test_basic_fields(self):
# This test just provides a sanity check, it's not that interesting as it only searches one field.
hits = self.searcher.search('information retrieval', k=42, fields={'contents': 2.0})
self.assertEqual(3204, self.searcher.num_docs)
self.assertTrue(isinstance(hits, List))
self.assertTrue(isinstance(hits[0], JLuceneSearcherResult))
self.assertEqual(len(hits), 42)
def test_batch_fields(self):
# This test just provides a sanity check, it's not that interesting as it only searches one field.
results = self.searcher.batch_search(['information retrieval', 'search'], ['q1', 'q2'], k=42,
threads=2, fields={'contents': 2.0})
self.assertEqual(3204, self.searcher.num_docs)
self.assertTrue(isinstance(results, Dict))
self.assertTrue(isinstance(results['q1'], List))
self.assertTrue(isinstance(results['q1'][0], JLuceneSearcherResult))
self.assertEqual(len(results['q1']), 42)
self.assertTrue(isinstance(results['q2'], List))
self.assertTrue(isinstance(results['q2'][0], JLuceneSearcherResult))
self.assertEqual(len(results['q2']), 42)
def test_different_similarity(self):
# qld, default mu
self.searcher.set_qld()
self.assertTrue(self.searcher.get_similarity().toString().startswith('LM Dirichlet'))
hits = self.searcher.search('information retrieval')
self.assertEqual(hits[0].docid, 'CACM-3134')
self.assertAlmostEqual(hits[0].score, 3.68030, places=5)
self.assertEqual(hits[9].docid, 'CACM-1927')
self.assertAlmostEqual(hits[9].score, 2.53240, places=5)
# bm25, default parameters
self.searcher.set_bm25()
self.assertTrue(self.searcher.get_similarity().toString().startswith('BM25'))
hits = self.searcher.search('information retrieval')
self.assertEqual(hits[0].docid, 'CACM-3134')
self.assertAlmostEqual(hits[0].score, 4.76550, places=5)
self.assertEqual(hits[9].docid, 'CACM-2516')
self.assertAlmostEqual(hits[9].score, 4.21740, places=5)
# qld, custom mu
self.searcher.set_qld(100)
self.assertTrue(self.searcher.get_similarity().toString().startswith('LM Dirichlet'))
hits = self.searcher.search('information retrieval')
self.assertEqual(hits[0].docid, 'CACM-3134')
self.assertAlmostEqual(hits[0].score, 6.35580, places=5)
self.assertEqual(hits[9].docid, 'CACM-2631')
self.assertAlmostEqual(hits[9].score, 5.18960, places=5)
# bm25, custom parameters
self.searcher.set_bm25(0.8, 0.3)
self.assertTrue(self.searcher.get_similarity().toString().startswith('BM25'))
hits = self.searcher.search('information retrieval')
self.assertEqual(hits[0].docid, 'CACM-3134')
self.assertAlmostEqual(hits[0].score, 4.86880, places=5)
self.assertEqual(hits[9].docid, 'CACM-2516')
self.assertAlmostEqual(hits[9].score, 4.33320, places=5)
def test_rm3(self):
self.searcher = LuceneSearcher(f'{self.index_dir}lucene9-index.cacm')
self.searcher.set_rm3()
self.assertTrue(self.searcher.is_using_rm3())
hits = self.searcher.search('information retrieval')
self.assertEqual(hits[0].docid, 'CACM-3134')
self.assertAlmostEqual(hits[0].score, 2.17350, places=5)
self.assertEqual(hits[9].docid, 'CACM-2516')
self.assertAlmostEqual(hits[9].score, 1.70180, places=5)
feedback_terms = self.searcher.get_feedback_terms('information retrieval')
self.assertEqual(len(feedback_terms), 10)
self.assertAlmostEqual(feedback_terms['storag'], 0.024701, places=5)
self.searcher.unset_rm3()
self.assertFalse(self.searcher.is_using_rm3())
self.assertIsNone(self.searcher.get_feedback_terms('information retrieval'))
hits = self.searcher.search('information retrieval')
self.assertEqual(hits[0].docid, 'CACM-3134')
self.assertAlmostEqual(hits[0].score, 4.76550, places=5)
self.assertEqual(hits[9].docid, 'CACM-2516')
self.assertAlmostEqual(hits[9].score, 4.21740, places=5)
self.searcher.set_rm3(fb_docs=4, fb_terms=6, original_query_weight=0.3)
self.assertTrue(self.searcher.is_using_rm3())
hits = self.searcher.search('information retrieval')
self.assertEqual(hits[0].docid, 'CACM-3134')
self.assertAlmostEqual(hits[0].score, 2.17150, places=5)
self.assertEqual(hits[9].docid, 'CACM-1457')
self.assertAlmostEqual(hits[9].score, 1.45560, places=5)
with self.assertRaises(TypeError):
self.no_vec_searcher.set_rm3()
def test_rocchio(self):
self.searcher = LuceneSearcher(f'{self.index_dir}lucene9-index.cacm')
self.searcher.set_rocchio()
self.assertTrue(self.searcher.is_using_rocchio())
hits = self.searcher.search('information retrieval')
self.assertEqual(hits[0].docid, 'CACM-3134')
self.assertAlmostEqual(hits[0].score, 7.18830, places=5)
self.assertEqual(hits[9].docid, 'CACM-2140')
self.assertAlmostEqual(hits[9].score, 5.57970, places=5)
feedback_terms = self.searcher.get_feedback_terms('information retrieval')
self.assertEqual(len(feedback_terms), 10)
self.assertAlmostEqual(feedback_terms['storag'], 0.132200, places=5)
self.searcher.unset_rocchio()
self.assertFalse(self.searcher.is_using_rocchio())
self.assertIsNone(self.searcher.get_feedback_terms('information retrieval'))
hits = self.searcher.search('information retrieval')
self.assertEqual(hits[0].docid, 'CACM-3134')
self.assertAlmostEqual(hits[0].score, 4.76550, places=5)
self.assertEqual(hits[9].docid, 'CACM-2516')
self.assertAlmostEqual(hits[9].score, 4.21740, places=5)
self.searcher.set_rocchio(top_fb_terms=10, top_fb_docs=8, bottom_fb_terms=10,
bottom_fb_docs=8, alpha=0.4, beta=0.5, gamma=0.1,
debug=False, use_negative=True)
self.assertTrue(self.searcher.is_using_rocchio())
hits = self.searcher.search('information retrieval')
self.assertEqual(hits[0].docid, 'CACM-3134')
self.assertAlmostEqual(hits[0].score, 3.64890, places=5)
self.assertEqual(hits[9].docid, 'CACM-1032')
self.assertAlmostEqual(hits[9].score, 2.57510, places=5)
self.searcher.set_rocchio(top_fb_terms=10, top_fb_docs=8, bottom_fb_terms=10,
bottom_fb_docs=8, alpha=0.4, beta=0.5, gamma=0.1,
debug=False, use_negative=False)
self.assertTrue(self.searcher.is_using_rocchio())
hits = self.searcher.search('information retrieval')
self.assertEqual(hits[0].docid, 'CACM-3134')
self.assertAlmostEqual(hits[0].score, 4.03900, places=5)
self.assertEqual(hits[9].docid, 'CACM-1032')
self.assertAlmostEqual(hits[9].score, 2.91550, places=5)
with self.assertRaises(TypeError):
self.no_vec_searcher.set_rocchio()
def test_doc_int(self):
# The doc method is overloaded: if input is int, it's assumed to be a Lucene internal docid.
doc = self.searcher.doc(1)
self.assertTrue(isinstance(doc, Document))
# These are all equivalent ways to get the docid.
self.assertEqual('CACM-0002', doc.id())
self.assertEqual('CACM-0002', doc.docid())
self.assertEqual('CACM-0002', doc.get('id'))
self.assertEqual('CACM-0002', doc.lucene_document().getField('id').stringValue())
# These are all equivalent ways to get the 'raw' field
self.assertEqual(186, len(doc.raw()))
self.assertEqual(186, len(doc.get('raw')))
self.assertEqual(186, len(doc.lucene_document().get('raw')))
self.assertEqual(186, len(doc.lucene_document().getField('raw').stringValue()))
# These are all equivalent ways to get the 'contents' field
self.assertEqual(154, len(doc.contents()))
self.assertEqual(154, len(doc.get('contents')))
self.assertEqual(154, len(doc.lucene_document().get('contents')))
self.assertEqual(154, len(doc.lucene_document().getField('contents').stringValue()))
# Should return None if we request a docid that doesn't exist
self.assertTrue(self.searcher.doc(314159) is None)
def test_doc_str(self):
# The doc method is overloaded: if input is str, it's assumed to be an external collection docid.
doc = self.searcher.doc('CACM-0002')
self.assertTrue(isinstance(doc, Document))
# These are all equivalent ways to get the docid.
self.assertEqual(doc.lucene_document().getField('id').stringValue(), 'CACM-0002')
self.assertEqual(doc.id(), 'CACM-0002')
self.assertEqual(doc.docid(), 'CACM-0002')
self.assertEqual(doc.get('id'), 'CACM-0002')
# These are all equivalent ways to get the 'raw' field
self.assertEqual(186, len(doc.raw()))
self.assertEqual(186, len(doc.get('raw')))
self.assertEqual(186, len(doc.lucene_document().get('raw')))
self.assertEqual(186, len(doc.lucene_document().getField('raw').stringValue()))
# These are all equivalent ways to get the 'contents' field
self.assertEqual(154, len(doc.contents()))
self.assertEqual(154, len(doc.get('contents')))
self.assertEqual(154, len(doc.lucene_document().get('contents')))
self.assertEqual(154, len(doc.lucene_document().getField('contents').stringValue()))
# Should return None if we request a docid that doesn't exist
self.assertTrue(self.searcher.doc('foo') is None)
def test_batch_doc(self):
docids = ['CACM-0002', 'CACM-3134', 'Fake Doc1', 'Fake Doc2']
batch_doc = self.searcher.batch_doc(docids, threads=2)
# Doc Id should have corresponding Document object in batch result
for docid in docids[:2]:
doc = batch_doc.get(docid)
self.assertTrue(isinstance(doc, Document))
self.assertEqual(doc.lucene_document().getField('id').stringValue(), docid)
self.assertEqual(doc.id(), docid)
self.assertEqual(doc.docid(), docid)
self.assertEqual(doc.get('id'), docid)
# None should be returned for docids that don't exist
for docid in docids[2:]:
doc = batch_doc.get(docid[-1])
self.assertEqual(doc, None)
def test_doc_by_field(self):
self.assertEqual(self.searcher.doc('CACM-3134').docid(),
self.searcher.doc_by_field('id', 'CACM-3134').docid())
# Should return None if we request a docid that doesn't exist
self.assertTrue(self.searcher.doc_by_field('foo', 'bar') is None)
@classmethod
def tearDownClass(cls):
cls.searcher.close()
cls.no_vec_searcher.close()
os.remove(cls.tarball_name)
shutil.rmtree(cls.index_dir)
shutil.rmtree(cls.no_vec_index_dir)
if __name__ == '__main__':
unittest.main()