Added custom pre embedding func and some other improvements

Co-authored-by: olgaoznovich <ol.oznovich@gmail.com>
Co-authored-by: Yuval-Roth <rothyuv@post.bgu.ac.il>

Author: Yuval-Roth

flask4modelcache.py

@@ -10,8 +10,7 @@
 from modelcache.manager.vector_data import manager
 from modelcache.manager import CacheBase, VectorBase, get_data_manager, data_manager
 from modelcache.similarity_evaluation.distance import SearchDistanceEvaluation
-from modelcache.processor.pre import query_multi_splicing
-from modelcache.processor.pre import insert_multi_splicing
+from modelcache.processor.pre import query_multi_splicing,insert_multi_splicing, query_with_role
 from concurrent.futures import ThreadPoolExecutor
 from modelcache.utils.model_filter import model_blacklist_filter
 from modelcache.embedding import Data2VecAudio
@@ -36,13 +35,17 @@ def response_hitquery(cache_resp):
 
 if manager.MPNet_base:
     mpnet_base = MPNet_Base()
-    embedding_func = lambda x: mpnet_base.embedding_func(x)
+    embedding_func = mpnet_base.to_embeddings
     dimension =  mpnet_base.dimension
     data_manager.NORMALIZE = False
+    query_pre_embedding_func=query_with_role
+    insert_pre_embedding_func=query_with_role
 else:
     data2vec = Data2VecAudio()
     embedding_func = data2vec.to_embeddings
     dimension = data2vec.dimension
+    query_pre_embedding_func=query_multi_splicing
+    insert_pre_embedding_func=insert_multi_splicing
 
 mysql_config = configparser.ConfigParser()
 mysql_config.read('modelcache/config/mysql_config.ini')
@@ -95,8 +98,8 @@ def response_hitquery(cache_resp):
     embedding_func=embedding_func,
     data_manager=data_manager,
     similarity_evaluation=SearchDistanceEvaluation(),
-    query_pre_embedding_func=query_multi_splicing,
-    insert_pre_embedding_func=insert_multi_splicing,
+    query_pre_embedding_func=query_pre_embedding_func,
+    insert_pre_embedding_func=insert_pre_embedding_func,
 )
 
 global executor

modelcache/embedding/mpnet_base.py

@@ -5,7 +5,7 @@ def __init__(self):
         self.dimension = 768
         self.model = SentenceTransformer('sentence-transformers/all-mpnet-base-v2')
 
-    def embedding_func(self, *args, **kwargs):
+    def to_embeddings(self, *args, **kwargs):
         if not args:
             raise ValueError("No word provided for embedding.")
         embeddings = self.model.encode(args)

modelcache/processor/pre.py

@@ -4,7 +4,7 @@
 
 
 def insert_last_content(data: Dict[str, Any], **_: Dict[str, Any]) -> Any:
-    return data.get("chat_info")[-1]["query"]
+    return data.get("query")[-1]["content"]
 
 
 def query_last_content(data: Dict[str, Any], **_: Dict[str, Any]) -> Any:
@@ -67,6 +67,11 @@ def insert_multi_splicing(data: Dict[str, Any], **_: Dict[str, Any]) -> Any:
     insert_query_list = data['query']
     return multi_splicing(insert_query_list)
 
+def query_with_role(data: Dict[str, Any], **_: Dict[str, Any]) -> Any:
+    query = data["query"][-1]
+    content = query["content"]
+    role = query["role"]
+    return role+": "+content
 
 def multi_splicing(data_list) -> Any:
     result_str = ""