1  """
 2  Reranker module tests
 3  """
 4  
 5  import unittest
 6  
 7  from txtai import Embeddings
 8  from txtai.pipeline import Reranker, Similarity
 9  
10  
11  class TestReranker(unittest.TestCase):
12      """
13      Reranker tests.
14      """
15  
16      @classmethod
17      def setUpClass(cls):
18          """
19          Create single labels instance.
20          """
21  
22          cls.data = [
23              "US tops 5 million confirmed virus cases",
24              "Canada's last fully intact ice shelf has suddenly collapsed, forming a Manhattan-sized iceberg",
25              "Beijing mobilises invasion craft along coast as Taiwan tensions escalate",
26              "The National Park Service warns against sacrificing slower friends in a bear attack",
27              "Maine man wins $1M from $25 lottery ticket",
28              "Make huge profits without work, earn up to $100,000 a day",
29          ]
30  
31      def testRanker(self):
32          """
33          Test re-ranking pipeline
34          """
35  
36          embeddings = Embeddings(content=True)
37          embeddings.index(self.data)
38  
39          similarity = Similarity("neuml/colbert-bert-tiny", lateencode=True)
40  
41          ranker = Reranker(embeddings, similarity)
42          self.assertEqual(ranker("lottery winner")[0]["id"], "4")