From e171c61c2d0028b298ad9a2048a6da6f22c74cef Mon Sep 17 00:00:00 2001 From: Sean MacAvaney Date: Mon, 2 Dec 2024 18:42:48 +0000 Subject: [PATCH] faster tests --- tests/test_models.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/test_models.py b/tests/test_models.py index 29b9c27..1f6d612 100644 --- a/tests/test_models.py +++ b/tests/test_models.py @@ -11,13 +11,13 @@ class TestModels(unittest.TestCase): def _base_test(self, model, test_query_encoder=True, test_doc_encoder=True, test_scorer=True, test_indexer=True, test_retriever=True): dataset = pt.get_dataset('irds:vaswani') + topics = dataset.get_topics().head(10) - docs = list(itertools.islice(pt.get_dataset('irds:vaswani').get_corpus_iter(), 200)) + docs = list(itertools.islice(pt.get_dataset('irds:vaswani').get_corpus_iter(), 50)) docs_df = pd.DataFrame(docs) if test_query_encoder: with self.subTest('query_encoder'): - topics = dataset.get_topics() enc_topics = model(topics) self.assertEqual(len(enc_topics), len(topics)) self.assertTrue('query_vec' in enc_topics.columns)