multimodal insert ability develop

peng3307165 · peng3307165 · commit 9fce2e47a45b · 2024-04-17T10:35:50.000+08:00
diff --git a/flask4modelcache.py b/flask4modelcache.py
@@ -177,7 +177,6 @@ def user_backend():
         return json.dumps(result)
 
     if request_type == 'register':
-        # iat_type = param_dict.get("iat_type")
         response = adapter.ChatCompletion.create_register(
             model=model
         )
diff --git a/modelcache/adapter_mm/adapter.py b/modelcache/adapter_mm/adapter.py
@@ -3,8 +3,8 @@
 
 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
+from modelcache.adapter_mm.adapter_remove import adapt_remove
+from modelcache.adapter_mm.adapter_register import adapt_register
 
 
 class ChatCompletion(object):
@@ -30,7 +30,8 @@ def create_mm_insert(cls, *args, **kwargs):
                 **kwargs
             )
         except Exception as e:
-            return str(e)
+            # return str(e)
+            raise e
 
     @classmethod
     def create_mm_remove(cls, *args, **kwargs):
@@ -51,7 +52,7 @@ def create_mm_register(cls, *args, **kwargs):
                 **kwargs
             )
         except Exception as e:
-            return str(e)
+            raise e
 
 
 def construct_resp_from_cache(return_message, return_query):
diff --git a/modelcache/adapter_mm/adapter_insert.py b/modelcache/adapter_mm/adapter_insert.py
@@ -1,4 +1,8 @@
 # -*- coding: utf-8 -*-
+import time
+import requests
+import base64
+import numpy as np
 from modelcache import cache
 from modelcache.utils.error import NotInitError
 from modelcache.utils.time import time_cal
