Added handling for bulk insert

Added a flag for normalize because the model we use already normalizes

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

Author: olgaoznovich

docker-compose.yaml

@@ -14,7 +14,7 @@ services:
       - ./data/mysql/db:/var/lib/mysql
       - ./data/mysql/my.cnf:/etc/mysql/conf.d/my.cnf
       - ./data/mysql/init:/docker-entrypoint-initdb.d
-    restart: on-failure
+#    restart: on-failure
     networks:
       - modelcache
 
@@ -36,15 +36,15 @@ services:
       - 19530:19530
       - 9091:9091
       - 2379:2379
-    healthcheck:
-      test: ["CMD", "curl", "-f", "http://localhost:9091/healthz"]
-      interval: 30s
-      start_period: 90s
-      timeout: 20s
-      retries: 3
+#    healthcheck:
+#      test: ["CMD", "curl", "-f", "http://localhost:9091/healthz"]
+#      interval: 30s
+#      start_period: 90s
+#      timeout: 20s
+#      retries: 3
     networks:
       - modelcache
-    restart: on-failure  
+#    restart: on-failure
     command: milvus run standalone
 
 #  modelcache:

flask4modelcache.py

@@ -6,7 +6,7 @@
 import json
 from modelcache import cache
 from modelcache.adapter import adapter
-from modelcache.manager import CacheBase, VectorBase, get_data_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
@@ -30,8 +30,10 @@ def save_query_info(result, model, query, delta_time_log):
 def response_hitquery(cache_resp):
     return cache_resp['hitQuery']
 
-
 data2vec = Data2VecAudio()
+embedding_func = data2vec.to_embeddings
+dimension = data2vec.dimension
+
 mysql_config = configparser.ConfigParser()
 mysql_config.read('modelcache/config/mysql_config.ini')
 
@@ -48,7 +50,7 @@ def response_hitquery(cache_resp):
 # chromadb_config.read('modelcache/config/chromadb_config.ini')
 
 data_manager = get_data_manager(CacheBase("mysql", config=mysql_config),
-                                VectorBase("milvus", dimension=data2vec.dimension, milvus_config=milvus_config))
+                                VectorBase("milvus", dimension=dimension, milvus_config=milvus_config))
 
 
 # data_manager = get_data_manager(CacheBase("mysql", config=mysql_config),
@@ -57,9 +59,8 @@ def response_hitquery(cache_resp):
 # data_manager = get_data_manager(CacheBase("mysql", config=mysql_config),
 #                                 VectorBase("redis", dimension=data2vec.dimension, redis_config=redis_config))
 
