mm cache func develope

Author: peng3307165

modelcache/core.py

@@ -17,6 +17,8 @@ class Cache:
     def __init__(self):
         self.has_init = False
         self.cache_enable_func = None
+        self.query_pre_embedding_func = None
+        self.insert_pre_embedding_func = None
         self.mm_query_pre_embedding_func = None
         self.mm_insert_pre_embedding_func = None
         self.embedding_func = None
@@ -31,6 +33,8 @@ def __init__(self):
     def init(
             self,
             cache_enable_func=cache_all,
+            query_pre_embedding_func=None,
+            insert_pre_embedding_func=None,
             mm_query_pre_embedding_func=None,
             mm_insert_pre_embedding_func=None,
             embedding_func=string_embedding,
@@ -43,6 +47,8 @@ def init(
     ):
         self.has_init = True
         self.cache_enable_func = cache_enable_func
+        self.query_pre_embedding_func = query_pre_embedding_func
+        self.insert_pre_embedding_func = insert_pre_embedding_func
         self.mm_query_pre_embedding_func = mm_query_pre_embedding_func
         self.mm_insert_pre_embedding_func = mm_insert_pre_embedding_func
         self.embedding_func = embedding_func

modelcache/manager/vector_data/redis.py

@@ -21,10 +21,7 @@ def __init__(
         port: str = "6379",
         username: str = "",
         password: str = "",
-        # dimension: int = 0,
         dimension: int = 0,
-        i_dimension: int = 0,
-        t_dimension: int = 0,
         top_k: int = 1,
         namespace: str = "",
     ):
@@ -37,8 +34,6 @@ def __init__(
         )
         self.top_k = top_k
         self.dimension = dimension
-        self.i_dimension = i_dimension
-        self.t_dimension = t_dimension
         self.namespace = namespace
         self.doc_prefix = f"{self.namespace}doc:"
 

modelcache/manager_mm/vector_data/manager.py

@@ -69,7 +69,7 @@ def get(name, **kwargs):
                 local_data=local_data
             )
         elif name == "redis":
-            from modelcache.manager.vector_data.redis import RedisVectorStore
+            from modelcache.manager_mm.vector_data.redis import RedisVectorStore
             redis_config = kwargs.get("redis_config")
 
             mm_dimension = kwargs.get("mm_dimension", DIMENSION)

modelcache/manager_mm/vector_data/redis.py

@@ -21,19 +21,23 @@ def __init__(
         port: str = "6379",
         username: str = "",
         password: str = "",
-        dimension: int = 0,
+        mm_dimension: int = 0,
+        i_dimension: int = 0,
+        t_dimension: int = 0,
         top_k: int = 1,
         namespace: str = "",
     ):
-        if dimension <= 0:
+        if mm_dimension <= 0:
             raise ValueError(
-                f"invalid `dim` param: {dimension} in the Milvus vector store."
+                f"invalid `dim` param: {mm_dimension} in the Milvus vector store."
             )
         self._client = Redis(
             host=host, port=int(port), username=username, password=password
         )
         self.top_k = top_k
-        self.dimension = dimension
+        self.mm_dimension = mm_dimension
+        self.i_dimension = i_dimension
+        self.t_dimension = t_dimension
         self.namespace = namespace
         self.doc_prefix = f"{self.namespace}doc:"
 
@@ -47,8 +51,16 @@ def _check_index_exists(self, index_name: str) -> bool:
         modelcache_log.info("Index already exists")
         return True
 
-    def create_index(self, index_name, index_prefix):
-        dimension = self.dimension
+    def create_index(self, index_name, mm_type, index_prefix):
+        # dimension = self.dimension
+        if mm_type == 'IMG_TEXT':
+            dimension = self.mm_dimension
+        elif mm_type == 'IMG':
+            dimension = self.i_dimension
+        elif mm_type == 'TEXT':
+            dimension = self.t_dimension
+        else:
+            raise ValueError('dimension type exception')
         print('dimension: {}'.format(dimension))
         if self._check_index_exists(index_name):
             modelcache_log.info(
@@ -77,13 +89,17 @@ def create_index(self, index_name, index_prefix):
             )
             return 'create_success'
 
-    def mul_add(self, datas: List[VectorData], model=None):
-        # pipe = self._client.pipeline()
+    def mul_add(self, datas: List[VectorData], model=None, mm_type=None):
         for data in datas:
             id: int = data.id
             embedding = data.data.astype(np.float32).tobytes()
+
+            collection_name = get_collection_iat_name(model, mm_type)
+            index_prefix = get_collection_iat_prefix(model, mm_type)
+
             id_field_name = "data_id"
             embedding_field_name = "data_vector"
+
             obj = {id_field_name: id, embedding_field_name: embedding}
             index_prefix = get_index_prefix(model)
             self._client.hset(f"{index_prefix}{id}", mapping=obj)

modelcache/utils/index_util.py

@@ -7,3 +7,27 @@ def get_index_name(model):
 
 def get_index_prefix(model):
     return 'prefix' + '_' + model
+
+
+def get_mm_index_name(model, iat_type):
+    if iat_type not in ['IMG_TEXT', 'iat', 'IMG', 'image', 'TEXT', 'text']:
+        raise ValueError('iat_type is not normal!')
+    if iat_type == 'IMG_TEXT':
+        iat_type = 'iat'
+    elif iat_type == 'IMG':
+        iat_type = 'image'
+    elif iat_type == 'TEXT':
+        iat_type = 'text'
+    return 'multicache' + '_' + model + '_' + iat_type
+
+
+def get_collection_iat_prefix(model, iat_type, table_suffix):
+    if iat_type not in ['IMG_TEXT', 'iat', 'IMG', 'image', 'TEXT', 'text']:
+        raise ValueError('iat_type is not normal!')
+    if iat_type == 'IMG_TEXT':
+        iat_type = 'iat'
+    elif iat_type == 'IMG':
+        iat_type = 'image'
+    elif iat_type == 'TEXT':
+        iat_type = 'text'
+    return 'prefix' + '_' + model + '_' + iat_type + '_' + table_suffix

multicache_serving.py

@@ -27,6 +27,8 @@
 from concurrent.futures import ThreadPoolExecutor
 from modelcache.maya_embedding_service.maya_multi_embedding_service import get_embedding_multi
 from modelcache.maya_embedding_service.maya_multi_embedding_service import get_embedding_multi_concurrent_sin
+from modelcache.processor.pre import query_multi_splicing
+from modelcache.processor.pre import insert_multi_splicing
 
 
 def save_query_info(result, model, query, delta_time_log):
@@ -73,6 +75,8 @@ def __init__(self):
             embedding_concurrent_func=get_embedding_multi_concurrent_sin,
             data_manager=data_manager,
             similarity_evaluation=SearchDistanceEvaluation(),
+            query_pre_embedding_func=query_multi_splicing,
+            insert_pre_embedding_func=insert_multi_splicing,
             mm_insert_pre_embedding_func=mm_insert_dict,
             mm_query_pre_embedding_func=mm_query_dict,
         )