@@ -15,26 +19,77 @@ def adapt_insert(*args, **kwargs):
     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(
+    pre_embedding_data_dict = chat_cache.mm_insert_pre_embedding_func(
         kwargs,
         extra_param=context.get("pre_embedding_func", None),
         prompts=chat_cache.config.prompts,
     )
+
+    print('pre_embedding_data_dict: {}'.format(pre_embedding_data_dict))
     chat_info = kwargs.pop("chat_info", [])
     llm_data = chat_info[-1]['answer']
 
+    pre_embedding_text = '###'.join(pre_embedding_data_dict['text'])
+    pre_embedding_image_url = pre_embedding_data_dict['imageUrl']
+    pre_embedding_image_raw = pre_embedding_data_dict['imageRaw']
+    pre_embedding_image_id = pre_embedding_data_dict.get('imageId', None)
+
+    if pre_embedding_image_url and pre_embedding_image_raw:
+        raise ValueError("Both pre_embedding_image_url and pre_embedding_image_raw cannot be non-empty at the same time.")
+
+    if pre_embedding_image_url:
+        url_start_time = time.time()
+        response = requests.get(pre_embedding_image_url)
+        image_data = response.content
+        pre_embedding_image = base64.b64encode(image_data).decode('utf-8')
+        get_image_time = '{}s'.format(round(time.time() - url_start_time, 2))
+        print('get_image_time: {}'.format(get_image_time))
+    elif pre_embedding_image_raw:
+        pre_embedding_image = pre_embedding_image_raw
+    else:
+        pre_embedding_image = None
+        if not pre_embedding_text:
+            raise ValueError(
+                "Both pre_embedding_image_url and pre_embedding_image_raw are empty. Please provide at least one.")
+
+    data_dict = {'text': [pre_embedding_text], 'image': pre_embedding_image}
+    embedding_data = None
+    mm_type = None
+
     if cache_enable:
-        embedding_data = time_cal(
+        embedding_data_resp = time_cal(
             chat_cache.embedding_func,
-            func_name="embedding",
+            func_name="image_embedding",
             report_func=chat_cache.report.embedding,
-        )(pre_embedding_data)
+        )(data_dict)
+
+        image_embeddings = embedding_data_resp['image_embedding']
+        text_embeddings = embedding_data_resp['text_embeddings']
+
+        if len(image_embeddings) > 0 and len(image_embeddings) > 0:
+            image_embedding = np.array(image_embeddings[0])
+            text_embedding = text_embeddings[0]
+            embedding_data = np.concatenate((image_embedding, text_embedding))
+            mm_type = 'mm'
+        elif len(image_embeddings) > 0:
+            image_embedding = np.array(image_embeddings[0])
+            embedding_data = image_embedding
+            mm_type = 'image'
+        elif len(text_embeddings) > 0:
+            text_embedding = np.array(text_embeddings[0])
+            embedding_data = text_embedding
+            mm_type = 'text'
+        else:
+            raise ValueError('maya embedding service return both empty list, please check!')
 
     chat_cache.data_manager.save(
-        pre_embedding_data,
+        pre_embedding_text,
+        pre_embedding_image_url,
+        pre_embedding_image_id,
         llm_data,
         embedding_data,
         model=model,
-        extra_param=context.get("save_func", None)
+        mm_type=mm_type,
+        extra_param=context.get("mm_save_func", None)
     )
-    return 'success'
+    return 'success'
diff --git a/modelcache/adapter_mm/adapter_register.py b/modelcache/adapter_mm/adapter_register.py
@@ -5,9 +5,12 @@
 def adapt_register(*args, **kwargs):
     chat_cache = kwargs.pop("cache_obj", cache)
     model = kwargs.pop("model", None)
+    mm_type = kwargs.pop("mm_type", None)
     if model is None or len(model) == 0:
         return ValueError('')
 
-    register_resp = chat_cache.data_manager.create_index(model)
+    print('mm_type: {}'.format(mm_type))
+    print('model: {}'.format(model))
+    register_resp = chat_cache.data_manager.create_index(model, mm_type)
     print('register_resp: {}'.format(register_resp))
     return register_resp
diff --git a/modelcache/manager_mm/data_manager.py b/modelcache/manager_mm/data_manager.py
@@ -25,18 +25,21 @@
 class DataManager(metaclass=ABCMeta):
     """DataManager manage the cache data, including save and search"""
 
+    # @abstractmethod
+    # def save(self, question, answer, embedding_data, **kwargs):
+    #     pass
+
     @abstractmethod
-    def save(self, question, answer, embedding_data, **kwargs):
+    def save(self, text, image_url, image_id,  answer, embedding, **kwargs):
         pass
 
     @abstractmethod
     def save_query_resp(self, query_resp_dict, **kwargs):
         pass
 
     @abstractmethod
-    def import_data(
-        self, questions: List[Any], answers: List[Any], embedding_datas: List[Any], model:Any
-    ):
+    def import_data(self, texts: List[Any], image_urls: List[Any], image_ids: List[Any], answers: List[Answer],
+                    embeddings: List[Any], model: Any, iat_type: Any):
         pass
 
     @abstractmethod
@@ -89,21 +92,20 @@ def init(self):
                 f"You don't have permission to access this file <{self.data_path}>."
             )
 
-    def save(self, question, answer, embedding_data, **kwargs):
-        if isinstance(question, Question):
-            question = question.content
-        self.data[embedding_data] = (question, answer, embedding_data)
+    # def save(self, question, answer, embedding_data, **kwargs):
+    #     if isinstance(question, Question):
+    #         question = question.content
+    #     self.data[embedding_data] = (question, answer, embedding_data)
+
+    def save(self, text, image_url, image_id,  answer, embedding, **kwargs):
+        pass
 
     def save_query_resp(self, query_resp_dict, **kwargs):
         pass
 
-    def import_data(
-        self, questions: List[Any], answers: List[Any], embedding_datas: List[Any], model: Any
-    ):
-        if len(questions) != len(answers) or len(questions) != len(embedding_datas):
-            raise ParamError("Make sure that all parameters have the same length")
-        for i, embedding_data in enumerate(embedding_datas):
-            self.data[embedding_data] = (questions[i], answers[i], embedding_datas[i])
+    def import_data(self, texts: List[Any], image_urls: List[Any], image_ids: List[Any], answers: List[Answer],
+                    embeddings: List[Any], model: Any, iat_type: Any):
+        pass
 
     def get_scalar_data(self, res_data, **kwargs) -> CacheData:
         return CacheData(question=res_data[0], answers=res_data[1])
