multimodel cache

Author: peng3307165

.gitignore

@@ -136,4 +136,6 @@ dmypy.json
 /flask_server
 *.bin
 **/modelcache_serving.py
-*.ini
+*.ini
+
+**/maya_embedding_service

modelcache/adapter_mm/__init__.py

@@ -0,0 +1 @@
+# -*- coding: utf-8 -*-

modelcache/adapter_mm/adapter.py

@@ -0,0 +1,63 @@
+# -*- coding: utf-8 -*-
+import logging
+
+from modelcache.adapter_mm.adapter_query import adapt_query
+from modelcache.adapter_mm.adapter_insert import adapt_insert
+from modelcache.adapter.adapter_remove import adapt_remove
+from modelcache.adapter.adapter_register import adapt_register
+
+
+class ChatCompletion(object):
+    """Openai ChatCompletion Wrapper"""
+    @classmethod
+    def create_mm_query(cls, *args, **kwargs):
+        def cache_data_convert(cache_data, cache_query):
+            return construct_resp_from_cache(cache_data, cache_query)
+        try:
+            return adapt_query(
+                cache_data_convert,
+                *args,
+                **kwargs
+            )
+        except Exception as e:
+            return str(e)
+
+    @classmethod
+    def create_mm_insert(cls, *args, **kwargs):
+        try:
+            return adapt_insert(
+                *args,
+                **kwargs
+            )
+        except Exception as e:
+            return str(e)
+
+    @classmethod
+    def create_mm_remove(cls, *args, **kwargs):
+        try:
+            return adapt_remove(
+                *args,
+                **kwargs
+            )
+        except Exception as e:
+            logging.info('adapt_remove_e: {}'.format(e))
+            return str(e)
+
+    @classmethod
+    def create_mm_register(cls, *args, **kwargs):
+        try:
+            return adapt_register(
+                *args,
+                **kwargs
+            )
+        except Exception as e:
+            return str(e)
+
+
+def construct_resp_from_cache(return_message, return_query):
+    return {
+        "modelcache": True,
+        "hitQuery": return_query,
+        "data": return_message,
+        "errorCode": 0
+    }

modelcache/adapter_mm/adapter_insert.py

@@ -0,0 +1,40 @@
+# -*- coding: utf-8 -*-
+from modelcache import cache
+from modelcache.utils.error import NotInitError
+from modelcache.utils.time import time_cal
+
+
+def adapt_insert(*args, **kwargs):
+    chat_cache = kwargs.pop("cache_obj", cache)
+    model = kwargs.pop("model", None)
+    require_object_store = kwargs.pop("require_object_store", False)
+    if require_object_store:
+        assert chat_cache.data_manager.o, "Object store is required for adapter."
+    if not chat_cache.has_init:
+        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:
+        embedding_data = time_cal(
+            chat_cache.embedding_func,
+            func_name="embedding",
+            report_func=chat_cache.report.embedding,
+        )(pre_embedding_data)
+
+    chat_cache.data_manager.save(
+        pre_embedding_data,
+        llm_data,
+        embedding_data,
+        model=model,
+        extra_param=context.get("save_func", None)
+    )
+    return 'success'

modelcache/adapter_mm/adapter_query.py