-
 cache.init(
-    embedding_func=data2vec.to_embeddings,
+    embedding_func=embedding_func,
     data_manager=data_manager,
     similarity_evaluation=SearchDistanceEvaluation(),
     query_pre_embedding_func=query_multi_splicing,

model/.gitignore

@@ -0,0 +1,2 @@
+*.tflite
+text2vec-base-chinese/*

model/clone_model_repository.bat

@@ -0,0 +1 @@
+curl -o embedder.tflite https://storage.googleapis.com/mediapipe-models/text_embedder/bert_embedder/float32/1/bert_embedder.tflite

model/download_bert_embedder.bat

@@ -14,26 +14,31 @@ def adapt_insert(*args, **kwargs):
         raise NotInitError()
     cache_enable = chat_cache.cache_enable_func(*args, **kwargs)
     context = kwargs.pop("cache_context", {})
-    embedding_data = None
-    pre_embedding_data = chat_cache.insert_pre_embedding_func(
-        kwargs,
-        extra_param=context.get("pre_embedding_func", None),
-        prompts=chat_cache.config.prompts,
-    )
     chat_info = kwargs.pop("chat_info", [])
-    llm_data = chat_info[-1]['answer']
 
-    if cache_enable:
+    pre_embedding_data_list = []
+    embedding_data_list = []
+    llm_data_list = []
+
+    for row in chat_info:
+        pre_embedding_data = chat_cache.insert_pre_embedding_func(
+            row,
+            extra_param=context.get("pre_embedding_func", None),
+            prompts=chat_cache.config.prompts,
+        )
+        pre_embedding_data_list.append(pre_embedding_data)
+        llm_data_list.append(row['answer'])
         embedding_data = time_cal(
             chat_cache.embedding_func,
             func_name="embedding",
             report_func=chat_cache.report.embedding,
         )(pre_embedding_data)
+        embedding_data_list.append(embedding_data)
 
     chat_cache.data_manager.save(
-        pre_embedding_data,
-        llm_data,
-        embedding_data,
+        pre_embedding_data_list,
+        llm_data_list,
+        embedding_data_list,
         model=model,
         extra_param=context.get("save_func", None)
     )

modelcache/adapter/adapter_insert.py

@@ -7,6 +7,9 @@
 import cachetools
 from abc import abstractmethod, ABCMeta
 from typing import List, Any, Optional, Union
+
+from numpy import ndarray
+
 from modelcache.manager.scalar_data.base import (
     CacheStorage,
     CacheData,
@@ -21,6 +24,7 @@
 from modelcache.manager.eviction_manager import EvictionManager
 from modelcache.utils.log import modelcache_log
 
+NORMALIZE = True
 
 class DataManager(metaclass=ABCMeta):
     """DataManager manage the cache data, including save and search"""
@@ -158,9 +162,9 @@ def __init__(
         self.v = v
         self.o = o
 
-    def save(self, question, answer, embedding_data, **kwargs):
+    def save(self, questions: List[any], answers: List[any], embedding_datas: List[any], **kwargs):
         model = kwargs.pop("model", None)
-        self.import_data([question], [answer], [embedding_data], model)
+        self.import_data(questions, answers, embedding_datas, model)
 
     def save_query_resp(self, query_resp_dict, **kwargs):
         save_query_start_time = time.time()
@@ -197,9 +201,10 @@ def import_data(
             raise ParamError("Make sure that all parameters have the same length")
         cache_datas = []
 
-        embedding_datas = [
-            normalize(embedding_data) for embedding_data in embedding_datas
-        ]
+        if NORMALIZE:
+            embedding_datas = [
+                normalize(embedding_data) for embedding_data in embedding_datas
+            ]
 
         for i, embedding_data in enumerate(embedding_datas):
             if self.o is not None:
@@ -212,11 +217,9 @@ def import_data(
             cache_datas.append([ans, question, embedding_data, model])
 
         ids = self.s.batch_insert(cache_datas)
+        datas_ = [VectorData(id=ids[i], data=embedding_data.astype("float32")) for i, embedding_data in enumerate(embedding_datas)]
         self.v.mul_add(
-            [
-                VectorData(id=ids[i], data=embedding_data)
-                for i, embedding_data in enumerate(embedding_datas)
-            ],
+            datas_,
             model
 
         )
@@ -235,7 +238,8 @@ def hit_cache_callback(self, res_data, **kwargs):
 
     def search(self, embedding_data, **kwargs):
         model = kwargs.pop("model", None)
-        embedding_data = normalize(embedding_data)
+        if NORMALIZE:
+            embedding_data = normalize(embedding_data)
         top_k = kwargs.get("top_k", -1)
         return self.v.search(data=embedding_data, top_k=top_k, model=model)
 

modelcache/manager/data_manager.py

@@ -64,7 +64,7 @@ def query_multi_splicing(data: Dict[str, Any], **_: Dict[str, Any]) -> Any:
 
 
 def insert_multi_splicing(data: Dict[str, Any], **_: Dict[str, Any]) -> Any:
-    insert_query_list = data.get("chat_info")[-1]['query']
+    insert_query_list = data['query']
     return multi_splicing(insert_query_list)
 
 

modelcache/processor/pre.py

@@ -18,4 +18,6 @@ chromadb==0.5.23
 elasticsearch==7.10.0
 snowflake-id==1.0.2
 flagembedding==1.3.4
-cryptography==45.0.2
+cryptography==45.0.2
+mediapipe==0.10.21
+protobuf==4.25.8