File tree Expand file tree Collapse file tree 1 file changed +28
-0
lines changed Expand file tree Collapse file tree 1 file changed +28
-0
lines changed Original file line number Diff line number Diff 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 ()
You can’t perform that action at this time.
0 commit comments