fixing test

dleemiller · Jul 28, 2024 · 54526c9 · 54526c9
1 parent bdf4621
commit 54526c9
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/tests/test_inference.py b/tests/test_inference.py
@@ -82,11 +82,11 @@ def setUp(self, mock_tokenizer):
         return_value=np.array([[0.1] * 64, [0.1] * 64, np.random.rand(64), [0.1] * 64]),
     )
     def test_deduplicate_cosine(self, mock_embed):
-        docs = ["doc1", "doc1_dup", "doc2", "doc1_dup2"]
+        docs = ["doc1", "doc1_dup", "a second document that is different", "doc1_dup2"]
         deduplicated_docs = self.model.deduplicate(docs, threshold=0.9)
         self.assertEqual(len(deduplicated_docs), 2)
         self.assertIn("doc1", deduplicated_docs)
-        self.assertIn("doc2", deduplicated_docs)
+        self.assertIn("a second document that is different", deduplicated_docs)
 
     @patch.object(
         WordLlamaInference,