@@ -158,9 +160,15 @@ def __init__(
         self.v = v
         self.o = o
 
-    def save(self, question, answer, embedding_data, **kwargs):
+    # def save(self, question, answer, embedding_data, **kwargs):
+    #     model = kwargs.pop("model", None)
+    #     self.import_data([question], [answer], [embedding_data], model)
+
+    def save(self, text, image_url, image_id,  answer, embedding, **kwargs):
         model = kwargs.pop("model", None)
-        self.import_data([question], [answer], [embedding_data], model)
+        mm_type = kwargs.pop("mm_type", None)
+        self.import_data([text], [image_url], [image_id], [answer],
+                             [embedding], model, mm_type)
 
     def save_query_resp(self, query_resp_dict, **kwargs):
         save_query_start_time = time.time()
@@ -190,36 +198,38 @@ def _process_question_data(self, question: Union[str, Question]):
 
         return Question(question)
 
-    def import_data(
-        self, questions: List[Any], answers: List[Answer], embedding_datas: List[Any], model: Any
-    ):
-        if len(questions) != len(answers) or len(questions) != len(embedding_datas):
+    def import_data(self, texts: List[Any], image_urls: List[Any], image_ids: List[Any], answers: List[Answer],
+                    embeddings: List[Any], model: Any, iat_type: Any):
+        if len(texts) != len(answers):
             raise ParamError("Make sure that all parameters have the same length")
         cache_datas = []
 
-        embedding_datas = [
-            normalize(embedding_data) for embedding_data in embedding_datas
+        embeddings = [
+            normalize(text_embedding) for text_embedding in embeddings
         ]
 
-        for i, embedding_data in enumerate(embedding_datas):
+        # print('embedding_datas: {}'.format(embedding_datas))
+        for i, embedding in enumerate(embeddings):
             if self.o is not None:
                 ans = self._process_answer_data(answers[i])
             else:
                 ans = answers[i]
-
-            question = questions[i]
-            embedding_data = embedding_data.astype("float32")
-            cache_datas.append([ans, question, embedding_data, model])
-
-        ids = self.s.batch_insert(cache_datas)
-        logging.info('ids: {}'.format(ids))
-        self.v.mul_add(
+            text = texts[i]
+            image_url = image_urls[i]
+            image_id = image_ids[i]
+            # iat_embedding = embedding.astype("float32")
+            cache_datas.append([ans, text, image_url, image_id, model])
+
+        # ids = self.s.batch_multimodal_insert(cache_datas)
+        ids = self.s.batch_iat_insert(cache_datas)
+        # self.v.multimodal_add(
+        self.v.iat_add(
             [
-                VectorData(id=ids[i], data=embedding_data)
-                for i, embedding_data in enumerate(embedding_datas)
+                VectorData(id=ids[i], data=embedding)
+                for i, embedding in enumerate(embeddings)
             ],
-            model
-
+            model,
+            iat_type
         )
 
     def get_scalar_data(self, res_data, **kwargs) -> Optional[CacheData]:
@@ -256,8 +266,8 @@ def delete(self, id_list, **kwargs):
         return {'status': 'success', 'milvus': 'delete_count: '+str(v_delete_count),
                 'mysql': 'delete_count: '+str(s_delete_count)}
 
-    def create_index(self, model, **kwargs):
-        return self.v.create(model)
+    def create_index(self, model, mm_type, **kwargs):
+        return self.v.create(model, mm_type)
 
     def truncate(self, model_name):
         # drop vector base data
diff --git a/modelcache/manager_mm/factory.py b/modelcache/manager_mm/factory.py
@@ -1,7 +1,7 @@
 # -*- coding: utf-8 -*-
 from typing import Union, Callable
-from modelcache.manager import CacheBase, VectorBase, ObjectBase
-from modelcache.manager.data_manager import SSDataManager, MapDataManager
+from modelcache.manager_mm import CacheBase, VectorBase, ObjectBase
+from modelcache.manager_mm.data_manager import SSDataManager, MapDataManager
 
 
 def get_data_manager(
@@ -25,26 +25,3 @@ def get_data_manager(
         object_base = ObjectBase(name=object_base)
     assert cache_base and vector_base
     return SSDataManager(cache_base, vector_base, object_base, max_size, clean_size, eviction)
-
-
-def get_data_manager_mm(
-    cache_base: Union[CacheBase, str] = None,
-    vector_base: Union[VectorBase, str] = None,
-    object_base: Union[ObjectBase, str] = None,
-    max_size: int = 1000,
-    clean_size: int = None,
-    eviction: str = "LRU",
-    data_path: str = "data_map.txt",
-    get_data_container: Callable = None,
-):
-    if not cache_base and not vector_base:
-        return MapDataManager(data_path, max_size, get_data_container)
-
-    if isinstance(cache_base, str):
-        cache_base = CacheBase(name=cache_base)
-    if isinstance(vector_base, str):
-        vector_base = VectorBase(name=vector_base)
-    if isinstance(object_base, str):
-        object_base = ObjectBase(name=object_base)
-    assert cache_base and vector_base
-    return SSDataManager(cache_base, vector_base, object_base, max_size, clean_size, eviction)
diff --git a/modelcache/manager_mm/scalar_data/sql_storage.py b/modelcache/manager_mm/scalar_data/sql_storage.py
@@ -36,25 +36,24 @@ def create(self):
 
     def _insert(self, data: List):
         answer = data[0]
-        question = data[1]
-        embedding_data = data[2]
-        model = data[3]
+        text = data[1]
+        image_url = data[2]
+        image_id = data[3]
+        model = data[4]
         answer_type = 0
-        embedding_data = embedding_data.tobytes()
-
-        table_name = "cache_codegpt_answer"
-        insert_sql = "INSERT INTO {} (question, answer, answer_type, model, embedding_data) VALUES (%s, %s, %s, %s, _binary%s)".format(table_name)
 
+        table_name = "multimodal_answer"
+        insert_sql = "INSERT INTO {} (question_text, image_url, image_id, answer, answer_type, model) VALUES (%s, %s, %s, %s, %s, %s)".format(table_name)
         conn = self.pool.connection()
         try:
             with conn.cursor() as cursor:
-                # 执行插入数据操作
-                values = (question, answer, answer_type, model, embedding_data)
+                # data insert operation
+                values = (text, image_url, image_id, answer, answer_type, model)
                 cursor.execute(insert_sql, values)
                 conn.commit()
                 id = cursor.lastrowid
         finally:
-            # 关闭连接，将连接返回给连接池
+            # Close the connection and return it back to the connection pool
             conn.close()
         return id
 
diff --git a/modelcache/manager_mm/vector_data/base.py b/modelcache/manager_mm/vector_data/base.py
@@ -31,7 +31,7 @@ def delete(self, ids) -> bool:
         pass
 
     @abstractmethod
-    def rebuild_col(self, model):
+    def rebuild_idx(self, model):
         pass
 
     def flush(self):
diff --git a/modelcache/manager_mm/vector_data/redis.py b/modelcache/manager_mm/vector_data/redis.py
diff --git a/modelcache/utils/index_util.py b/modelcache/utils/index_util.py
diff --git a/multicache_serving.py b/multicache_serving.py

Original file line number	Diff line number	Diff line change
`@@ -177,7 +177,6 @@ def user_backend():`
`177`	`177`	`return json.dumps(result)`
`178`	`178`
`179`	`179`	`if request_type == 'register':`
`180`		`- # iat_type = param_dict.get("iat_type")`
`181`	`180`	`response = adapter.ChatCompletion.create_register(`
`182`	`181`	`model=model`
`183`	`182`	`)`