Skip to content

Commit 61ee6b4

Browse files
committed
reproduce ignoring top_n
1 parent b87b41e commit 61ee6b4

File tree

1 file changed

+28
-0
lines changed

1 file changed

+28
-0
lines changed

backend/python/rerankers/test.py

Lines changed: 28 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -86,5 +86,33 @@ def test_rerank(self):
8686
except Exception as err:
8787
print(err)
8888
self.fail("Reranker service failed")
89+
finally:
90+
self.tearDown()
91+
92+
def test_rerank_crop(self):
93+
"""
94+
This method tests if the embeddings are generated successfully
95+
"""
96+
try:
97+
self.setUp()
98+
with grpc.insecure_channel("localhost:50051") as channel:
99+
stub = backend_pb2_grpc.BackendStub(channel)
100+
request = backend_pb2.RerankRequest(
101+
query="I love you",
102+
documents=["I hate you", "I really like you", "I hate ignoring top_n"],
103+
top_n=2
104+
)
105+
response = stub.LoadModel(backend_pb2.ModelOptions(Model="cross-encoder"))
106+
self.assertTrue(response.success)
107+
108+
rerank_response = stub.Rerank(request)
109+
print(rerank_response.results[0])
110+
self.assertIsNotNone(rerank_response.results)
111+
self.assertEqual(len(rerank_response.results), 2)
112+
self.assertEqual(rerank_response.results[0].text, "I really like you")
113+
self.assertEqual(rerank_response.results[1].text, "I hate you")
114+
except Exception as err:
115+
print(err)
116+
self.fail("Reranker service failed")
89117
finally:
90118
self.tearDown()

0 commit comments

Comments
 (0)