Fix the wrong return value of onnx similarity evaluation

Signed-off-by: SimFG <bang.fu@zilliz.com>

Author: SimFG

gptcache/adapter/adapter.py

@@ -166,7 +166,7 @@ def adapt(llm_handler, cache_data_convert, update_cache_callback, *args, **kwarg
                 rank,
             )
             if rank_threshold <= rank:
-                cache_answers.append((rank, cache_data.answers[0].answer, search_data))
+                cache_answers.append((float(rank), cache_data.answers[0].answer, search_data))
                 chat_cache.data_manager.hit_cache_callback(search_data)
         cache_answers = sorted(cache_answers, key=lambda x: x[0], reverse=True)
         answers_dict = dict((d[1], d[2]) for d in cache_answers)
@@ -397,7 +397,7 @@ async def aadapt(llm_handler, cache_data_convert, update_cache_callback, *args,
                 rank,
             )
             if rank_threshold <= rank:
-                cache_answers.append((rank, cache_data.answers[0].answer, search_data))
+                cache_answers.append((float(rank), cache_data.answers[0].answer, search_data))
                 chat_cache.data_manager.hit_cache_callback(search_data)
         cache_answers = sorted(cache_answers, key=lambda x: x[0], reverse=True)
         answers_dict = dict((d[1], d[2]) for d in cache_answers)

gptcache/manager/scalar_data/manager.py

@@ -1,4 +1,3 @@
-from gptcache.manager.scalar_data.mongo import MongoStorage
 from gptcache.utils import import_sql_client
 from gptcache.utils.error import NotFoundError
 
@@ -85,6 +84,8 @@ def get(name, **kwargs):
                 table_len_config=table_len_config,
             )
         elif name == "mongo":
+            from gptcache.manager.scalar_data.mongo import MongoStorage
+
             return MongoStorage(
                 host=kwargs.get("mongo_host", "localhost"),
                 port=kwargs.get("mongo_port", 27017),

gptcache/processor/post.py

@@ -1,5 +1,6 @@
 import random
 from typing import List, Any
+
 import numpy
 
 from gptcache.utils import softmax

gptcache/similarity_evaluation/distance.py

@@ -1,4 +1,5 @@
 from typing import Tuple, Dict, Any
+
 from gptcache.similarity_evaluation import SimilarityEvaluation
 
 

gptcache/similarity_evaluation/onnx.py

@@ -1,11 +1,13 @@
 from typing import Dict, List, Tuple, Any
+
 import numpy as np
+
+from gptcache.similarity_evaluation import SimilarityEvaluation
 from gptcache.utils import (
     import_onnxruntime,
     import_huggingface_hub,
     import_huggingface,
 )
-from gptcache.similarity_evaluation import SimilarityEvaluation
 
 import_onnxruntime()
 import_huggingface_hub()
@@ -130,4 +132,4 @@ def inference(self, reference: str, candidates: List[str]) -> np.ndarray:
         }
         ort_outputs = self.ort_session.run(None, ort_inputs)
         scores = ort_outputs[0][:, 1]
-        return scores
+        return float(scores[0])

gptcache/utils/softmax.py

@@ -3,7 +3,7 @@
 
 def softmax(x: list):
     x = np.array(x)
-    assert len(x.shape) == 1, f"Expect to get a shape of (len,) but got {x.shape}."
+    assert len(x.shape) == 1, f"Expect to get a shape of (len,) but got {x.shape}, x value: {x}."
     max_val = x.max()
     e_x = np.exp(x - max_val)
     return e_x / e_x.sum()

tests/unit_tests/similarity_evaluation/test_evaluation_onnx.py

@@ -17,6 +17,7 @@ def _test_evaluation(evaluation):
     candidate_2 = "how old are you?"
 
     score = evaluation.evaluation({"question": query}, {"question": candidate_1})
+    assert isinstance(score, float), type(score)
     assert score > 0.8
 
     score = evaluation.evaluation({"question": query}, {"question": candidate_2})