@@ -0,0 +1,148 @@
+# -*- coding: utf-8 -*-
+import logging
+import time
+from modelcache import cache
+from modelcache.utils.error import NotInitError
+from modelcache.utils.time import time_cal
+from modelcache.processor.pre import multi_analysis
+
+
+def adapt_query(cache_data_convert, *args, **kwargs):
+    chat_cache = kwargs.pop("cache_obj", cache)
+    scope = kwargs.pop("scope", None)
+    model = scope['model']
+    if not chat_cache.has_init:
+        raise NotInitError()
+    cache_enable = chat_cache.cache_enable_func(*args, **kwargs)
+    context = kwargs.pop("cache_context", {})
+    embedding_data = None
+    cache_factor = kwargs.pop("cache_factor", 1.0)
+    pre_embedding_data = chat_cache.query_pre_embedding_func(
+        kwargs,
+        extra_param=context.get("pre_embedding_func", None),
+        prompts=chat_cache.config.prompts,
+    )
+
+    if cache_enable:
+        embedding_data = time_cal(
+            chat_cache.embedding_func,
+            func_name="embedding",
+            report_func=chat_cache.report.embedding,
+        )(pre_embedding_data)
+
+    if cache_enable:
+        cache_data_list = time_cal(
+            chat_cache.data_manager.search,
+            func_name="milvus_search",
+            report_func=chat_cache.report.search,
+        )(
+            embedding_data,
+            extra_param=context.get("search_func", None),
+            top_k=kwargs.pop("top_k", -1),
+            model=model
+        )
+        cache_answers = []
+        cache_questions = []
+        cache_ids = []
+        similarity_threshold = chat_cache.config.similarity_threshold
+        similarity_threshold_long = chat_cache.config.similarity_threshold_long
+
+        min_rank, max_rank = chat_cache.similarity_evaluation.range()
+        rank_threshold = (max_rank - min_rank) * similarity_threshold * cache_factor
+        rank_threshold_long = (max_rank - min_rank) * similarity_threshold_long * cache_factor
+        rank_threshold = (
+            max_rank
+            if rank_threshold > max_rank
+            else min_rank
+            if rank_threshold < min_rank
+            else rank_threshold
+        )
+        rank_threshold_long = (
+            max_rank
+            if rank_threshold_long > max_rank
+            else min_rank
+            if rank_threshold_long < min_rank
+            else rank_threshold_long
+        )
+
+        if cache_data_list is None or len(cache_data_list) == 0:
+            rank_pre = -1.0
+        else:
+            cache_data_dict = {'search_result': cache_data_list[0]}
+            rank_pre = chat_cache.similarity_evaluation.evaluation(
+                None,
+                cache_data_dict,
+                extra_param=context.get("evaluation_func", None),
+            )
+        if rank_pre < rank_threshold:
+            return
+
+        for cache_data in cache_data_list:
+            primary_id = cache_data[1]
+            start_time = time.time()
+            ret = chat_cache.data_manager.get_scalar_data(
+                cache_data, extra_param=context.get("get_scalar_data", None)
+            )
+            if ret is None:
+                continue
+
+            if "deps" in context and hasattr(ret.question, "deps"):
+                eval_query_data = {
+                    "question": context["deps"][0]["data"],
+                    "embedding": None
+                }
+                eval_cache_data = {
+                    "question": ret.question.deps[0].data,
+                    "answer": ret.answers[0].answer,
+                    "search_result": cache_data,
+                    "embedding": None,
+                }
+            else:
+                eval_query_data = {
+                    "question": pre_embedding_data,
+                    "embedding": embedding_data,
+                }
+
+                eval_cache_data = {
+                    "question": ret[0],
+                    "answer": ret[1],
+                    "search_result": cache_data,
+                    "embedding": None
+                }
+            rank = chat_cache.similarity_evaluation.evaluation(
+                eval_query_data,
+                eval_cache_data,
+                extra_param=context.get("evaluation_func", None),
+            )
+
+            if len(pre_embedding_data) <= 256:
+                if rank_threshold <= rank:
+                    cache_answers.append((rank, ret[1]))
+                    cache_questions.append((rank, ret[0]))
+                    cache_ids.append((rank, primary_id))
+            else:
+                if rank_threshold_long <= rank:
+                    cache_answers.append((rank, ret[1]))
+                    cache_questions.append((rank, ret[0]))
+                    cache_ids.append((rank, primary_id))
+        cache_answers = sorted(cache_answers, key=lambda x: x[0], reverse=True)
+        cache_questions = sorted(cache_questions, key=lambda x: x[0], reverse=True)
+        cache_ids = sorted(cache_ids, key=lambda x: x[0], reverse=True)
+        if len(cache_answers) != 0:
+            return_message = chat_cache.post_process_messages_func(
+                [t[1] for t in cache_answers]
+            )
+            return_query = chat_cache.post_process_messages_func(
+                [t[1] for t in cache_questions]
+            )
+            return_id = chat_cache.post_process_messages_func(
+                [t[1] for t in cache_ids]
+            )
+            # 更新命中次数
+            try:
+                chat_cache.data_manager.update_hit_count(return_id)
+            except Exception:
+                logging.info('update_hit_count except, please check!')
+
+            chat_cache.report.hint_cache()
+            return cache_data_convert(return_message, return_query)

