|
| 1 | +# -*- coding: utf-8 -*- |
| 2 | +import time |
| 3 | +import requests |
| 4 | +import numpy as np |
| 5 | +import base64 |
| 6 | +from modelcache_mm import cache |
| 7 | +from modelcache_mm.utils.error import NotInitError |
| 8 | +from modelcache_mm.utils.error import MultiTypeError |
| 9 | +from modelcache_mm.utils.time import time_cal |
| 10 | + |
| 11 | + |
| 12 | +def adapt_query(cache_data_convert, *args, **kwargs): |
| 13 | + chat_cache = kwargs.pop("cache_obj", cache) |
| 14 | + scope = kwargs.pop("scope", None) |
| 15 | + model = scope['model'] |
| 16 | + if not chat_cache.has_init: |
| 17 | + raise NotInitError() |
| 18 | + |
| 19 | + cache_enable = chat_cache.cache_enable_func(*args, **kwargs) |
| 20 | + context = kwargs.pop("cache_context", {}) |
| 21 | + cache_factor = kwargs.pop("cache_factor", 1.0) |
| 22 | + |
| 23 | + pre_embedding_data_dict = chat_cache.query_pre_embedding_func( |
| 24 | + kwargs, |
| 25 | + extra_param=context.get("pre_embedding_func", None), |
| 26 | + prompts=chat_cache.config.prompts, |
| 27 | + ) |
| 28 | + |
| 29 | + pre_embedding_text = '###'.join(pre_embedding_data_dict['text']) |
| 30 | + pre_embedding_image_raw = pre_embedding_data_dict['imageRaw'] |
| 31 | + pre_embedding_image_url = pre_embedding_data_dict['imageUrl'] |
| 32 | + pre_multi_type = pre_embedding_data_dict['multiType'] |
| 33 | + # print('pre_embedding_image_url: {}'.format(pre_embedding_image_url)) |
| 34 | + # print('pre_embedding_text: {}'.format(pre_embedding_text)) |
| 35 | + |
| 36 | + # 判断逻辑 |
| 37 | + if pre_multi_type == 'IMG_TEXT': |
| 38 | + if pre_embedding_image_raw and pre_embedding_image_url: |
| 39 | + raise ValueError( |
| 40 | + "Both pre_embedding_imageUrl and pre_embedding_imageRaw cannot be non-empty at the same time.") |
| 41 | + if pre_embedding_image_url: |
| 42 | + url_start_time = time.time() |
| 43 | + response = requests.get(pre_embedding_image_url) |
| 44 | + image_data = response.content |
| 45 | + pre_embedding_image = base64.b64encode(image_data).decode('utf-8') |
| 46 | + get_image_time = '{}s'.format(round(time.time() - url_start_time, 2)) |
| 47 | + print('get_image_time: {}'.format(get_image_time)) |
| 48 | + elif pre_embedding_image_raw: |
| 49 | + pre_embedding_image = pre_embedding_image_raw |
| 50 | + else: |
| 51 | + raise ValueError( |
| 52 | + "Both pre_embedding_imageUrl and pre_embedding_imageRaw are empty. Please provide at least one.") |
| 53 | + data_dict = {'text': [pre_embedding_text], 'image': pre_embedding_image} |
| 54 | + # print('data_dict: {}'.format(data_dict)) |
| 55 | + elif pre_multi_type == 'TEXT': |
| 56 | + data_dict = {'text': [pre_embedding_text], 'image': None} |
| 57 | + else: |
| 58 | + raise MultiTypeError |
| 59 | + # print('data_dict: {}'.format(data_dict)) |
| 60 | + |
| 61 | + embedding_data = None |
| 62 | + mm_type = None |
| 63 | + if cache_enable: |
| 64 | + if pre_multi_type == 'IMG_TEXT': |
| 65 | + embedding_data_resp = time_cal( |
| 66 | + chat_cache.embedding_concurrent_func, |
| 67 | + func_name="iat_embedding", |
| 68 | + report_func=chat_cache.report.embedding, |
| 69 | + )(data_dict) |
| 70 | + else: |
| 71 | + embedding_data_resp = time_cal( |
| 72 | + chat_cache.embedding_func, |
| 73 | + func_name="iat_embedding", |
| 74 | + report_func=chat_cache.report.embedding, |
| 75 | + )(data_dict) |
| 76 | + image_embeddings = embedding_data_resp['image_embedding'] |
| 77 | + text_embeddings = embedding_data_resp['text_embeddings'] |
| 78 | + |
| 79 | + if len(image_embeddings) > 0 and len(image_embeddings) > 0: |
| 80 | + image_embedding = np.array(image_embeddings[0]) |
| 81 | + text_embedding = np.array(text_embeddings[0]) |
| 82 | + embedding_data = np.concatenate((image_embedding, text_embedding)) |
| 83 | + mm_type = 'mm' |
| 84 | + elif len(image_embeddings) > 0: |
| 85 | + image_embedding = np.array(image_embeddings[0]) |
| 86 | + embedding_data = image_embedding |
| 87 | + mm_type = 'image' |
| 88 | + elif len(text_embeddings) > 0: |
| 89 | + text_embedding = np.array(text_embeddings[0]) |
| 90 | + embedding_data = text_embedding |
| 91 | + mm_type = 'text' |
| 92 | + else: |
| 93 | + raise ValueError('maya embedding service return both empty list, please check!') |
| 94 | + |
| 95 | + if cache_enable: |
| 96 | + cache_data_list = time_cal( |
| 97 | + chat_cache.data_manager.search, |
| 98 | + func_name="vector_search", |
| 99 | + report_func=chat_cache.report.search, |
| 100 | + )( |
| 101 | + embedding_data, |
| 102 | + extra_param=context.get("search_func", None), |
| 103 | + top_k=kwargs.pop("top_k", -1), |
| 104 | + model=model, |
| 105 | + mm_type=pre_multi_type, |
| 106 | + ) |
| 107 | + |
| 108 | + cache_answers = [] |
| 109 | + cache_questions = [] |
| 110 | + cache_image_urls = [] |
| 111 | + cache_image_ids = [] |
| 112 | + cache_ids = [] |
| 113 | + similarity_threshold = chat_cache.config.similarity_threshold |
| 114 | + similarity_threshold_long = chat_cache.config.similarity_threshold_long |
| 115 | + |
| 116 | + min_rank, max_rank = chat_cache.similarity_evaluation.range() |
| 117 | + rank_threshold = (max_rank - min_rank) * similarity_threshold * cache_factor |
| 118 | + rank_threshold_long = (max_rank - min_rank) * similarity_threshold_long * cache_factor |
| 119 | + rank_threshold = ( |
| 120 | + max_rank |
| 121 | + if rank_threshold > max_rank |
| 122 | + else min_rank |
| 123 | + if rank_threshold < min_rank |
| 124 | + else rank_threshold |
| 125 | + ) |
| 126 | + rank_threshold_long = ( |
| 127 | + max_rank |
| 128 | + if rank_threshold_long > max_rank |
| 129 | + else min_rank |
| 130 | + if rank_threshold_long < min_rank |
| 131 | + else rank_threshold_long |
| 132 | + ) |
| 133 | + |
| 134 | + if cache_data_list is None or len(cache_data_list) == 0: |
| 135 | + rank_pre = -1.0 |
| 136 | + else: |
| 137 | + cache_data_dict = {'search_result': cache_data_list[0]} |
| 138 | + rank_pre = chat_cache.similarity_evaluation.evaluation( |
| 139 | + None, |
| 140 | + cache_data_dict, |
| 141 | + extra_param=context.get("evaluation_func", None), |
| 142 | + ) |
| 143 | + |
| 144 | + print('rank_pre: {}'.format(rank_pre)) |
| 145 | + print('rank_threshold: {}'.format(rank_threshold)) |
| 146 | + if rank_pre < rank_threshold: |
| 147 | + return |
| 148 | + |
| 149 | + for cache_data in cache_data_list: |
| 150 | + print('cache_data: {}'.format(cache_data)) |
| 151 | + primary_id = cache_data[1] |
| 152 | + ret = chat_cache.data_manager.get_scalar_data( |
| 153 | + cache_data, extra_param=context.get("get_scalar_data", None) |
| 154 | + ) |
| 155 | + if ret is None: |
| 156 | + continue |
| 157 | + |
| 158 | + if "deps" in context and hasattr(ret.question, "deps"): |
| 159 | + eval_query_data = { |
| 160 | + "question": context["deps"][0]["data"], |
| 161 | + "embedding": None |
| 162 | + } |
| 163 | + eval_cache_data = { |
| 164 | + "question": ret.question.deps[0].data, |
| 165 | + "answer": ret.answers[0].answer, |
| 166 | + "search_result": cache_data, |
| 167 | + "embedding": None, |
| 168 | + } |
| 169 | + else: |
| 170 | + eval_query_data = { |
| 171 | + "question": pre_embedding_text, |
| 172 | + "embedding": embedding_data, |
| 173 | + } |
| 174 | + |
| 175 | + eval_cache_data = { |
| 176 | + "question": ret[0], |
| 177 | + "image_url": ret[1], |
| 178 | + "image_raw": ret[2], |
| 179 | + "answer": ret[3], |
| 180 | + "search_result": cache_data, |
| 181 | + "embedding": None |
| 182 | + } |
| 183 | + rank = chat_cache.similarity_evaluation.evaluation( |
| 184 | + eval_query_data, |
| 185 | + eval_cache_data, |
| 186 | + extra_param=context.get("evaluation_func", None), |
| 187 | + ) |
| 188 | + print('rank_threshold: {}'.format(rank_threshold)) |
| 189 | + print('rank_threshold_long: {}'.format(rank_threshold_long)) |
| 190 | + print('rank: {}'.format(rank)) |
| 191 | + |
| 192 | + if len(pre_embedding_text) <= 50: |
| 193 | + if rank_threshold <= rank: |
| 194 | + cache_answers.append((rank, ret[3])) |
| 195 | + cache_image_urls.append((rank, ret[1])) |
| 196 | + cache_image_ids.append((rank, ret[2])) |
| 197 | + cache_questions.append((rank, ret[0])) |
| 198 | + cache_ids.append((rank, primary_id)) |
| 199 | + else: |
| 200 | + if rank_threshold_long <= rank: |
| 201 | + cache_answers.append((rank, ret[3])) |
| 202 | + cache_image_urls.append((rank, ret[1])) |
| 203 | + cache_image_ids.append((rank, ret[2])) |
| 204 | + cache_questions.append((rank, ret[0])) |
| 205 | + cache_ids.append((rank, primary_id)) |
| 206 | + |
| 207 | + cache_answers = sorted(cache_answers, key=lambda x: x[0], reverse=True) |
| 208 | + cache_image_urls = sorted(cache_image_urls, key=lambda x: x[0], reverse=True) |
| 209 | + cache_image_ids = sorted(cache_image_ids, key=lambda x: x[0], reverse=True) |
| 210 | + cache_questions = sorted(cache_questions, key=lambda x: x[0], reverse=True) |
| 211 | + cache_ids = sorted(cache_ids, key=lambda x: x[0], reverse=True) |
| 212 | + |
| 213 | + print('cache_answers: {}'.format(cache_answers)) |
| 214 | + |
| 215 | + if len(cache_answers) != 0: |
| 216 | + return_message = chat_cache.post_process_messages_func( |
| 217 | + [t[1] for t in cache_answers] |
| 218 | + ) |
| 219 | + return_image_url = chat_cache.post_process_messages_func( |
| 220 | + [t[1] for t in cache_image_urls] |
| 221 | + ) |
| 222 | + return_image_id = chat_cache.post_process_messages_func( |
| 223 | + [t[1] for t in cache_image_ids] |
| 224 | + ) |
| 225 | + return_query = chat_cache.post_process_messages_func( |
| 226 | + [t[1] for t in cache_questions] |
| 227 | + ) |
| 228 | + return_id = chat_cache.post_process_messages_func( |
| 229 | + [t[1] for t in cache_ids] |
| 230 | + ) |
| 231 | + # 更新命中次数 |
| 232 | + try: |
| 233 | + chat_cache.data_manager.update_hit_count(return_id) |
| 234 | + except Exception: |
| 235 | + print('update_hit_count except, please check!') |
| 236 | + |
| 237 | + chat_cache.report.hint_cache() |
| 238 | + return_query_dict = {"image_url": return_image_url, "image_id": return_image_id, "question": return_query} |
| 239 | + return cache_data_convert(return_message, return_query_dict) |
0 commit comments