modelcache/adapter_mm/adapter_register.py

@@ -0,0 +1,13 @@
+# -*- coding: utf-8 -*-
+from modelcache import cache
+
+
+def adapt_register(*args, **kwargs):
+    chat_cache = kwargs.pop("cache_obj", cache)
+    model = kwargs.pop("model", None)
+    if model is None or len(model) == 0:
+        return ValueError('')
+
+    register_resp = chat_cache.data_manager.create_index(model)
+    print('register_resp: {}'.format(register_resp))
+    return register_resp

modelcache/adapter_mm/adapter_remove.py

@@ -0,0 +1,26 @@
+# -*- coding: utf-8 -*-
+from modelcache import cache
+from modelcache.utils.error import NotInitError, RemoveError
+
+
+def adapt_remove(*args, **kwargs):
+    chat_cache = kwargs.pop("cache_obj", cache)
+    model = kwargs.pop("model", None)
+    remove_type = kwargs.pop("remove_type", None)
+    require_object_store = kwargs.pop("require_object_store", False)
+    if require_object_store:
+        assert chat_cache.data_manager.o, "Object store is required for adapter."
+    if not chat_cache.has_init:
+        raise NotInitError()
+
+    # delete data
+    if remove_type == 'delete_by_id':
+        id_list = kwargs.pop("id_list", [])
+        resp = chat_cache.data_manager.delete(id_list, model=model)
+    elif remove_type == 'truncate_by_model':
+        resp = chat_cache.data_manager.truncate(model)
+    else:
+        # resp = "remove_type_error"
+        raise RemoveError()
+    return resp
+

modelcache/core.py

@@ -17,9 +17,10 @@ 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
+        self.embedding_concurrent_func = None
         self.data_manager: Optional[DataManager] = None
         self.similarity_evaluation: Optional[SimilarityEvaluation] = None
         self.post_process_messages_func = None
@@ -30,9 +31,10 @@ 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,
+            embedding_concurrent_func=string_embedding,
             data_manager: DataManager = get_data_manager(),
             similarity_evaluation=ExactMatchEvaluation(),
             post_process_messages_func=first,
@@ -41,9 +43,10 @@ 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
+        self.embedding_concurrent_func = embedding_concurrent_func
         self.data_manager: DataManager = data_manager
         self.similarity_evaluation = similarity_evaluation
         self.post_process_messages_func = post_process_messages_func

modelcache/manager/vector_data/redis.py

@@ -21,19 +21,25 @@ def __init__(
         port: str = "6379",
         username: str = "",
         password: str = "",
-        dimension: int = 0,
+        # 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 Redis vector store."
             )
         self._client = Redis(
             host=host, port=int(port), username=username, password=password
         )
         self.top_k = top_k
-        self.dimension = dimension
+        # 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:"
 

modelcache/manager_mm/__init__.py

@@ -0,0 +1,5 @@
+# -*- coding: utf-8 -*-
+from modelcache.manager_mm.scalar_data import CacheBase
+from modelcache.manager_mm.vector_data import VectorBase
+from modelcache.manager_mm.object_data import ObjectBase
+from modelcache.manager_mm.factory import get_data_manager