+
+ | Module |
+ Function |
+
+
+
+ | ModelCache |
+ GPTCache |
+
+
+ | Basic Interface |
+ Data query interface |
+ ☑ |
+ ☑ |
+
+
+ | Data writing interface |
+ ☑ |
+ ☑ |
+
+
+ | Embedding |
+ Embedding model configuration |
+ ☑ |
+ ☑ |
+
+
+ | Large model embedding layer |
+ ☑ |
+ |
+
+
+ | BERT model long text processing |
+ ☑ |
+ |
+
+
+ | Large model invocation |
+ Decoupling from large models |
+ ☑ |
+ |
+
+
+ | Local loading of embedding model |
+ ☑ |
+ |
+
+
+ | Data isolation |
+ Model data isolation |
+ ☑ |
+ ☑ |
+
+
+ | Hyperparameter isolation |
+ |
+ |
+
+
+ | Databases |
+ MySQL |
+ ☑ |
+ ☑ |
+
+
+ | Milvus |
+ ☑ |
+ ☑ |
+
+
+ | OceanBase |
+ ☑ |
+ |
+
+
+ | Session management |
+ Single-turn dialogue |
+ ☑ |
+ ☑ |
+
+
+ | System commands |
+ ☑ |
+ |
+
+
+ | Multi-turn dialogue |
+ ☑ |
+ |
+
+
+ | Data management |
+ Data persistence |
+ ☑ |
+ ☑ |
+
+
+ | One-click cache clearance |
+ ☑ |
+ |
+
+
+ | Tenant management |
+ Support for multi-tenancy |
+ ☑ |
+ |
+
+
+ | Milvus multi-collection capability |
+ ☑ |
+ |
+
+
+ | Other |
+ Long-short dialogue distinction |
+ ☑ |
+ |
+
+
diff --git a/docs/2.model-cache-features.md b/docs/2.model-cache-features.md
new file mode 100644
index 0000000..2294d03
--- /dev/null
+++ b/docs/2.model-cache-features.md
@@ -0,0 +1,29 @@
+# ModelCache features
+
+This topic describes ModelCache features. In ModelCache, we incorporated the core principles of GPTCache. ModelCache has four modules: adapter, embedding, similarity, and data_manager.
+
+- The adapter module orchestrates the business logic for various tasks, integrate the embedding, similarity, and data_manager modules.
+- The embedding module converts text into semantic vector representations, and transforms user queries into vectors.
+- The rank module ranks and evaluate the similarity of recalled vectors.
+- The data_manager module manages the databases.
+
+To make ModelCache more suitable for industrial use, we made several improvements to its architecture and functionality:
+
+- [x] Architectural adjustment (lightweight integration):
+ - Embedded into LLM products using a Redis-like caching mode.
+ - Provided semantic caching without interfering with LLM calls, security audits, and other functions.
+ - Compatible with all LLM services.
+- [x] Multiple model loading:
+ - Supported local embedding model loading, and resolved Hugging Face network connectivity issues.
+ - Supported loading embedding layers from various pre-trained models.
+- [x] Data isolation
+ - Environment isolation: Read different database configurations based on the environment. Isolate development, staging, and production environments.
+ - Multi-tenant data isolation: Dynamically create collections based on models for data isolation, addressing data separation issues in multi-model/service scenarios within large language model products.
+- [x] Supported system instruction: Adopted a concatenation approach to resolve issues with system instructions in the prompt paradigm.
+- [x] Long and short text differentiation: Long texts bring more challenges for similarity assessment. Added differentiation between long and short texts, allowing for separate threshold configurations.
+- [x] Milvus performance optimization: Adjusted Milvus consistency level to "Session" level for better performance.
+- [x] Data management:
+ - One-click cache clearing to enable easy data management after model upgrades.
+ - Recall of hit queries for subsequent data analysis and model iteration reference.
+ - Asynchronous log write-back for data analysis and statistics.
+ - Added model field and data statistics field to enhance features.
diff --git a/docs/3.model-cache-quick-start.md b/docs/3.model-cache-quick-start.md
new file mode 100644
index 0000000..7cd3903
--- /dev/null
+++ b/docs/3.model-cache-quick-start.md
@@ -0,0 +1,97 @@
+# Quick start
+
+This topic describes how to set up and use ModelCache.
+
+You can find the start script in `flask4modelcache.py` and `flask4modelcache_demo.py`.
+
+- `flask4modelcache_demo.py`: A quick test service that embeds SQLite and FAISS. No database configuration required.
+- `flask4modelcache.py`: The standard service that requires MySQL and Milvus configuration.
+
+## Dependencies
+
+- Python: V3.8 or above
+- Package installation
+
+ ```shell
+ pip install -r requirements.txt
+ ```
+
+## Start service
+
+### Start demo
+
+1. Download the embedding model bin file from [Hugging Face](https://huggingface.co/shibing624/text2vec-base-chinese/tree/main). Place it in the `model/text2vec-base-chinese` folder.
+2. Start the backend service:
+
+ ```shell
+ cd CodeFuse-ModelCache
+ ```
+
+ ```shell
+ python flask4modelcache_demo.py
+ ```
+
+### Start standard service
+
+Before you start standard service, do these steps:
+
+1. Install MySQL and import the SQL file from `reference_doc/create_table.sql`.
+2. Install vector database Milvus.
+3. Configure database access in:
+ - `modelcache/config/milvus_config.ini`
+ - `modelcache/config/mysql_config.ini`
+4. Download the embedding model bin file from [Hugging Face](https://huggingface.co/shibing624/text2vec-base-chinese/tree/main). Put it in `model/text2vec-base-chinese`.
+5. Start the backend service:
+
+ ```bash
+ python flask4modelcache.py
+ ```
+
+## Visit the service
+
+The service provides three core RESTful API functionalities: Cache-Writing, Cache-Querying, and Cache-Clearing.
+
+### Write cache
+
+```python
+import json
+import requests
+url = 'http://127.0.0.1:5000/modelcache'
+type = 'insert'
+scope = {"model": "CODEGPT-1008"}
+chat_info = [{"query": [{"role": "system", "content": "You are an AI code assistant and you must provide neutral and harmless answers to help users solve code-related problems."}, {"role": "user", "content": "你是谁?"}],
+ "answer": "Hello, I am an intelligent assistant. How can I assist you?"}]
+data = {'type': type, 'scope': scope, 'chat_info': chat_info}
+headers = {"Content-Type": "application/json"}
+res = requests.post(url, headers=headers, json=json.dumps(data))
+```
+
+### Query cache
+
+```python
+import json
+import requests
+url = 'http://127.0.0.1:5000/modelcache'
+type = 'query'
+scope = {"model": "CODEGPT-1008"}
+query = [{"role": "system", "content": "You are an AI code assistant and you must provide neutral and harmless answers to help users solve code-related problems."}, {"role": "user", "content": "Who are you?"}]
+data = {'type': type, 'scope': scope, 'query': query}
+
+headers = {"Content-Type": "application/json"}
+res = requests.post(url, headers=headers, json=json.dumps(data))
+```
+
+### Clear cache
+
+```python
+import json
+import requests
+url = 'http://127.0.0.1:5000/modelcache'
+type = 'remove'
+scope = {"model": "CODEGPT-1008"}
+remove_type = 'truncate_by_model'
+data = {'type': type, 'scope': scope, 'remove_type': remove_type}
+
+headers = {"Content-Type": "application/json"}
+res = requests.post(url, headers=headers, json=json.dumps(data))
+```
diff --git a/docs/4.create-cache.md b/docs/4.create-cache.md
new file mode 100644
index 0000000..e0506ba
--- /dev/null
+++ b/docs/4.create-cache.md
@@ -0,0 +1,42 @@
+# Create cache
+
+This topic describes how to create cache.
+
+## Default cache interface
+
+```py
+class Cache:
+ # ModelCache calls it whe you start the cache system
+ def __init__(self):
+ self.has_init = False
+ self.cache_enable_func = None
+ self.embedding_func = None
+ self.post_process_messages_func = None
+ self.config = Config()
+```
+
+This function embeds text into dense vectors for context similarity search. ModelCache supports these embedding context methods: Huggingface, ONNX, and SentenceTransformers. The default model is text2vec Hugging Face because it performs better for Chinese. Simply initialize your embedding function as `text2vec.to_embeddings`.
+
+```py
+data_manager = get_data_manager(CacheBase("mysql", config=mysql_config),
+ VectorBase("milvus", dimension=data2vec.dimension, milvus_config=milvus_config))
+
+cache.init(
+ embedding_func=data2vec.to_embeddings,
+ data_manager=data_manager,
+ similarity_evaluation=SearchDistanceEvaluation(),
+ query_pre_embedding_func=query_multi_splicing,
+ insert_pre_embedding_func=insert_multi_splicing,
+)
+```
+
+data_manager CacheVase stores all scalar data, such as original questions, prompts, answers, and access times. ModelCache supports multiple cache storages like SQLite, MySQL, and OceanBase. NoSQL databases will be supported in the future.
+
+data_manager VectorBase stores and searches all embedding vectors to find semantically similar results. ModelCache supports using vector search libraries like FAISS or vector databases like Milvus. More vector database and cloud service will be supported in the future.
+
+## Examples
+
+```py
+data_manager = get_data_manager(CacheBase("sqlite"), VectorBase("faiss", dimension=data2vec.dimension))
+data_manager = get_data_manager(CacheBase("oceanbase"), VectorBase("milvus", dimension=data2vec.dimension))
+```
diff --git a/docs/cache-service-cost-time-distribution.webp b/docs/cache-service-cost-time-distribution.webp
new file mode 100644
index 0000000..d22a1f4
Binary files /dev/null and b/docs/cache-service-cost-time-distribution.webp differ
diff --git a/docs/modelcache_modules_20240409.png b/docs/modelcache_modules_20240409.png
new file mode 100644
index 0000000..e7c4442
Binary files /dev/null and b/docs/modelcache_modules_20240409.png differ
diff --git a/docs/time-cost-comparison.webp b/docs/time-cost-comparison.webp
new file mode 100644
index 0000000..e00c014
Binary files /dev/null and b/docs/time-cost-comparison.webp differ
diff --git a/examples/embedding/__init__.py b/examples/embedding/__init__.py
new file mode 100644
index 0000000..40a96af
--- /dev/null
+++ b/examples/embedding/__init__.py
@@ -0,0 +1 @@
+# -*- coding: utf-8 -*-
diff --git a/examples/embedding/huggingface_tei_example.py b/examples/embedding/huggingface_tei_example.py
new file mode 100644
index 0000000..152834f
--- /dev/null
+++ b/examples/embedding/huggingface_tei_example.py
@@ -0,0 +1,17 @@
+# -*- coding: utf-8 -*-
+import sys
+sys.path.append(".")
+from modelcache.embedding.huggingface_tei import HuggingfaceTEI
+
+'''
+run tei server:
+text-embeddings-router --model-id BAAI/bge-large-zh-v1.5 --port 8080
+'''
+
+def run():
+ tei_instance = HuggingfaceTEI('http://127.0.0.1:8080/v1/embeddings', 'BAAI/bge-large-zh-v1.5')
+ print('dimenson', tei_instance.dimension)
+ print('embedding', tei_instance.to_embeddings('hello'))
+
+if __name__ == '__main__':
+ run()
\ No newline at end of file
diff --git a/examples/flask/llms_cache/__init__.py b/examples/flask/llms_cache/__init__.py
new file mode 100644
index 0000000..40a96af
--- /dev/null
+++ b/examples/flask/llms_cache/__init__.py
@@ -0,0 +1 @@
+# -*- coding: utf-8 -*-
diff --git a/examples/flask/data_insert.py b/examples/flask/llms_cache/data_insert.py
similarity index 92%
rename from examples/flask/data_insert.py
rename to examples/flask/llms_cache/data_insert.py
index 094082d..e26bec7 100644
--- a/examples/flask/data_insert.py
+++ b/examples/flask/llms_cache/data_insert.py
@@ -13,8 +13,8 @@ def run():
headers = {"Content-Type": "application/json"}
res = requests.post(url, headers=headers, json=json.dumps(data))
res_text = res.text
- print('res_text: {}'.format(res_text))
+ print("data_insert:", res.status_code, res_text)
if __name__ == '__main__':
run()
diff --git a/examples/flask/data_query.py b/examples/flask/llms_cache/data_query.py
similarity index 90%
rename from examples/flask/data_query.py
rename to examples/flask/llms_cache/data_query.py
index 3eccae3..6fe8313 100644
--- a/examples/flask/data_query.py
+++ b/examples/flask/llms_cache/data_query.py
@@ -13,8 +13,8 @@ def run():
headers = {"Content-Type": "application/json"}
res = requests.post(url, headers=headers, json=json.dumps(data))
res_text = res.text
- print('res_text: {}'.format(res_text))
+ print("data_query:", res.status_code, res_text)
if __name__ == '__main__':
run()
diff --git a/examples/flask/data_query_long.py b/examples/flask/llms_cache/data_query_long.py
similarity index 90%
rename from examples/flask/data_query_long.py
rename to examples/flask/llms_cache/data_query_long.py
index 2ebfcc8..7aa4f14 100644
--- a/examples/flask/data_query_long.py
+++ b/examples/flask/llms_cache/data_query_long.py
@@ -18,8 +18,8 @@ def run():
headers = {"Content-Type": "application/json"}
res = requests.post(url, headers=headers, json=json.dumps(data))
res_text = res.text
- print('res_text: {}'.format(res_text))
+ print("data_query_long:", res.status_code, res_text)
if __name__ == '__main__':
run()
diff --git a/examples/flask/llms_cache/register.py b/examples/flask/llms_cache/register.py
new file mode 100644
index 0000000..23f0849
--- /dev/null
+++ b/examples/flask/llms_cache/register.py
@@ -0,0 +1,20 @@
+# -*- coding: utf-8 -*-
+"""
+register index for redis
+"""
+import json
+import requests
+
+
+def run():
+ url = 'http://127.0.0.1:5000/modelcache'
+ type = 'register'
+ scope = {"model": "CODEGPT-1117"}
+ data = {'type': type, 'scope': scope}
+ headers = {"Content-Type": "application/json"}
+ res = requests.post(url, headers=headers, json=json.dumps(data))
+ res_text = res.text
+
+
+if __name__ == '__main__':
+ run()
\ No newline at end of file
diff --git a/examples/flask/multi_cache/__init__.py b/examples/flask/multi_cache/__init__.py
new file mode 100644
index 0000000..40a96af
--- /dev/null
+++ b/examples/flask/multi_cache/__init__.py
@@ -0,0 +1 @@
+# -*- coding: utf-8 -*-
diff --git a/examples/flask/multi_cache/data_insert.py b/examples/flask/multi_cache/data_insert.py
new file mode 100644
index 0000000..3d173b8
--- /dev/null
+++ b/examples/flask/multi_cache/data_insert.py
@@ -0,0 +1,31 @@
+# -*- coding: utf-8 -*-
+import time
+import json
+import uuid
+import requests
+
+
+def run():
+ url = 'http://127.0.0.1:5000/multicache'
+
+ request_type = 'insert'
+ scope = {"model": "multimodal_test"}
+ # UUID = "820b0052-d9d8-11ee-95f1-52775e3e6fd1" + "==>" + str(time.time())
+ UUID = str(uuid.uuid1()) + "==>" + str(time.time())
+ img_data = "https://img0.baidu.com/it/u=1436460262,4166266890&fm=253&fmt=auto&app=138&f=JPEG?w=500&h=282"
+ query = {'text': ['父母带着孩子来这个地方可能会有什么顾虑'],
+ 'imageRaw': '',
+ 'imageUrl': img_data,
+ 'imageId': 'ccc'}
+ answer = "应该注意小孩不要跑到铁轨上"
+ chat_info = [{"query": query, "answer": answer}]
+ data_dict = {'request_type': request_type, 'scope': scope, 'chat_info': chat_info, 'UUID': UUID}
+
+ headers = {"Content-Type": "application/json"}
+ res = requests.post(url, headers=headers, json=json.dumps(data_dict))
+ res_text = res.text
+ print('res_text: {}'.format(res_text))
+
+
+if __name__ == '__main__':
+ run()
diff --git a/examples/flask/multi_cache/data_query.py b/examples/flask/multi_cache/data_query.py
new file mode 100644
index 0000000..ccfc335
--- /dev/null
+++ b/examples/flask/multi_cache/data_query.py
@@ -0,0 +1,28 @@
+# -*- coding: utf-8 -*-
+import json
+import requests
+import uuid
+import time
+
+
+def run():
+ url = 'http://127.0.0.1:5000/multicache'
+ request_type = 'query'
+ UUID = str(uuid.uuid1()) + "==>" + str(time.time())
+ scope = {"model": "multimodal_test"}
+ img_data = "https://img0.baidu.com/it/u=1436460262,4166266890&fm=253&fmt=auto&app=138&f=JPEG?w=500&h=282"
+ query = {'text': ['父母带着孩子来这个地方可能会有什么顾虑'],
+ 'imageRaw': '',
+ 'imageUrl': img_data,
+ 'multiType': 'IMG_TEXT'}
+
+ data = {'request_type': request_type, 'scope': scope, 'query': query, 'UUID': UUID}
+
+ headers = {"Content-Type": "application/json"}
+ res = requests.post(url, headers=headers, json=json.dumps(data))
+ res_text = res.text
+ print('res_text: {}'.format(res_text))
+
+
+if __name__ == '__main__':
+ run()
diff --git a/examples/flask/multi_cache/register.py b/examples/flask/multi_cache/register.py
new file mode 100644
index 0000000..4ca830b
--- /dev/null
+++ b/examples/flask/multi_cache/register.py
@@ -0,0 +1,22 @@
+# -*- coding: utf-8 -*-
+"""
+register index for redis
+"""
+import json
+import requests
+
+
+def run():
+ url = 'http://127.0.0.1:5000/multicache'
+ request_type = 'register'
+ scope = {"model": "multimodal_test"}
+ type = 'IMG_TEXT'
+ data = {'request_type': request_type, 'scope': scope, 'type': type}
+ headers = {"Content-Type": "application/json"}
+ res = requests.post(url, headers=headers, json=json.dumps(data))
+ res_text = res.text
+ print('res_text: {}'.format(res_text))
+
+
+if __name__ == '__main__':
+ run()
\ No newline at end of file
diff --git a/examples/flask/multi_cache/remove.py b/examples/flask/multi_cache/remove.py
new file mode 100644
index 0000000..ffad449
--- /dev/null
+++ b/examples/flask/multi_cache/remove.py
@@ -0,0 +1,23 @@
+# -*- coding: utf-8 -*-
+"""
+register index for redis
+"""
+import json
+import requests
+
+
+def run():
+ url = 'http://127.0.0.1:5000/multicache'
+ request_type = 'remove'
+ scope = {"model": "multimodal_test"}
+ remove_type = 'truncate_by_model'
+ data = {'request_type': request_type, 'scope': scope, 'remove_type': remove_type}
+
+ headers = {"Content-Type": "application/json"}
+ res = requests.post(url, headers=headers, json=json.dumps(data))
+ res_text = res.text
+ print('res_text: {}'.format(res_text))
+
+
+if __name__ == '__main__':
+ run()
diff --git a/fastapi4modelcache.py b/fastapi4modelcache.py
new file mode 100644
index 0000000..b260eae
--- /dev/null
+++ b/fastapi4modelcache.py
@@ -0,0 +1,46 @@
+# -*- coding: utf-8 -*-
+import asyncio
+from contextlib import asynccontextmanager
+import uvicorn
+import json
+from fastapi.responses import JSONResponse
+from fastapi import FastAPI, Request
+from modelcache.cache import Cache
+from modelcache.embedding import EmbeddingModel
+
+@asynccontextmanager
+async def lifespan(app: FastAPI):
+ global cache
+ cache, _ = await Cache.init(
+ sql_storage="mysql",
+ vector_storage="milvus",
+ embedding_model=EmbeddingModel.HUGGINGFACE_ALL_MPNET_BASE_V2,
+ embedding_workers_num=2
+ )
+ yield
+
+app = FastAPI(lifespan=lifespan)
+cache: Cache = None
+
+@app.get("/welcome")
+async def first_fastapi():
+ return "hello, modelcache!"
+
+@app.post("/modelcache")
+async def user_backend(request: Request):
+
+ try:
+ request_data = await request.json()
+ except Exception:
+ result = {"errorCode": 400, "errorDesc": "bad request", "cacheHit": False, "delta_time": 0, "hit_query": '', "answer": ''}
+ return JSONResponse(status_code=400, content=result)
+
+ try:
+ return await cache.handle_request(request_data)
+ except Exception as e:
+ result = {"errorCode": 500, "errorDesc": str(e), "cacheHit": False, "delta_time": 0, "hit_query": '', "answer": ''}
+ cache.save_query_resp(result, model='', query='', delta_time=0)
+ return JSONResponse(status_code=500, content=result)
+
+if __name__ == '__main__':
+ uvicorn.run(app, host='0.0.0.0', port=5000, loop="asyncio", http="httptools")
diff --git a/fastapi4modelcache_demo.py b/fastapi4modelcache_demo.py
new file mode 100644
index 0000000..c1f6a3f
--- /dev/null
+++ b/fastapi4modelcache_demo.py
@@ -0,0 +1,46 @@
+# -*- coding: utf-8 -*-
+import asyncio
+from contextlib import asynccontextmanager
+import uvicorn
+import json
+from fastapi.responses import JSONResponse
+from fastapi import FastAPI, Request
+from modelcache.cache import Cache
+from modelcache.embedding import EmbeddingModel
+
+@asynccontextmanager
+async def lifespan(app: FastAPI):
+ global cache
+ cache, _ = await Cache.init(
+ sql_storage="sqlite",
+ vector_storage="faiss",
+ embedding_model=EmbeddingModel.HUGGINGFACE_ALL_MPNET_BASE_V2,
+ embedding_workers_num=2
+ )
+ yield
+
+app = FastAPI(lifespan=lifespan)
+cache: Cache = None
+
+@app.get("/welcome")
+async def first_fastapi():
+ return "hello, modelcache!"
+
+@app.post("/modelcache")
+async def user_backend(request: Request):
+
+ try:
+ request_data = await request.json()
+ except Exception:
+ result = {"errorCode": 400, "errorDesc": "bad request", "cacheHit": False, "delta_time": 0, "hit_query": '', "answer": ''}
+ return JSONResponse(status_code=400, content=result)
+
+ try:
+ return await cache.handle_request(request_data)
+ except Exception as e:
+ result = {"errorCode": 500, "errorDesc": str(e), "cacheHit": False, "delta_time": 0, "hit_query": '', "answer": ''}
+ cache.save_query_resp(result, model='', query='', delta_time=0)
+ return JSONResponse(status_code=500, content=result)
+
+if __name__ == '__main__':
+ uvicorn.run(app, host='0.0.0.0', port=5000, loop="asyncio", http="httptools")
diff --git a/flask4modelcache.py b/flask4modelcache.py
index 18c5a58..3d18518 100644
--- a/flask4modelcache.py
+++ b/flask4modelcache.py
@@ -1,175 +1,48 @@
# -*- coding: utf-8 -*-
-import time
-from datetime import datetime
-from flask import Flask, request
-import logging
-import configparser
-import json
-from modelcache import cache
-from modelcache.adapter import adapter
-from modelcache.manager import CacheBase, VectorBase, get_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
-from concurrent.futures import ThreadPoolExecutor
-from modelcache.utils.model_filter import model_blacklist_filter
-from modelcache.embedding import Data2VecAudio
+import asyncio
+from flask import Flask, request, jsonify
+from modelcache.cache import Cache
+from modelcache.embedding import EmbeddingModel
-# 创建一个Flask实例
-app = Flask(__name__)
+async def main():
-def response_text(cache_resp):
- return cache_resp['data']
+ # 创建一个Flask实例
+ app = Flask(__name__)
+ cache,loop = await Cache.init(
+ sql_storage="mysql",
+ vector_storage="milvus",
+ embedding_model=EmbeddingModel.HUGGINGFACE_ALL_MPNET_BASE_V2,
+ embedding_workers_num=2
+ )
-def save_query_info(result, model, query, delta_time_log):
- cache.data_manager.save_query_resp(result, model=model, query=json.dumps(query, ensure_ascii=False),
- delta_time=delta_time_log)
+ @app.route('/welcome')
+ def first_flask(): # 视图函数
+ return 'hello, modelcache!'
-def response_hitquery(cache_resp):
- return cache_resp['hitQuery']
-
-
-data2vec = Data2VecAudio()
-mysql_config = configparser.ConfigParser()
-mysql_config.read('modelcache/config/mysql_config.ini')
-milvus_config = configparser.ConfigParser()
-milvus_config.read('modelcache/config/milvus_config.ini')
-data_manager = get_data_manager(CacheBase("mysql", config=mysql_config),
- VectorBase("milvus", dimension=data2vec.dimension, milvus_config=milvus_config))
-
-
-cache.init(
- embedding_func=data2vec.to_embeddings,
- data_manager=data_manager,
- similarity_evaluation=SearchDistanceEvaluation(),
- query_pre_embedding_func=query_multi_splicing,
- insert_pre_embedding_func=insert_multi_splicing,
-)
-
-# cache.set_openai_key()
-global executor
-executor = ThreadPoolExecutor(max_workers=6)
-
-
-@app.route('/welcome')
-def first_flask(): # 视图函数
- return 'hello, modelcache!'
-
-
-@app.route('/modelcache', methods=['GET', 'POST'])
-def user_backend():
- try:
- if request.method == 'POST':
- request_data = request.json
- elif request.method == 'GET':
- request_data = request.args
- param_dict = json.loads(request_data)
- except Exception as e:
- result = {"errorCode": 101, "errorDesc": str(e), "cacheHit": False, "delta_time": 0, "hit_query": '',
- "answer": ''}
- cache.data_manager.save_query_resp(result, model='', query='', delta_time=0)
- return json.dumps(result)
-
- # param parsing
- try:
- request_type = param_dict.get("type")
-
- scope = param_dict.get("scope")
- if scope is not None:
- model = scope.get('model')
- model = model.replace('-', '_')
- model = model.replace('.', '_')
- query = param_dict.get("query")
- chat_info = param_dict.get("chat_info")
- if request_type is None or request_type not in ['query', 'insert', 'detox', 'remove']:
- result = {"errorCode": 102,
- "errorDesc": "type exception, should one of ['query', 'insert', 'detox', 'remove']",
- "cacheHit": False, "delta_time": 0, "hit_query": '', "answer": ''}
- cache.data_manager.save_query_resp(result, model=model, query='', delta_time=0)
- return json.dumps(result)
- except Exception as e:
- result = {"errorCode": 103, "errorDesc": str(e), "cacheHit": False, "delta_time": 0, "hit_query": '',
- "answer": ''}
- return json.dumps(result)
-
- # model filter
- filter_resp = model_blacklist_filter(model, request_type)
- if isinstance(filter_resp, dict):
- return json.dumps(filter_resp)
-
- if request_type == 'query':
+ @app.post('/modelcache')
+ def user_backend():
try:
- start_time = time.time()
- response = adapter.ChatCompletion.create_query(
- scope={"model": model},
- query=query
- )
- delta_time = '{}s'.format(round(time.time() - start_time, 2))
- if response is None:
- result = {"errorCode": 0, "errorDesc": '', "cacheHit": False, "delta_time": delta_time, "hit_query": '',
- "answer": ''}
- # elif response in ['adapt_query_exception']:
- elif isinstance(response, str):
- result = {"errorCode": 201, "errorDesc": response, "cacheHit": False, "delta_time": delta_time,
- "hit_query": '', "answer": ''}
- else:
- answer = response_text(response)
- hit_query = response_hitquery(response)
- result = {"errorCode": 0, "errorDesc": '', "cacheHit": True, "delta_time": delta_time,
- "hit_query": hit_query, "answer": answer}
- delta_time_log = round(time.time() - start_time, 2)
- future = executor.submit(save_query_info, result, model, query, delta_time_log)
- except Exception as e:
- result = {"errorCode": 202, "errorDesc": str(e), "cacheHit": False, "delta_time": 0,
- "hit_query": '', "answer": ''}
- logging.info('result: {}'.format(result))
+ param_dict = request.json
+ except Exception:
+ result = {"errorCode": 400, "errorDesc": "bad request", "cacheHit": False, "delta_time": 0, "hit_query": '',"answer": ''}
+ return jsonify(result), 400
- return json.dumps(result, ensure_ascii=False)
-
- if request_type == 'insert':
try:
- try:
- response = adapter.ChatCompletion.create_insert(
- model=model,
- chat_info=chat_info
- )
- except Exception as e:
- result = {"errorCode": 302, "errorDesc": str(e), "writeStatus": "exception"}
- return json.dumps(result, ensure_ascii=False)
-
- if response == 'success':
- result = {"errorCode": 0, "errorDesc": "", "writeStatus": "success"}
- else:
- result = {"errorCode": 301, "errorDesc": response, "writeStatus": "exception"}
- return json.dumps(result, ensure_ascii=False)
+ result = asyncio.run_coroutine_threadsafe(
+ cache.handle_request(param_dict), loop
+ ).result()
+ return jsonify(result), 200
except Exception as e:
- result = {"errorCode": 303, "errorDesc": str(e), "writeStatus": "exception"}
- return json.dumps(result, ensure_ascii=False)
-
- if request_type == 'remove':
- remove_type = param_dict.get("remove_type")
- id_list = param_dict.get("id_list", [])
-
- response = adapter.ChatCompletion.create_remove(
- model=model,
- remove_type=remove_type,
- id_list=id_list
- )
- if not isinstance(response, dict):
- result = {"errorCode": 401, "errorDesc": "", "response": response, "removeStatus": "exception"}
- return json.dumps(result)
+ result = {"errorCode": 500, "errorDesc": str(e), "cacheHit": False, "delta_time": 0, "hit_query": '',"answer": ''}
+ cache.save_query_resp(result, model='', query='', delta_time=0)
+ return jsonify(result), 500
- state = response.get('status')
- if state == 'success':
- result = {"errorCode": 0, "errorDesc": "", "response": response, "writeStatus": "success"}
- else:
- result = {"errorCode": 402, "errorDesc": "", "response": response, "writeStatus": "exception"}
- return json.dumps(result)
+ await asyncio.to_thread(app.run, host='0.0.0.0', port=5000)
if __name__ == '__main__':
- app.run(host='0.0.0.0', port=5000, debug=True)
+ asyncio.run(main())
diff --git a/flask4modelcache_demo.py b/flask4modelcache_demo.py
index 66ce8b0..4251f94 100644
--- a/flask4modelcache_demo.py
+++ b/flask4modelcache_demo.py
@@ -1,169 +1,48 @@
# -*- coding: utf-8 -*-
-import time
-from flask import Flask, request
-import logging
-import json
-from modelcache import cache
-from modelcache.adapter import adapter
-from modelcache.manager import CacheBase, VectorBase, get_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
-from concurrent.futures import ThreadPoolExecutor
-from modelcache.utils.model_filter import model_blacklist_filter
-from modelcache.embedding import Data2VecAudio
+import asyncio
-# 创建一个Flask实例
-app = Flask(__name__)
+from flask import Flask, request, jsonify
+from modelcache.cache import Cache
+from modelcache.embedding import EmbeddingModel
-def response_text(cache_resp):
- return cache_resp['data']
+async def main():
+ # 创建一个Flask实例
+ app = Flask(__name__)
-def save_query_info(result, model, query, delta_time_log):
- cache.data_manager.save_query_resp(result, model=model, query=json.dumps(query, ensure_ascii=False),
- delta_time=delta_time_log)
+ cache,loop = await Cache.init(
+ sql_storage="sqlite",
+ vector_storage="faiss",
+ embedding_model=EmbeddingModel.HUGGINGFACE_ALL_MPNET_BASE_V2,
+ embedding_workers_num=2
+ )
+ @app.route('/welcome')
+ def first_flask(): # 视图函数
+ return 'hello, modelcache!'
-def response_hitquery(cache_resp):
- return cache_resp['hitQuery']
-
-data2vec = Data2VecAudio()
-data_manager = get_data_manager(CacheBase("sqlite"), VectorBase("faiss", dimension=data2vec.dimension))
-
-
-cache.init(
- embedding_func=data2vec.to_embeddings,
- data_manager=data_manager,
- similarity_evaluation=SearchDistanceEvaluation(),
- query_pre_embedding_func=query_multi_splicing,
- insert_pre_embedding_func=insert_multi_splicing,
-)
-
-# cache.set_openai_key()
-global executor
-executor = ThreadPoolExecutor(max_workers=6)
-
-
-@app.route('/welcome')
-def first_flask(): # 视图函数
- return 'hello, modelcache!'
-
-
-@app.route('/modelcache', methods=['GET', 'POST'])
-def user_backend():
- try:
- if request.method == 'POST':
- request_data = request.json
- elif request.method == 'GET':
- request_data = request.args
- param_dict = json.loads(request_data)
- except Exception as e:
- result = {"errorCode": 101, "errorDesc": str(e), "cacheHit": False, "delta_time": 0, "hit_query": '',
- "answer": ''}
- cache.data_manager.save_query_resp(result, model='', query='', delta_time=0)
- return json.dumps(result)
-
- # param parsing
- try:
- request_type = param_dict.get("type")
- scope = param_dict.get("scope")
- if scope is not None:
- model = scope.get('model')
- model = model.replace('-', '_')
- model = model.replace('.', '_')
- query = param_dict.get("query")
- chat_info = param_dict.get("chat_info")
- if request_type is None or request_type not in ['query', 'insert', 'detox', 'remove']:
- result = {"errorCode": 102,
- "errorDesc": "type exception, should one of ['query', 'insert', 'detox', 'remove']",
- "cacheHit": False, "delta_time": 0, "hit_query": '', "answer": ''}
- cache.data_manager.save_query_resp(result, model=model, query='', delta_time=0)
- return json.dumps(result)
- except Exception as e:
- result = {"errorCode": 103, "errorDesc": str(e), "cacheHit": False, "delta_time": 0, "hit_query": '',
- "answer": ''}
- return json.dumps(result)
-
- # model filter
- filter_resp = model_blacklist_filter(model, request_type)
- if isinstance(filter_resp, dict):
- return json.dumps(filter_resp)
-
- if request_type == 'query':
+ @app.post('/modelcache')
+ def user_backend():
try:
- start_time = time.time()
- response = adapter.ChatCompletion.create_query(
- scope={"model": model},
- query=query
- )
- delta_time = '{}s'.format(round(time.time() - start_time, 2))
- if response is None:
- result = {"errorCode": 0, "errorDesc": '', "cacheHit": False, "delta_time": delta_time, "hit_query": '',
- "answer": ''}
- elif response in ['adapt_query_exception']:
- result = {"errorCode": 201, "errorDesc": response, "cacheHit": False, "delta_time": delta_time,
- "hit_query": '', "answer": ''}
- else:
- answer = response_text(response)
- hit_query = response_hitquery(response)
- result = {"errorCode": 0, "errorDesc": '', "cacheHit": True, "delta_time": delta_time,
- "hit_query": hit_query, "answer": answer}
- delta_time_log = round(time.time() - start_time, 2)
- future = executor.submit(save_query_info, result, model, query, delta_time_log)
- except Exception as e:
- result = {"errorCode": 202, "errorDesc": e, "cacheHit": False, "delta_time": 0,
- "hit_query": '', "answer": ''}
- logging.info('result: {}'.format(result))
-
- return json.dumps(result, ensure_ascii=False)
+ param_dict = request.json
+ except Exception:
+ result = {"errorCode": 400, "errorDesc": "bad request", "cacheHit": False, "delta_time": 0, "hit_query": '',"answer": ''}
+ return jsonify(result), 400
- if request_type == 'insert':
try:
- try:
- response = adapter.ChatCompletion.create_insert(
- model=model,
- chat_info=chat_info
- )
- except Exception as e:
- result = {"errorCode": 303, "errorDesc": e, "writeStatus": "exception"}
- return json.dumps(result, ensure_ascii=False)
-
- if response in ['adapt_insert_exception']:
- result = {"errorCode": 301, "errorDesc": response, "writeStatus": "exception"}
- elif response == 'success':
- result = {"errorCode": 0, "errorDesc": "", "writeStatus": "success"}
- else:
- result = {"errorCode": 302, "errorDesc": response,
- "writeStatus": "exception"}
- return json.dumps(result, ensure_ascii=False)
+ result = asyncio.run_coroutine_threadsafe(
+ cache.handle_request(param_dict), loop
+ ).result()
+ return jsonify(result), 200
except Exception as e:
- result = {"errorCode": 304, "errorDesc": e, "writeStatus": "exception"}
- return json.dumps(result, ensure_ascii=False)
-
- if request_type == 'remove':
- remove_type = param_dict.get("remove_type")
- id_list = param_dict.get("id_list", [])
-
- response = adapter.ChatCompletion.create_remove(
- model=model,
- remove_type=remove_type,
- id_list=id_list
- )
-
- if not isinstance(response, dict):
- result = {"errorCode": 401, "errorDesc": "", "response": response, "removeStatus": "exception"}
- return json.dumps(result)
+ result = {"errorCode": 500, "errorDesc": str(e), "cacheHit": False, "delta_time": 0, "hit_query": '',"answer": ''}
+ cache.save_query_resp(result, model='', query='', delta_time=0)
+ return jsonify(result), 500
- state = response.get('status')
- if state == 'success':
- result = {"errorCode": 0, "errorDesc": "", "response": response, "writeStatus": "success"}
- else:
- result = {"errorCode": 402, "errorDesc": "", "response": response, "writeStatus": "exception"}
- return json.dumps(result)
+ await asyncio.to_thread(app.run, host='0.0.0.0', port=5000)
if __name__ == '__main__':
- app.run(host='0.0.0.0', port=5000, debug=True)
+ asyncio.run(main())
diff --git a/flask4multicache.py b/flask4multicache.py
new file mode 100644
index 0000000..f59cf1f
--- /dev/null
+++ b/flask4multicache.py
@@ -0,0 +1,198 @@
+import time
+from flask import Flask, request
+import logging
+import json
+import configparser
+from concurrent.futures import ThreadPoolExecutor
+from modelcache_mm import cache
+from modelcache_mm.adapter import adapter
+from modelcache_mm.manager import CacheBase, VectorBase, get_data_manager
+from modelcache_mm.similarity_evaluation.distance import SearchDistanceEvaluation
+# from modelcache.processor.pre import query_multi_splicing
+# from modelcache.processor.pre import insert_multi_splicing
+# from modelcache.utils.model_filter import model_blacklist_filter
+# from modelcache.embedding import Data2VecAudio
+from modelcache_mm.processor.pre import mm_insert_dict
+from modelcache_mm.processor.pre import mm_query_dict
+from modelcache_mm.embedding import Clip2Vec
+
+# 创建一个Flask实例
+app = Flask(__name__)
+
+
+def response_text(cache_resp):
+ return cache_resp['data']
+
+
+def save_query_info(result, model, query, delta_time_log):
+ cache.data_manager.save_query_resp(result, model=model, query=json.dumps(query, ensure_ascii=False),
+ delta_time=delta_time_log)
+
+
+def response_hitquery(cache_resp):
+ return cache_resp['hitQuery']
+
+
+# data2vec = Data2VecAudio()
+mysql_config = configparser.ConfigParser()
+mysql_config.read('modelcache/config/mysql_config.ini')
+
+# milvus_config = configparser.ConfigParser()
+# milvus_config.read('modelcache/config/milvus_config.ini')
+
+redis_config = configparser.ConfigParser()
+redis_config.read('modelcache/config/redis_config.ini')
+
+
+image_dimension = 512
+text_dimension = 512
+clip2vec = Clip2Vec()
+data_manager = get_data_manager(CacheBase("mysql", config=mysql_config),
+ VectorBase("redis", mm_dimension=image_dimension+text_dimension,
+ i_dimension=image_dimension, t_dimension=text_dimension,
+ redis_config=redis_config))
+# data_manager = get_data_manager(CacheBase("mysql", config=mysql_config),
+# VectorBase("redis", dimension=data2vec.dimension, redis_config=redis_config))
+
+
+cache.init(
+ embedding_func=clip2vec.to_embeddings,
+ data_manager=data_manager,
+ similarity_evaluation=SearchDistanceEvaluation(),
+ insert_pre_embedding_func=mm_insert_dict,
+ query_pre_embedding_func=mm_query_dict,
+ )
+
+global executor
+executor = ThreadPoolExecutor(max_workers=6)
+
+
+@app.route('/welcome')
+def first_flask(): # 视图函数
+ return 'hello, llms_cache!'
+
+
+@app.route('/multicache', methods=['GET', 'POST'])
+def user_backend():
+ try:
+ if request.method == 'POST':
+ request_data = request.json
+ elif request.method == 'GET':
+ request_data = request.args
+ param_dict = json.loads(request_data)
+ except Exception as e:
+ result = {"errorCode": 301, "errorDesc": str(e), "cacheHit": False, "delta_time": 0, "hit_query": '',
+ "answer": ''}
+ cache.data_manager.save_query_resp(result, model='', query='', delta_time=0)
+ return json.dumps(result)
+
+ # param parsing
+ try:
+ request_type = param_dict.get("request_type")
+ scope = param_dict.get("scope")
+ if scope is not None:
+ model = scope.get('model')
+ model = model.replace('-', '_')
+ model = model.replace('.', '_')
+
+ if request_type in ['query', 'insert']:
+ if request_type == 'query':
+ query = param_dict.get("query")
+ elif request_type == 'insert':
+ chat_info = param_dict.get("chat_info")
+ query = chat_info[-1]['query']
+
+ if request_type is None or request_type not in ['query', 'remove', 'insert', 'register']:
+ result = {"errorCode": 102,
+ "errorDesc": "type exception, should one of ['query', 'insert', 'remove', 'register']",
+ "cacheHit": False, "delta_time": 0, "hit_query": '', "answer": ''}
+ cache.data_manager.save_query_resp(result, model=model, query='', delta_time=0)
+ return json.dumps(result)
+ except Exception as e:
+ result = {"errorCode": 103, "errorDesc": str(e), "cacheHit": False, "delta_time": 0, "hit_query": '',
+ "answer": ''}
+ return json.dumps(result)
+
+ if request_type == 'query':
+ try:
+ start_time = time.time()
+ response = adapter.ChatCompletion.create_query(
+ scope={"model": model},
+ query=query,
+ )
+ delta_time = '{}s'.format(round(time.time() - start_time, 2))
+ if response is None:
+ result = {"errorCode": 0, "errorDesc": '', "cacheHit": False, "delta_time": delta_time,
+ "hit_query": '', "answer": ''}
+ elif isinstance(response, dict):
+ answer = response_text(response)
+ hit_query = response_hitquery(response)
+ result = {"errorCode": 0, "errorDesc": '', "cacheHit": True, "delta_time": delta_time,
+ "hit_query": hit_query, "answer": answer}
+ else:
+ result = {"errorCode": 201, "errorDesc": response, "cacheHit": False, "delta_time": delta_time,
+ "hit_query": '', "answer": ''}
+ delta_time_log = round(time.time() - start_time, 3)
+
+ future = executor.submit(save_query_info, result, model, query, delta_time_log)
+ except Exception as e:
+ raise e
+ return json.dumps(result, ensure_ascii=False)
+
+ if request_type == 'insert':
+ try:
+ start_time = time.time()
+ try:
+ response = adapter.ChatCompletion.create_insert(
+ model=model,
+ chat_info=chat_info,
+ )
+ except Exception as e:
+ raise e
+
+ if response == 'success':
+ result = {"errorCode": 0, "errorDesc": "", "writeStatus": "success"}
+ else:
+ result = {"errorCode": 301, "errorDesc": response, "writeStatus": "exception"}
+ insert_time = round(time.time() - start_time, 2)
+ return json.dumps(result, ensure_ascii=False)
+ except Exception as e:
+ raise e
+
+ if request_type == 'remove':
+ remove_type = param_dict.get("remove_type")
+ id_list = param_dict.get("id_list", [])
+
+ response = adapter.ChatCompletion.create_remove(
+ model=model,
+ remove_type=remove_type,
+ id_list=id_list
+ )
+
+ if not isinstance(response, dict):
+ result = {"errorCode": 401, "errorDesc": "", "response": response, "removeStatus": "exception"}
+ return json.dumps(result)
+
+ state = response.get('status')
+ # if response == 'success':
+ if state == 'success':
+ result = {"errorCode": 0, "errorDesc": "", "response": response, "writeStatus": "success"}
+ else:
+ result = {"errorCode": 402, "errorDesc": "", "response": response, "writeStatus": "exception"}
+ return json.dumps(result)
+
+ if request_type == 'register':
+ type = param_dict.get("type")
+ response = adapter.ChatCompletion.create_register(
+ model=model,
+ type=type
+ )
+ if response in ['create_success', 'already_exists']:
+ result = {"errorCode": 0, "errorDesc": "", "response": response, "writeStatus": "success"}
+ else:
+ result = {"errorCode": 502, "errorDesc": "", "response": response, "writeStatus": "exception"}
+ return json.dumps(result)
+
+
+if __name__ == '__main__':
+ app.run(host='0.0.0.0', port=5000, debug=True)
diff --git a/flask4multicache_demo.py b/flask4multicache_demo.py
new file mode 100644
index 0000000..3f205fd
--- /dev/null
+++ b/flask4multicache_demo.py
@@ -0,0 +1,190 @@
+# -*- coding: utf-8 -*-
+import time
+from flask import Flask, request
+import logging
+import json
+from concurrent.futures import ThreadPoolExecutor
+from modelcache_mm import cache
+from modelcache_mm.adapter import adapter
+from modelcache_mm.manager import CacheBase, VectorBase, get_data_manager
+from modelcache_mm.similarity_evaluation.distance import SearchDistanceEvaluation
+# from modelcache.processor.pre import query_multi_splicing
+# from modelcache.processor.pre import insert_multi_splicing
+# from modelcache.utils.model_filter import model_blacklist_filter
+# from modelcache.embedding import Data2VecAudio
+from modelcache_mm.processor.pre import mm_insert_dict
+from modelcache_mm.processor.pre import mm_query_dict
+from modelcache_mm.embedding import Clip2Vec
+
+# 创建一个Flask实例
+app = Flask(__name__)
+
+
+def response_text(cache_resp):
+ return cache_resp['data']
+
+
+def save_query_info(result, model, query, delta_time_log):
+ cache.data_manager.save_query_resp(result, model=model, query=json.dumps(query, ensure_ascii=False),
+ delta_time=delta_time_log)
+
+
+def response_hitquery(cache_resp):
+ return cache_resp['hitQuery']
+
+
+# data2vec = Data2VecAudio()
+# data_manager = get_data_manager(CacheBase("sqlite"), VectorBase("faiss", dimension=data2vec.dimension))
+
+image_dimension = 512
+text_dimension = 512
+clip2vec = Clip2Vec()
+data_manager = get_data_manager(CacheBase("sqlite"), VectorBase("faiss",
+ mm_dimension=image_dimension+text_dimension,
+ i_dimension=image_dimension,
+ t_dimension=text_dimension))
+
+
+cache.init(
+ embedding_func=clip2vec.to_embeddings,
+ data_manager=data_manager,
+ similarity_evaluation=SearchDistanceEvaluation(),
+ insert_pre_embedding_func=mm_insert_dict,
+ query_pre_embedding_func=mm_query_dict,
+ )
+
+# cache.set_openai_key()
+global executor
+executor = ThreadPoolExecutor(max_workers=6)
+
+
+@app.route('/welcome')
+def first_flask(): # 视图函数
+ return 'hello, llms_cache!'
+
+
+@app.route('/llms_cache', methods=['GET', 'POST'])
+def user_backend():
+ try:
+ if request.method == 'POST':
+ request_data = request.json
+ elif request.method == 'GET':
+ request_data = request.args
+ param_dict = json.loads(request_data)
+ except Exception as e:
+ result = {"errorCode": 301, "errorDesc": str(e), "cacheHit": False, "delta_time": 0, "hit_query": '',
+ "answer": ''}
+ cache.data_manager.save_query_resp(result, model='', query='', delta_time=0)
+ return json.dumps(result)
+
+ # param parsing
+ try:
+ request_type = param_dict.get("request_type")
+ scope = param_dict.get("scope")
+ if scope is not None:
+ model = scope.get('model')
+ model = model.replace('-', '_')
+ model = model.replace('.', '_')
+
+ if request_type in ['query', 'insert']:
+ if request_type == 'query':
+ query = param_dict.get("query")
+ elif request_type == 'insert':
+ chat_info = param_dict.get("chat_info")
+ query = chat_info[-1]['query']
+
+ if request_type is None or request_type not in ['query', 'remove', 'insert', 'register']:
+ result = {"errorCode": 102,
+ "errorDesc": "type exception, should one of ['query', 'insert', 'remove', 'register']",
+ "cacheHit": False, "delta_time": 0, "hit_query": '', "answer": ''}
+ cache.data_manager.save_query_resp(result, model=model, query='', delta_time=0)
+ return json.dumps(result)
+ except Exception as e:
+ result = {"errorCode": 103, "errorDesc": str(e), "cacheHit": False, "delta_time": 0, "hit_query": '',
+ "answer": ''}
+ return json.dumps(result)
+
+ if request_type == 'query':
+ try:
+ start_time = time.time()
+ response = adapter.ChatCompletion.create_query(
+ scope={"model": model},
+ query=query,
+ )
+ delta_time = '{}s'.format(round(time.time() - start_time, 2))
+ if response is None:
+ result = {"errorCode": 0, "errorDesc": '', "cacheHit": False, "delta_time": delta_time,
+ "hit_query": '', "answer": ''}
+ elif isinstance(response, dict):
+ answer = response_text(response)
+ hit_query = response_hitquery(response)
+ result = {"errorCode": 0, "errorDesc": '', "cacheHit": True, "delta_time": delta_time,
+ "hit_query": hit_query, "answer": answer}
+ else:
+ result = {"errorCode": 201, "errorDesc": response, "cacheHit": False, "delta_time": delta_time,
+ "hit_query": '', "answer": ''}
+ delta_time_log = round(time.time() - start_time, 3)
+
+ future = executor.submit(save_query_info, result, model, query, delta_time_log)
+ except Exception as e:
+ raise e
+ return json.dumps(result, ensure_ascii=False)
+
+ if request_type == 'insert':
+ try:
+ start_time = time.time()
+ try:
+ response = adapter.ChatCompletion.create_insert(
+ model=model,
+ chat_info=chat_info,
+ )
+ except Exception as e:
+ raise e
+
+ if response == 'success':
+ result = {"errorCode": 0, "errorDesc": "", "writeStatus": "success"}
+ else:
+ result = {"errorCode": 301, "errorDesc": response, "writeStatus": "exception"}
+ insert_time = round(time.time() - start_time, 2)
+ return json.dumps(result, ensure_ascii=False)
+ except Exception as e:
+ raise e
+
+ if request_type == 'remove':
+ remove_type = param_dict.get("remove_type")
+ id_list = param_dict.get("id_list", [])
+
+ response = adapter.ChatCompletion.create_remove(
+ model=model,
+ remove_type=remove_type,
+ id_list=id_list
+ )
+
+ if not isinstance(response, dict):
+ result = {"errorCode": 401, "errorDesc": "", "response": response, "removeStatus": "exception"}
+ return json.dumps(result)
+
+ state = response.get('status')
+ # if response == 'success':
+ if state == 'success':
+ result = {"errorCode": 0, "errorDesc": "", "response": response, "writeStatus": "success"}
+ else:
+ result = {"errorCode": 402, "errorDesc": "", "response": response, "writeStatus": "exception"}
+ return json.dumps(result)
+
+ if request_type == 'register':
+ type = param_dict.get("type")
+ response = adapter.ChatCompletion.create_register(
+ model=model,
+ type=type
+ )
+ if response in ['create_success', 'already_exists']:
+ result = {"errorCode": 0, "errorDesc": "", "response": response, "writeStatus": "success"}
+ else:
+ result = {"errorCode": 502, "errorDesc": "", "response": response, "writeStatus": "exception"}
+ return json.dumps(result)
+
+
+if __name__ == '__main__':
+ # app.run(host='0.0.0.0', port=5000, debug=True)
+ app.run(host='0.0.0.0', port=5000)
diff --git a/model/.gitignore b/model/.gitignore
new file mode 100644
index 0000000..08ad955
--- /dev/null
+++ b/model/.gitignore
@@ -0,0 +1,2 @@
+*.tflite
+text2vec-base-chinese/*
\ No newline at end of file
diff --git a/model/clip_zh/__init__.py b/model/clip_zh/__init__.py
new file mode 100644
index 0000000..40a96af
--- /dev/null
+++ b/model/clip_zh/__init__.py
@@ -0,0 +1 @@
+# -*- coding: utf-8 -*-
diff --git a/modelcache/__init__.py b/modelcache/__init__.py
index 9a82563..8bcb18a 100644
--- a/modelcache/__init__.py
+++ b/modelcache/__init__.py
@@ -1,5 +1,7 @@
# -*- coding: utf-8 -*-
-from modelcache.core import Cache
-from modelcache.core import cache
-from modelcache.config import Config
-import modelcache.adapter
+
+
+
+
+
+
diff --git a/modelcache/adapter/adapter.py b/modelcache/adapter/adapter.py
index f5e38eb..2425539 100644
--- a/modelcache/adapter/adapter.py
+++ b/modelcache/adapter/adapter.py
@@ -1,47 +1,59 @@
# -*- coding: utf-8 -*-
import logging
-
-import openai
from modelcache.adapter.adapter_query import adapt_query
from modelcache.adapter.adapter_insert import adapt_insert
from modelcache.adapter.adapter_remove import adapt_remove
+from modelcache.adapter.adapter_register import adapt_register
-class ChatCompletion(openai.ChatCompletion):
+class ChatCompletion(object):
"""Openai ChatCompletion Wrapper"""
@classmethod
- def create_query(cls, *args, **kwargs):
+ async def create_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(
+ return await adapt_query(
cache_data_convert,
*args,
**kwargs
)
except Exception as e:
+ print(e)
+ return str(e)
+
+ @classmethod
+ async def create_insert(cls, *args, **kwargs):
+ try:
+ return await adapt_insert(
+ *args,
+ **kwargs
+ )
+ except Exception as e:
+ print(e)
return str(e)
@classmethod
- def create_insert(cls, *args, **kwargs):
+ async def create_remove(cls, *args, **kwargs):
try:
- return adapt_insert(
+ return await adapt_remove(
*args,
**kwargs
)
except Exception as e:
+ print(e)
return str(e)
@classmethod
- def create_remove(cls, *args, **kwargs):
+ async def create_register(cls, *args, **kwargs):
try:
- return adapt_remove(
+ return await adapt_register(
*args,
**kwargs
)
except Exception as e:
- logging.info('adapt_remove_e: {}'.format(e))
+ print(e)
return str(e)
diff --git a/modelcache/adapter/adapter_insert.py b/modelcache/adapter/adapter_insert.py
index 74aa619..a2f224a 100644
--- a/modelcache/adapter/adapter_insert.py
+++ b/modelcache/adapter/adapter_insert.py
@@ -1,39 +1,56 @@
# -*- coding: utf-8 -*-
-from modelcache import cache
+import asyncio
+
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)
+async def adapt_insert(*args, **kwargs):
+ chat_cache = kwargs.pop("cache_obj")
model = kwargs.pop("model", None)
require_object_store = kwargs.pop("require_object_store", False)
+
+ # Validate object store availability if required
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(
+ # Initialize collections for parallel processing
+ pre_embedding_data_list = [] # Preprocessed data ready for embedding
+ embedding_futures_list = [] # Async embedding generation tasks
+ llm_data_list = [] # Extracted LLM response data
+
+ # Process each chat entry and prepare for parallel embedding generation
+ for row in chat_info:
+ # Preprocess chat data using configured preprocessing function
+ pre_embedding_data = chat_cache.insert_pre_embedding_func(
+ row,
+ extra_param=context.get("pre_embedding_func", None),
+ prompts=chat_cache.prompts,
+ )
+ pre_embedding_data_list.append(pre_embedding_data)
+ llm_data_list.append(row['answer']) # Extract answer text for storage
+
+ # Create async embedding generation task with performance monitoring
+ embedding_future = time_cal(
chat_cache.embedding_func,
func_name="embedding",
report_func=chat_cache.report.embedding,
+ cache_obj=chat_cache
)(pre_embedding_data)
+ embedding_futures_list.append(embedding_future)
+
+ # Wait for all embedding generation tasks to complete in parallel
+ embedding_data_list = await asyncio.gather(*embedding_futures_list)
- chat_cache.data_manager.save(
- pre_embedding_data,
- llm_data,
- embedding_data,
+ # Save all processed data to the data manager asynchronously
+ await asyncio.to_thread(
+ chat_cache.data_manager.save,
+ pre_embedding_data_list,
+ llm_data_list,
+ embedding_data_list,
model=model,
extra_param=context.get("save_func", None)
)
diff --git a/modelcache/adapter/adapter_query.py b/modelcache/adapter/adapter_query.py
index 934c644..d98e503 100644
--- a/modelcache/adapter/adapter_query.py
+++ b/modelcache/adapter/adapter_query.py
@@ -1,55 +1,72 @@
# -*- coding: utf-8 -*-
+import asyncio
import logging
-import time
-from modelcache import cache
-from modelcache.utils.error import NotInitError
+from modelcache.embedding import MetricType
from modelcache.utils.time import time_cal
-from modelcache.processor.pre import multi_analysis
+from FlagEmbedding import FlagReranker
+USE_RERANKER = False # 如果为 True 则启用 reranker，否则使用原有逻辑
-def adapt_query(cache_data_convert, *args, **kwargs):
- chat_cache = kwargs.pop("cache_obj", cache)
- scope = kwargs.pop("scope", None)
+async def adapt_query(cache_data_convert, *args, **kwargs):
+ # Extract query parameters
+ chat_cache = kwargs.pop("cache_obj")
+ scope = kwargs.pop("scope")
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)
+
+ # Preprocess query for embedding generation
pre_embedding_data = chat_cache.query_pre_embedding_func(
kwargs,
extra_param=context.get("pre_embedding_func", None),
- prompts=chat_cache.config.prompts,
+ prompts=chat_cache.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
+ # Generate embedding with performance monitoring
+ embedding_data = await time_cal(
+ chat_cache.embedding_func,
+ func_name="embedding",
+ report_func=chat_cache.report.embedding,
+ cache_obj=chat_cache
+ )(pre_embedding_data)
+
+ search_time_cal = time_cal(
+ chat_cache.data_manager.search,
+ func_name="vector_search",
+ report_func=chat_cache.report.search,
+ cache_obj=chat_cache
+ )
+ cache_data_list = await asyncio.to_thread(
+ search_time_cal,
+ embedding_data,
+ extra_param=context.get("search_func", None),
+ top_k=kwargs.pop("top_k", -1),
+ model=model
+ )
+
+ # Initialize result containers
+ cache_answers = []
+ cache_questions = []
+ cache_ids = []
+ cosine_similarity = None
+
+ # Similarity evaluation based on metric type
+ if chat_cache.similarity_metric_type == MetricType.COSINE:
+ cosine_similarity = cache_data_list[0][0]
+ # This code uses the built-in cosine similarity evaluation in milvus
+ if cosine_similarity < chat_cache.similarity_threshold:
+ return None # No suitable match found
+
+ elif chat_cache.similarity_metric_type == MetricType.L2:
+ # this is the code that uses L2 for similarity evaluation
+ similarity_threshold = chat_cache.similarity_threshold
+ similarity_threshold_long = chat_cache.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
+
+ # Clamp thresholds to valid range
rank_threshold = (
max_rank
if rank_threshold > max_rank
@@ -65,6 +82,7 @@ def adapt_query(cache_data_convert, *args, **kwargs):
else rank_threshold_long
)
+ # Evaluate similarity score
if cache_data_list is None or len(cache_data_list) == 0:
rank_pre = -1.0
else:
@@ -75,17 +93,26 @@ def adapt_query(cache_data_convert, *args, **kwargs):
extra_param=context.get("evaluation_func", None),
)
if rank_pre < rank_threshold:
- return
+ return None # Similarity too low
+ else:
+ raise ValueError(
+ f"Unsupported similarity metric type: {chat_cache.similarity_metric_type}"
+ )
+ # Process search results with optional reranking
+ if USE_RERANKER:
+ reranker = FlagReranker('BAAI/bge-reranker-v2-m3', use_fp16=False)
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)
+ ret = await asyncio.to_thread(
+ chat_cache.data_manager.get_scalar_data,
+ cache_data, extra_param=context.get("get_scalar_data", None), model=model
)
if ret is None:
continue
+ rank = reranker.compute_score([pre_embedding_data, ret[0]], normalize=True)[0]
+
if "deps" in context and hasattr(ret.question, "deps"):
eval_query_data = {
"question": context["deps"][0]["data"],
@@ -102,47 +129,109 @@ def adapt_query(cache_data_convert, *args, **kwargs):
"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_answers.append((rank, ret[0]))
+ cache_questions.append((rank, ret[1]))
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_answers.append((rank, ret[0]))
+ cache_questions.append((rank, ret[1]))
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]
+ else:
+ # Original logic without reranking
+ for cache_data in cache_data_list:
+ primary_id = cache_data[1]
+ # Retrieve full cache entry data
+ ret = await asyncio.to_thread(
+ chat_cache.data_manager.get_scalar_data,
+ cache_data, extra_param=context.get("get_scalar_data", None), model=model
)
- # 更新命中次数
- 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)
+ if ret is None:
+ continue
+
+ if chat_cache.similarity_metric_type == MetricType.COSINE:
+ assert cosine_similarity is not None, "cosine_similarity should not be None"
+ cache_answers.append((cosine_similarity, ret[0]))
+ cache_questions.append((cosine_similarity, ret[1]))
+ cache_ids.append((cosine_similarity, primary_id))
+
+ elif chat_cache.similarity_metric_type == MetricType.L2:
+ 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
+ }
+
+ # Evaluate similarity for this specific result
+ 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[0]))
+ cache_questions.append((rank, ret[1]))
+ cache_ids.append((rank, primary_id))
+ else:
+ if rank_threshold_long <= rank:
+ cache_answers.append((rank, ret[0]))
+ cache_questions.append((rank, ret[1]))
+ cache_ids.append((rank, primary_id))
+ else:
+ raise ValueError(
+ f"Unsupported similarity metric type: {chat_cache.similarity_metric_type}"
+ )
+
+ # Sort results by similarity score (highest first)
+ 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]
+ )
+
+ # Update hit count for analytics (async to avoid blocking)
+ try:
+ asyncio.create_task(asyncio.to_thread(chat_cache.data_manager.update_hit_count,return_id))
+ except Exception:
+ logging.info('update_hit_count except, please check!')
+
+ # Record cache hit for reporting
+ chat_cache.report.hint_cache()
+ return cache_data_convert(return_message, return_query)
+ return None
\ No newline at end of file
diff --git a/modelcache/adapter/adapter_register.py b/modelcache/adapter/adapter_register.py
new file mode 100644
index 0000000..c4d0317
--- /dev/null
+++ b/modelcache/adapter/adapter_register.py
@@ -0,0 +1,16 @@
+# -*- coding: utf-8 -*-
+import asyncio
+
+
+async def adapt_register(*args, **kwargs):
+ chat_cache = kwargs.pop("cache_obj")
+ model = kwargs.pop("model", None)
+ if model is None or len(model) == 0:
+ return ValueError('')
+
+ register_resp = await asyncio.to_thread(
+ chat_cache.data_manager.create_index,
+ model
+ )
+
+ return register_resp
diff --git a/modelcache/adapter/adapter_remove.py b/modelcache/adapter/adapter_remove.py
index 25f1ba3..aea0284 100644
--- a/modelcache/adapter/adapter_remove.py
+++ b/modelcache/adapter/adapter_remove.py
@@ -1,24 +1,29 @@
# -*- coding: utf-8 -*-
-from modelcache import cache
-from modelcache.utils.error import NotInitError, RemoveError
+import asyncio
+from modelcache.utils.error import RemoveError
-def adapt_remove(*args, **kwargs):
- chat_cache = kwargs.pop("cache_obj", cache)
+
+async def adapt_remove(*args, **kwargs):
+ chat_cache = kwargs.pop("cache_obj")
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)
+ resp = await asyncio.to_thread(
+ chat_cache.data_manager.delete,
+ id_list, model=model
+ )
elif remove_type == 'truncate_by_model':
- resp = chat_cache.data_manager.truncate(model)
+ resp = await asyncio.to_thread(
+ chat_cache.data_manager.truncate,
+ model
+ )
else:
# resp = "remove_type_error"
raise RemoveError()
diff --git a/modelcache/cache.py b/modelcache/cache.py
new file mode 100644
index 0000000..9149a3c
--- /dev/null
+++ b/modelcache/cache.py
@@ -0,0 +1,399 @@
+# -*- coding: utf-8 -*-
+import asyncio
+import atexit
+import json
+import logging
+import time
+from asyncio import AbstractEventLoop
+from typing import Callable, Optional, List, Any, Coroutine
+from modelcache.adapter import adapter
+from modelcache.embedding.embedding_dispatcher import EmbeddingDispatcher
+from modelcache.utils.model_filter import model_blacklist_filter
+from concurrent.futures import ThreadPoolExecutor, Future
+import configparser
+from modelcache.embedding.base import BaseEmbedding, EmbeddingModel, MetricType
+from modelcache.manager.scalar_data.sql_storage import SQLStorage
+from modelcache.manager.vector_data.base import VectorStorage
+from modelcache.processor.post import first
+from modelcache.processor.pre import query_with_role, query_multi_splicing, insert_multi_splicing
+from modelcache.similarity_evaluation.base import SimilarityEvaluation
+from modelcache.report import Report
+from modelcache.similarity_evaluation.distance import SearchDistanceEvaluation
+from modelcache.utils.error import CacheError
+from modelcache.utils.log import modelcache_log
+from modelcache.manager.data_manager import DataManager
+
+ #=====================================================================#
+ #==================== Cache class definition =========================#
+ #=====================================================================#
+
+
+# noinspection PyMethodMayBeStatic
+class Cache:
+ def __init__(
+ self,
+ embedding_model: EmbeddingModel,
+ similarity_metric_type: MetricType,
+ data_manager: DataManager,
+ query_pre_embedding_func: Callable,
+ insert_pre_embedding_func: Callable,
+ embedding_func: Callable[[str], Future],
+ report: Report, # TODO: figure out why this is needed
+ similarity_evaluation: Optional[SimilarityEvaluation],
+ post_process_messages_func: Callable,
+
+ similarity_threshold: float = 0.95,
+ similarity_threshold_long: float = 0.95,
+ prompts: Optional[List[str]] = None,
+ log_time_func: Callable[[str, float], None] = None,
+ ):
+ if similarity_threshold < 0 or similarity_threshold > 1:
+ raise CacheError(
+ "Invalid the similarity threshold param, reasonable range: 0-1"
+ )
+ self.data_manager: DataManager = data_manager
+ self.embedding_model: EmbeddingModel = embedding_model
+ self.similarity_metric_type: MetricType = similarity_metric_type
+ self.report: Report = report
+ self.query_pre_embedding_func: Callable = query_pre_embedding_func
+ self.insert_pre_embedding_func: Callable = insert_pre_embedding_func
+ self.embedding_func: Callable = embedding_func
+ self.similarity_evaluation: Optional[SimilarityEvaluation] = similarity_evaluation
+ self.post_process_messages_func: Callable = post_process_messages_func
+ self.similarity_threshold = similarity_threshold
+ self.similarity_threshold_long = similarity_threshold_long
+ self.prompts = prompts
+ self.log_time_func: Callable[[str, float], None] = log_time_func
+
+ @atexit.register
+ def close():
+ try:
+ self.data_manager.close()
+ except Exception as e:
+ modelcache_log.error(e)
+
+ def save_query_resp(self, query_resp_dict, **kwargs):
+ """
+ Save query response asynchronously to avoid blocking main thread.
+ Used for logging and analytics purposes.
+ """
+ # Execute save operation in a separate thread to maintain async performance
+ asyncio.create_task(asyncio.to_thread(
+ self.data_manager.save_query_resp,
+ query_resp_dict, **kwargs
+ ))
+
+ def save_query_info(self, result, model, query, delta_time_log):
+ """
+ Save query information with execution timing for performance analysis.
+ Serializes query data to JSON for storage.
+ """
+ # Convert query to JSON and save asynchronously
+ asyncio.create_task(asyncio.to_thread(
+ self.data_manager.save_query_resp,
+ result, model=model, query=json.dumps(query, ensure_ascii=False), delta_time=delta_time_log
+ ))
+
+ async def handle_request(self, param_dict: dict):
+ """
+ Main entry point for processing cache requests.
+
+ Routes requests to appropriate handlers based on request type.
+ Supports: query, insert, remove, register operations.
+
+ Args:
+ param_dict: Request parameters containing type, scope, query, etc.
+
+ Returns:
+ dict: Response dictionary with errorCode, result data, and metadata
+ """
+ # Parse and validate request parameters
+ try:
+ request_type = param_dict.get("type")
+
+ # Extract model information from scope
+ scope = param_dict.get("scope")
+ model = None
+ if scope is not None:
+ model = scope.get('model')
+ # Normalize model name for consistent storage (replace special chars)
+ model = model.replace('-', '_')
+ model = model.replace('.', '_')
+ query = param_dict.get("query")
+ chat_info = param_dict.get("chat_info")
+
+ # Validate request type against supported operations
+ if request_type is None or request_type not in ['query', 'insert', 'remove', 'register']:
+ result = {"errorCode": 102,
+ "errorDesc": "type exception, should one of ['query', 'insert', 'remove', 'register']",
+ "cacheHit": False, "delta_time": 0, "hit_query": '', "answer": ''}
+ self.save_query_resp(result, model=model, query='', delta_time=0)
+ return result
+ except Exception as e:
+ # Return error response for parameter parsing failures
+ return {"errorCode": 103, "errorDesc": str(e), "cacheHit": False, "delta_time": 0, "hit_query": '',
+ "answer": ''}
+
+ # Apply model-based filtering (blacklist check)
+ filter_resp = model_blacklist_filter(model, request_type)
+ if isinstance(filter_resp, dict):
+ return filter_resp
+
+ # Route to appropriate handler based on request type
+ if request_type == 'query':
+ return await self.handle_query(model, query)
+ elif request_type == 'insert':
+ return await self.handle_insert(chat_info, model)
+ elif request_type == 'remove':
+ return await self.handle_remove(model, param_dict)
+ elif request_type == 'register':
+ return await self.handle_register(model)
+ else:
+ return {"errorCode": 400, "errorDesc": "bad request"}
+
+ async def handle_register(self, model):
+ response = await adapter.ChatCompletion.create_register(
+ model=model,
+ cache_obj=self
+ )
+ # Process registration response and return standardized result
+ if response in ['create_success', 'already_exists']:
+ result = {"errorCode": 0, "errorDesc": "", "response": response, "writeStatus": "success"}
+ else:
+ result = {"errorCode": 502, "errorDesc": "", "response": response, "writeStatus": "exception"}
+ return result
+
+ async def handle_remove(self, model, param_dict):
+ remove_type = param_dict.get("remove_type")
+ id_list = param_dict.get("id_list", [])
+
+ # Execute removal operation through adapter
+ response = await adapter.ChatCompletion.create_remove(
+ model=model,
+ remove_type=remove_type,
+ id_list=id_list,
+ cache_obj=self
+ )
+
+ # Process removal response and standardize result format
+ if not isinstance(response, dict):
+ return {"errorCode": 401, "errorDesc": "", "response": response, "removeStatus": "exception"}
+ state = response.get('status')
+ if state == 'success':
+ result = {"errorCode": 0, "errorDesc": "", "response": response, "writeStatus": "success"}
+ else:
+ result = {"errorCode": 402, "errorDesc": "", "response": response, "writeStatus": "exception"}
+ return result
+
+ async def handle_insert(self, chat_info, model):
+ try:
+ try:
+ # Execute insertion through adapter with error handling
+ response = await adapter.ChatCompletion.create_insert(
+ model=model,
+ chat_info=chat_info,
+ cache_obj=self
+ )
+ except Exception as e:
+ return {"errorCode": 302, "errorDesc": str(e), "writeStatus": "exception"}
+
+ # Process insertion response
+ if response == 'success':
+ result = {"errorCode": 0, "errorDesc": "", "writeStatus": "success"}
+ else:
+ result = {"errorCode": 301, "errorDesc": response, "writeStatus": "exception"}
+ return result
+ except Exception as e:
+ return {"errorCode": 303, "errorDesc": str(e), "writeStatus": "exception"}
+
+ async def handle_query(self, model, query):
+ try:
+ start_time = time.time() # Start performance timer
+
+ # Execute query through adapter system
+ response = await adapter.ChatCompletion.create_query(
+ scope={"model": model},
+ query=query,
+ cache_obj=self
+ )
+
+ # Calculate query execution time
+ delta_time = '{}s'.format(round(time.time() - start_time, 2))
+
+ # Process different response types
+ if response is None:
+ # No cache hit found
+ result = {"errorCode": 0, "errorDesc": '', "cacheHit": False, "delta_time": delta_time, "hit_query": '',
+ "answer": ''}
+ elif isinstance(response, str):
+ # Error occurred during query processing
+ result = {"errorCode": 201, "errorDesc": response, "cacheHit": False, "delta_time": delta_time,
+ "hit_query": '', "answer": ''}
+ else:
+ # Cache hit found - extract response data
+ answer = response['data']
+ hit_query = response['hitQuery']
+ result = {"errorCode": 0, "errorDesc": '', "cacheHit": True, "delta_time": delta_time,
+ "hit_query": hit_query, "answer": answer}
+
+ # Log query performance data asynchronously
+ delta_time_log = round(time.time() - start_time, 2)
+ self.save_query_info(result, model, query, delta_time_log)
+ except Exception as e:
+ result = {"errorCode": 202, "errorDesc": str(e), "cacheHit": False, "delta_time": 0,
+ "hit_query": '', "answer": ''}
+ logging.info('result: {}'.format(result))
+ return result
+
+ def flush(self):
+ """Flush all cached data to persistent storage backends."""
+ self.data_manager.flush()
+
+ @staticmethod
+ async def init(
+ sql_storage: str,
+ vector_storage: str,
+ embedding_model: EmbeddingModel,
+ embedding_workers_num: int
+ ) -> tuple['Cache' , AbstractEventLoop]:
+ """
+ Initialize a complete Cache system with all required components.
+
+ Args:
+ sql_storage: SQL backend type ("mysql", "sqlite", "elasticsearch")
+ vector_storage: Vector backend type ("milvus", "faiss", "chromadb", "redis")
+ embedding_model: Embedding model enum value
+ embedding_workers_num: Number of parallel embedding worker processes
+
+ Returns:
+ tuple: (Cache instance, event loop) ready for async operations
+
+ Raises:
+ CacheError: If unsupported storage backends or model configuration issues
+
+ Example:
+ cache, loop = await Cache.init(
+ sql_storage="mysql",
+ vector_storage="milvus",
+ embedding_model=EmbeddingModel.HUGGINGFACE_ALL_MPNET_BASE_V2,
+ embedding_workers_num=4
+ )
+ """
+ #================= Database configuration loading ===================#
+
+ sql_config = configparser.ConfigParser()
+ vector_config = configparser.ConfigParser()
+
+ # Load SQL storage configuration from INI files
+ if sql_storage == "mysql":
+ sql_config.read('modelcache/config/mysql_config.ini')
+ elif sql_storage == "elasticsearch":
+ sql_config.read('modelcache/config/elasticsearch_config.ini')
+ elif sql_storage == "sqlite":
+ sql_config.read('modelcache/config/sqlite_config.ini')
+ else:
+ modelcache_log.error(f"Unsupported cache storage: {sql_storage}.")
+ raise CacheError(f"Unsupported cache storage: {sql_storage}.")
+
+ # Load vector storage configuration from INI files
+ if vector_storage == "milvus" :
+ vector_config.read('modelcache/config/milvus_config.ini')
+ elif vector_storage == "chromadb" :
+ vector_config.read('modelcache/config/chromadb_config.ini')
+ elif vector_storage == "redis" :
+ vector_config.read('modelcache/config/redis_config.ini')
+ elif vector_storage == "faiss" :
+ vector_config = None # FAISS does not require external configuration
+ else:
+ modelcache_log.error(f"Unsupported vector storage: {vector_storage}.")
+ raise CacheError(f"Unsupported vector storage: {vector_storage}.")
+
+
+ #=============== embedding-model-specific configuration =====================#
+
+ event_loop = asyncio.get_running_loop() # Get current async event loop
+ model_path = embedding_model.value['model_path']
+ dimension = embedding_model.value['dimension']
+
+ # Validate that embedding model has required configuration
+ if model_path is None or dimension is None:
+ modelcache_log.error(f"Please set the model_path and dimension for {embedding_model} in modelcache/embedding/base.py.")
+ raise CacheError(f"Please set the model_path and dimension for {embedding_model} in modelcache/embedding/base.py.")
+
+ # Initialize parallel embedding generation system
+ embedding_dispatcher = EmbeddingDispatcher(embedding_model, model_path, event_loop, embedding_workers_num)
+
+ #=== These will be used to initialize the cache ===#
+ query_pre_embedding_func: Callable = None
+ insert_pre_embedding_func: Callable = None
+ post_process_messages_func: Callable = None
+ similarity_evaluation: Optional[SimilarityEvaluation] = None
+ similarity_metric_type: MetricType = None
+ similarity_threshold: float = None
+ similarity_threshold_long: float = None
+ normalize: bool = None
+ #==================================================#
+
+ # Configure cache behavior based on embedding model type
+ if (embedding_model == EmbeddingModel.HUGGINGFACE_ALL_MPNET_BASE_V2
+ or embedding_model == EmbeddingModel.HUGGINGFACE_ALL_MINILM_L6_V2
+ or embedding_model == EmbeddingModel.HUGGINGFACE_ALL_MINILM_L12_V2):
+ query_pre_embedding_func = query_with_role
+ insert_pre_embedding_func = query_with_role
+ post_process_messages_func = first
+ similarity_evaluation = None # Uses the built-in cosine similarity evaluation in milvus
+ similarity_metric_type = MetricType.COSINE
+ similarity_threshold = 0.9
+ similarity_threshold_long = 0.9
+ normalize = False
+
+ elif embedding_model == EmbeddingModel.DATA2VEC_AUDIO:
+ query_pre_embedding_func = query_multi_splicing
+ insert_pre_embedding_func = insert_multi_splicing
+ post_process_messages_func = first
+ similarity_evaluation = SearchDistanceEvaluation()
+ similarity_metric_type = MetricType.L2
+ similarity_threshold = 0.95
+ similarity_threshold_long = 0.95
+ normalize = True
+
+ # Add configurations for additional embedding models as needed
+ else:
+ modelcache_log.error(f"Please add configuration for {embedding_model} in modelcache/cache.py.")
+ raise CacheError(f"Please add configuration for {embedding_model} in modelcache/cache.py.")
+
+ # ====================== Data manager ==============================#
+
+ # Create coordinated data manager with all storage backends
+ data_manager = DataManager.get(
+ SQLStorage.get(sql_storage, config=sql_config),
+ VectorStorage.get(
+ name=vector_storage,
+ dimension=dimension,
+ config=vector_config,
+ metric_type=similarity_metric_type,
+ ),
+ memory_cache_policy='ARC',
+ max_size=10000,
+ normalize=normalize,
+ )
+
+ #================== Cache Initialization ====================#
+
+ # Create fully configured Cache instance
+ cache = Cache(
+ embedding_model = embedding_model,
+ similarity_metric_type = similarity_metric_type,
+ data_manager = data_manager,
+ report = Report(),
+ embedding_func = embedding_dispatcher.embed,
+ query_pre_embedding_func = query_pre_embedding_func,
+ insert_pre_embedding_func = insert_pre_embedding_func,
+ similarity_evaluation = similarity_evaluation,
+ post_process_messages_func = post_process_messages_func,
+ similarity_threshold = similarity_threshold,
+ similarity_threshold_long = similarity_threshold_long,
+ prompts = None,
+ log_time_func = None,
+ )
+ return cache, event_loop
\ No newline at end of file
diff --git a/modelcache/config/chromadb_config.ini b/modelcache/config/chromadb_config.ini
new file mode 100644
index 0000000..7400f80
--- /dev/null
+++ b/modelcache/config/chromadb_config.ini
@@ -0,0 +1,2 @@
+[chromadb]
+persist_directory=''
diff --git a/modelcache/config/elasticsearch_config.ini b/modelcache/config/elasticsearch_config.ini
new file mode 100644
index 0000000..ba87a5a
--- /dev/null
+++ b/modelcache/config/elasticsearch_config.ini
@@ -0,0 +1,5 @@
+[elasticsearch]
+host = ''
+port = ''
+user = ''
+password = ''
\ No newline at end of file
diff --git a/modelcache/config/milvus_config.ini b/modelcache/config/milvus_config.ini
index f5bd532..1e9aa66 100644
--- a/modelcache/config/milvus_config.ini
+++ b/modelcache/config/milvus_config.ini
@@ -1,5 +1,5 @@
[milvus]
-host = ''
-port = ''
+host = localhost
+port = 19530
user = ''
password = ''
\ No newline at end of file
diff --git a/modelcache/config/mysql_config.ini b/modelcache/config/mysql_config.ini
index 2c63f0e..3d92385 100644
--- a/modelcache/config/mysql_config.ini
+++ b/modelcache/config/mysql_config.ini
@@ -1,6 +1,6 @@
[mysql]
-host = ''
-port = ''
-username = ''
-password = ''
-database = ''
+host = localhost
+port = 3306
+username = modelcache
+password = modelcache
+database = modelcache
diff --git a/modelcache/config/redis_config.ini b/modelcache/config/redis_config.ini
new file mode 100644
index 0000000..a1cdb3d
--- /dev/null
+++ b/modelcache/config/redis_config.ini
@@ -0,0 +1,5 @@
+[redis]
+host = ''
+port = ''
+user = ''
+password = ''
diff --git a/modelcache/embedding/__init__.py b/modelcache/embedding/__init__.py
index 03b6762..19f9b94 100644
--- a/modelcache/embedding/__init__.py
+++ b/modelcache/embedding/__init__.py
@@ -1,32 +1,2 @@
# -*- coding: utf-8 -*-
-from modelcache.utils.lazy_import import LazyImport
-huggingface = LazyImport("huggingface", globals(), "modelcache.embedding.huggingface")
-data2vec = LazyImport("data2vec", globals(), "modelcache.embedding.data2vec")
-llmEmb = LazyImport("llmEmb", globals(), "modelcache.embedding.llmEmb")
-fasttext = LazyImport("fasttext", globals(), "modelcache.embedding.fasttext")
-paddlenlp = LazyImport("paddlenlp", globals(), "modelcache.embedding.paddlenlp")
-timm = LazyImport("timm", globals(), "modelcache.embedding.timm")
-
-
-def Huggingface(model="sentence-transformers/all-mpnet-base-v2"):
- return huggingface.Huggingface(model)
-
-
-def Data2VecAudio(model="facebook/data2vec-audio-base-960h"):
- return data2vec.Data2VecAudio(model)
-
-
-def LlmEmb2vecAudio():
- return llmEmb.LlmEmb2Vec()
-
-
-def FastText(model="en", dim=None):
- return fasttext.FastText(model, dim)
-
-
-def PaddleNLP(model="ernie-3.0-medium-zh"):
- return paddlenlp.PaddleNLP(model)
-
-
-def Timm(model="resnet50", device="default"):
- return timm.Timm(model, device)
+from modelcache.embedding.base import EmbeddingModel, MetricType, BaseEmbedding
\ No newline at end of file
diff --git a/modelcache/embedding/base.py b/modelcache/embedding/base.py
index 4b316aa..70ae240 100644
--- a/modelcache/embedding/base.py
+++ b/modelcache/embedding/base.py
@@ -1,5 +1,46 @@
# -*- coding: utf-8 -*-
-from abc import ABCMeta, abstractmethod
+from abc import abstractmethod, ABCMeta
+
+from modelcache.utils.error import CacheError
+from modelcache.utils.lazy_import import LazyImport
+from enum import Enum
+
+from modelcache.utils.log import modelcache_log
+
+huggingface = LazyImport("huggingface", globals(), "modelcache.embedding.huggingface")
+data2vec = LazyImport("data2vec", globals(), "modelcache.embedding.data2vec")
+llmEmb = LazyImport("llmEmb", globals(), "modelcache.embedding.llmEmb")
+fasttext = LazyImport("fasttext", globals(), "modelcache.embedding.fasttext")
+paddlenlp = LazyImport("paddlenlp", globals(), "modelcache.embedding.paddlenlp")
+timm = LazyImport("timm", globals(), "modelcache.embedding.timm")
+huggingface_tei = LazyImport("huggingface_tei", globals(), "modelcache.embedding.huggingface_tei")
+bge_m3 = LazyImport("bge_m3", globals(), "modelcache.embedding.bge_m3")
+
+# define the embedding model enum
+class EmbeddingModel(Enum):
+ """
+ Enum for different embedding models.
+ """
+ # todo: fill in the dimension and model_path for each embedding model as needed
+ HUGGINGFACE_ALL_MPNET_BASE_V2 = {"dimension":768, "model_path":"sentence-transformers/all-mpnet-base-v2"}
+ HUGGINGFACE_ALL_MINILM_L6_V2 = {"dimension":384, "model_path":"sentence-transformers/all-MiniLM-L6-v2"}
+ HUGGINGFACE_ALL_MINILM_L12_V2 = {"dimension":384, "model_path":"sentence-transformers/all-MiniLM-L12-v2"}
+ DATA2VEC_AUDIO = {"dimension":768, "model_path":"model/text2vec-base-chinese/"}
+ LLM_EMB2VEC_AUDIO = {"dimension":None, "model_path":None}
+ FASTTEXT = {"dimension":None, "model_path":None}
+ PADDLE_NLP = {"dimension":None, "model_path":None}
+ TIMM = {"dimension":None, "model_path":None}
+ HUGGINGFACE_TEI = {"dimension":None, "model_path":None}
+ BGE_M3 = {"dimension":None, "model_path":None}
+
+
+class MetricType(Enum):
+ """
+ Enum for different metric types used in similarity evaluation.
+ Different models may require different metrics for optimal performance.
+ """
+ COSINE = "COSINE"
+ L2 = "L2"
class BaseEmbedding(metaclass=ABCMeta):
@@ -15,3 +56,62 @@ def to_embeddings(self, data, **kwargs):
@abstractmethod
def dimension(self) -> int:
return 0
+
+ @staticmethod
+ def get(model:EmbeddingModel, **kwargs):
+ """
+ Get the embedding model instance based on the specified model type.
+ :param model: The embedding model type.
+ :type model: EmbeddingModel
+ :param kwargs: Additional parameters for the model.
+ :return: An instance of the specified embedding model.
+ :rtype: BaseEmbedding
+ :raises ValueError: If the specified model type is not supported.
+ """
+ if model == EmbeddingModel.HUGGINGFACE_ALL_MPNET_BASE_V2:
+ model_path = kwargs.pop("model_path","sentence-transformers/all-mpnet-base-v2")
+ return huggingface.Huggingface(model_path)
+
+ elif model == EmbeddingModel.HUGGINGFACE_ALL_MINILM_L6_V2:
+ model_path = kwargs.pop("model_path","sentence-transformers/all-MiniLM-L6-v2")
+ return huggingface.Huggingface(model_path)
+
+ elif model == EmbeddingModel.HUGGINGFACE_ALL_MINILM_L12_V2:
+ model_path = kwargs.pop("model_path","sentence-transformers/all-MiniLM-L12-v2")
+ return huggingface.Huggingface(model_path)
+
+ elif model == EmbeddingModel.DATA2VEC_AUDIO:
+ model_path = kwargs.pop("model_path","model/text2vec-base-chinese/")
+ return data2vec.Data2VecAudio(model_path)
+
+ elif model == EmbeddingModel.LLM_EMB2VEC_AUDIO:
+ return llmEmb.LlmEmb2Vec()
+
+ elif model == EmbeddingModel.FASTTEXT:
+ model_path = kwargs.pop("model_path","en")
+ dim = kwargs.pop("dim", None)
+ return fasttext.FastText(model_path, dim)
+
+ elif model == EmbeddingModel.PADDLE_NLP:
+ model_path = kwargs.pop("model_path", "ernie-3.0-medium-zh")
+ return paddlenlp.PaddleNLP(model_path)
+
+ elif model == EmbeddingModel.TIMM:
+ model_path = kwargs.pop("model_path", "resnet50")
+ device = kwargs.pop("device", "default")
+ return timm.Timm(model_path, device)
+
+ elif model == EmbeddingModel.HUGGINGFACE_TEI:
+ base_url = kwargs.pop("base_url")
+ model_path = kwargs.pop("model_path")
+ return huggingface_tei.HuggingfaceTEI(base_url, model_path)
+
+ elif model == EmbeddingModel.BGE_M3:
+ model_path = kwargs.pop("model_path","model/bge-m3")
+ return bge_m3.BgeM3Embedding(model_path)
+
+ else:
+ modelcache_log.error(f"Please add configuration for {model} in modelcache/embedding/base.py.")
+ raise CacheError(f"Please add configuration for {model} in modelcache/embedding/base.py.")
+
+
diff --git a/modelcache/embedding/bge_m3.py b/modelcache/embedding/bge_m3.py
new file mode 100644
index 0000000..10cc15c
--- /dev/null
+++ b/modelcache/embedding/bge_m3.py
@@ -0,0 +1,28 @@
+# -*- coding: utf-8 -*-
+import numpy as np
+from modelcache.embedding.base import BaseEmbedding
+from transformers import AutoTokenizer, AutoModel
+from FlagEmbedding import BGEM3FlagModel
+
+class BgeM3Embedding(BaseEmbedding):
+ def __init__(self, model_path: str = "model/bge-m3"):
+ self.tokenizer = AutoTokenizer.from_pretrained(model_path)
+ self.model = AutoModel.from_pretrained(model_path)
+
+ self.bge_model = BGEM3FlagModel(model_name_or_path=model_path,
+ model=self.model,
+ tokenizer=self.tokenizer,
+ use_fp16=False)
+
+ self.__dimension = 768
+
+ def to_embeddings(self, data, **_):
+ if not isinstance(data, list):
+ data = [data]
+
+ embeddings = self.bge_model.encode(data, batch_size=12, max_length=8192)['dense_vecs']
+ return np.array(embeddings).astype("float32")
+
+ @property
+ def dimension(self):
+ return self.__dimension
\ No newline at end of file
diff --git a/modelcache/embedding/data2vec.py b/modelcache/embedding/data2vec.py
index 274eb5a..8c1ab2f 100644
--- a/modelcache/embedding/data2vec.py
+++ b/modelcache/embedding/data2vec.py
@@ -8,30 +8,29 @@
def mean_pooling(model_output, attention_mask):
- token_embeddings = model_output[0] # First element of model_output contains all token embeddings
+ token_embeddings = model_output[0]
input_mask_expanded = attention_mask.unsqueeze(-1).expand(token_embeddings.size()).float()
return torch.sum(token_embeddings * input_mask_expanded, 1) / torch.clamp(input_mask_expanded.sum(1), min=1e-9)
class Data2VecAudio(BaseEmbedding):
- def __init__(self, model: str = "sentence-transformers/all-MiniLM-L6-v2"):
+ def __init__(self, model):
current_dir = os.path.dirname(os.path.abspath(__file__))
parent_dir = os.path.dirname(current_dir)
model_dir = os.path.dirname(parent_dir)
- model = os.path.join(model_dir, 'model/text2vec-base-chinese/')
+ model_path = os.path.join(model_dir, model)
+
+ self.device = 'cuda' if torch.cuda.is_available() else 'cpu'
+ self.tokenizer = BertTokenizer.from_pretrained(model_path, local_files_only=True)
+ self.model = BertModel.from_pretrained(model_path, local_files_only=True)
try:
self.__dimension = self.model.config.hidden_size
except Exception:
from transformers import AutoConfig
-
config = AutoConfig.from_pretrained(model)
self.__dimension = config.hidden_size
- self.device = 'cuda' if torch.cuda.is_available() else 'cpu'
- self.tokenizer = BertTokenizer.from_pretrained(model, local_files_only=True)
- self.model = BertModel.from_pretrained(model, local_files_only=True)
-
def to_embeddings(self, data, **_):
encoded_input = self.tokenizer(data, padding=True, truncation=True, return_tensors='pt')
num_tokens = sum(map(len, encoded_input['input_ids']))
diff --git a/modelcache/embedding/embedding_dispatcher.py b/modelcache/embedding/embedding_dispatcher.py
new file mode 100644
index 0000000..16f4beb
--- /dev/null
+++ b/modelcache/embedding/embedding_dispatcher.py
@@ -0,0 +1,84 @@
+import multiprocessing
+import threading
+import uuid
+import asyncio
+import psutil
+from asyncio import Future, AbstractEventLoop
+
+from modelcache.embedding import EmbeddingModel
+from modelcache.embedding.base import BaseEmbedding
+
+
+def worker_func(embedding_model: EmbeddingModel, model_path, task_queue, result_queue, worker_id):
+ """Worker function that runs in separate processes to generate embeddings."""
+ base_embedding = BaseEmbedding.get(embedding_model, model_path=model_path)
+ print(f"Embedding worker {worker_id} started.")
+ try:
+ while True:
+ job_id, data = task_queue.get() # Get task from queue
+ try:
+ result = base_embedding.to_embeddings(data) # Generate embedding
+ except Exception as e:
+ result = e
+ result_queue.put((job_id, result)) # Send result back
+ except KeyboardInterrupt:
+ print(f"Embedding worker {worker_id} stopped.")
+ except Exception as e:
+ print(f"Embedding worker {worker_id} encountered an error: {e}")
+
+
+class EmbeddingDispatcher:
+ """Manages a pool of worker processes for parallel embedding generation."""
+
+ def __init__(
+ self,
+ embedding_model: EmbeddingModel,
+ model_path: str,
+ event_loop: AbstractEventLoop,
+ num_workers: int
+ ):
+ """Initialize the dispatcher with worker processes."""
+ if num_workers <= 0:
+ raise ValueError("Number of workers must be greater than 0.")
+
+ self.task_queue = multiprocessing.Queue() # Tasks to workers
+ self.result_queue = multiprocessing.Queue() # Results from workers
+ self.futures: dict[str, asyncio.Future] = {} # Pending futures
+ self.event_loop = event_loop
+ self._start_result_collector_thread() # Start result collection thread
+
+ # Start worker processes
+ self.workers = []
+ for i in range(num_workers):
+ p = multiprocessing.Process(
+ target=worker_func,
+ args=(embedding_model, model_path, self.task_queue, self.result_queue, i)
+ )
+ p.daemon = True
+ p.start()
+ psutil.Process(p.pid).nice(psutil.HIGH_PRIORITY_CLASS)
+ self.workers.append(p)
+
+ def _start_result_collector_thread(self):
+ """Start a thread to collect results from worker processes."""
+ def collect():
+ while True:
+ job_id, result = self.result_queue.get() # Get result from queue
+ future = self.futures.pop(job_id, None) # Retrieve future
+ if future:
+ self.event_loop.call_soon_threadsafe(
+ future.set_exception if isinstance(result, Exception) else future.set_result,
+ result
+ )
+
+ t = threading.Thread(target=collect, daemon=True)
+ t.start()
+
+ def embed(self, data: str) -> Future:
+ """Submit a task for embedding generation."""
+ job_id = str(uuid.uuid4()) # Generate unique job ID
+ future = asyncio.get_running_loop().create_future() # Create future
+ self.futures[job_id] = future # Store future
+ self.task_queue.put((job_id, data)) # Add task to queue
+ return future
+
diff --git a/modelcache/embedding/huggingface.py b/modelcache/embedding/huggingface.py
index 8c1434d..77ca08f 100644
--- a/modelcache/embedding/huggingface.py
+++ b/modelcache/embedding/huggingface.py
@@ -1,34 +1,21 @@
# -*- coding: utf-8 -*-
-import numpy as np
-
-from modelcache.utils import import_huggingface, import_torch
from modelcache.embedding.base import BaseEmbedding
-
-import_torch()
-import_huggingface()
-
-import torch # pylint: disable=C0413
-from transformers import AutoTokenizer, AutoModel # pylint: disable=C0413
-
+from sentence_transformers import SentenceTransformer
class Huggingface(BaseEmbedding):
- def __init__(self, model: str = "sentence-transformers/all-MiniLM-L6-v2"):
- self.model = AutoModel.from_pretrained(model, local_files_only=True)
- self.model.eval()
-
- # self.tokenizer = AutoTokenizer.from_pretrained(model)
- self.tokenizer = AutoTokenizer.from_pretrained(model, local_files_only=True)
- if not self.tokenizer.pad_token:
- self.tokenizer.pad_token = "[PAD]"
+ def __init__(self, model: str):
+ self.model = SentenceTransformer(model,tokenizer_kwargs={
+ "clean_up_tokenization_spaces":False
+ })
try:
self.__dimension = self.model.config.hidden_size
- except Exception: # pylint: disable=W0703
- from transformers import AutoConfig # pylint: disable=C0415
+ except Exception:
+ from transformers import AutoConfig
config = AutoConfig.from_pretrained(model)
self.__dimension = config.hidden_size
- def to_embeddings(self, data, **_):
+ def to_embeddings(self, data: str, **_):
"""Generate embedding given text input
:param data: text in string.
@@ -36,24 +23,11 @@ def to_embeddings(self, data, **_):
:return: a text embedding in shape of (dim,).
"""
- if not isinstance(data, list):
- data = [data]
- inputs = self.tokenizer(
- data, padding=True, truncation=True, return_tensors="pt"
- )
- outs = self.model(**inputs).last_hidden_state
- emb = self.post_proc(outs, inputs).squeeze(0).detach().numpy()
- return np.array(emb).astype("float32")
- def post_proc(self, token_embeddings, inputs):
- attention_mask = inputs["attention_mask"]
- input_mask_expanded = (
- attention_mask.unsqueeze(-1).expand(token_embeddings.size()).float()
- )
- sentence_embs = torch.sum(
- token_embeddings * input_mask_expanded, 1
- ) / torch.clamp(input_mask_expanded.sum(1), min=1e-9)
- return sentence_embs
+ if not data:
+ raise ValueError("No data provided for embedding.")
+ embeddings = self.model.encode(data)
+ return embeddings[0] if len(data) == 1 else embeddings
@property
def dimension(self):
diff --git a/modelcache/embedding/huggingface_tei.py b/modelcache/embedding/huggingface_tei.py
new file mode 100644
index 0000000..94075fe
--- /dev/null
+++ b/modelcache/embedding/huggingface_tei.py
@@ -0,0 +1,32 @@
+# -*- coding: utf-8 -*-
+import requests
+import numpy as np
+from modelcache.embedding.base import BaseEmbedding
+
+class HuggingfaceTEI(BaseEmbedding):
+ def __init__(self, base_url: str, model: str):
+ self.base_url = base_url
+ self.model = model
+ self.headers = {
+ 'accept': 'application/json',
+ 'Content-Type': 'application/json',
+ }
+ self.__dimension = self.to_embeddings('test').shape[0]
+
+ def to_embeddings(self, data, **_):
+ json_data = {
+ 'input': data,
+ 'model': self.model,
+ }
+
+ response = requests.post(self.base_url, headers=self.headers, json=json_data)
+ embedding = response.json()['data'][0]['embedding']
+ return np.array(embedding)
+
+ @property
+ def dimension(self):
+ """Embedding dimension.
+
+ :return: embedding dimension
+ """
+ return self.__dimension
diff --git a/modelcache/embedding/timm.py b/modelcache/embedding/timm_embedding.py
similarity index 100%
rename from modelcache/embedding/timm.py
rename to modelcache/embedding/timm_embedding.py
diff --git a/modelcache/manager/__init__.py b/modelcache/manager/__init__.py
index 59a5ffe..40a96af 100644
--- a/modelcache/manager/__init__.py
+++ b/modelcache/manager/__init__.py
@@ -1,5 +1 @@
# -*- coding: utf-8 -*-
-from modelcache.manager.scalar_data import CacheBase
-from modelcache.manager.vector_data import VectorBase
-from modelcache.manager.object_data import ObjectBase
-from modelcache.manager.factory import get_data_manager
diff --git a/modelcache/manager/data_manager.py b/modelcache/manager/data_manager.py
index db8b776..7f019cd 100644
--- a/modelcache/manager/data_manager.py
+++ b/modelcache/manager/data_manager.py
@@ -6,19 +6,13 @@
import numpy as np
import cachetools
from abc import abstractmethod, ABCMeta
-from typing import List, Any, Optional, Union
-from modelcache.manager.scalar_data.base import (
- CacheStorage,
- CacheData,
- DataType,
- Answer,
- Question
-)
+from typing import List, Any, Optional
+from typing import Union, Callable
+from modelcache.manager.scalar_data.base import CacheStorage,CacheData,DataType,Answer,Question
from modelcache.utils.error import CacheError, ParamError
-from modelcache.manager.vector_data.base import VectorBase, VectorData
+from modelcache.manager.vector_data.base import VectorStorage, VectorData
from modelcache.manager.object_data.base import ObjectBase
-from modelcache.manager.eviction import EvictionBase
-from modelcache.manager.eviction_manager import EvictionManager
+from modelcache.manager.eviction.memory_cache import MemoryCacheEviction
from modelcache.utils.log import modelcache_log
@@ -34,9 +28,7 @@ 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, questions: List[Any], answers: List[Any], embedding_datas: List[Any], model:Any):
pass
@abstractmethod
@@ -58,9 +50,11 @@ def search(self, embedding_data, **kwargs):
def delete(self, id_list, **kwargs):
pass
+ @abstractmethod
def truncate(self, model_name):
pass
+ @abstractmethod
def flush(self):
pass
@@ -68,6 +62,30 @@ def flush(self):
def close(self):
pass
+ @staticmethod
+ def get(
+ cache_base: Union[CacheStorage, str] = None,
+ vector_base: Union[VectorStorage, str] = None,
+ object_base: Union[ObjectBase, str] = None,
+ max_size: int = 3,
+ clean_size: int = 1,
+ memory_cache_policy: str = "ARC",
+ data_path: str = "data_map.txt",
+ get_data_container: Callable = None,
+ normalize: bool = True
+ ):
+ if not cache_base and not vector_base:
+ return MapDataManager(data_path, max_size, get_data_container)
+
+ if isinstance(cache_base, str):
+ cache_base = CacheStorage.get(name=cache_base)
+ if isinstance(vector_base, str):
+ vector_base = VectorStorage.get(name=vector_base)
+ if isinstance(object_base, str):
+ object_base = ObjectBase.get(name=object_base)
+ assert cache_base and vector_base
+ return SSDataManager(cache_base, vector_base, object_base, max_size, clean_size,normalize, memory_cache_policy)
+
class MapDataManager(DataManager):
def __init__(self, data_path, max_size, get_data_container=None):
@@ -146,23 +164,33 @@ class SSDataManager(DataManager):
def __init__(
self,
s: CacheStorage,
- v: VectorBase,
+ v: VectorStorage,
o: Optional[ObjectBase],
max_size,
clean_size,
+ normalize: bool,
policy="LRU",
):
self.max_size = max_size
self.clean_size = clean_size
- self.s = s
- self.v = v
- self.o = o
-
- def save(self, question, answer, embedding_data, **kwargs):
+ self.s = s # SQL storage
+ self.v = v # Vector storage
+ self.o = o # Object storage (optional)
+ self.normalize = normalize
+
+ # Initialize memory cache with specified eviction policy
+ self.eviction_base = MemoryCacheEviction(
+ policy=policy,
+ maxsize=max_size,
+ clean_size=clean_size)
+
+ def save(self, questions: List[any], answers: List[any], embedding_datas: List[any], **kwargs):
+ """Save multiple questions, answers, and embeddings to storage."""
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 response log to SQL storage for analytics."""
save_query_start_time = time.time()
self.s.insert_query_resp(query_resp_dict, **kwargs)
save_query_delta_time = '{}s'.format(round(time.time() - save_query_start_time, 2))
@@ -193,61 +221,97 @@ def _process_question_data(self, question: Union[str, Question]):
def import_data(
self, questions: List[Any], answers: List[Answer], embedding_datas: List[Any], model: Any
):
+ """
+ Add multiple cache entries into all storage backends.
+
+ Coordinates data insertion across SQL, vector, and object storage,
+ with memory cache population and optional vector normalization.
+ """
if len(questions) != len(answers) or len(questions) != len(embedding_datas):
raise ParamError("Make sure that all parameters have the same length")
cache_datas = []
- embedding_datas = [
- normalize(embedding_data) for embedding_data in embedding_datas
- ]
+ # Normalize embedding vectors if configured
+ if self.normalize:
+ embedding_datas = [
+ normalize(embedding_data) for embedding_data in embedding_datas
+ ]
- for i, embedding_data in enumerate(embedding_datas):
+ for embedding_data, answer, question in zip(embedding_datas,answers,questions):
if self.o is not None:
- ans = self._process_answer_data(answers[i])
- else:
- ans = answers[i]
+ answer = self._process_answer_data(answer)
- question = questions[i]
embedding_data = embedding_data.astype("float32")
- cache_datas.append([ans, question, embedding_data, model])
+ cache_datas.append([answer, question, embedding_data, model])
+ # Insert into SQL storage and get generated IDs
ids = self.s.batch_insert(cache_datas)
- logging.info('ids: {}'.format(ids))
- self.v.mul_add(
- [
- VectorData(id=ids[i], data=embedding_data)
- for i, embedding_data in enumerate(embedding_datas)
- ],
- model
- )
+ # Prepare vector data and populate memory cache
+ datas = []
+ for _id,embedding_data,cache_data in zip(ids,embedding_datas,cache_datas):
+ datas.append(VectorData(id=_id, data=embedding_data.astype("float32")))
+ self.eviction_base.put([(_id, cache_data)],model=model)
+ self.v.mul_add(datas,model)
def get_scalar_data(self, res_data, **kwargs) -> Optional[CacheData]:
- cache_data = self.s.get_data_by_id(res_data[1])
+ """
+ Retrieve scalar data with multi-level caching strategy.
+
+ First checks memory cache, then falls back to SQL storage.
+ """
+ model = kwargs.pop("model")
+ _id = res_data[1]
+
+ # Try to get from memory cache first (fastest)
+ cache_hit = self.eviction_base.get(_id, model=model)
+ if cache_hit is not None:
+ return cache_hit
+ cache_data = self.s.get_data_by_id(_id)
if cache_data is None:
return None
+ self.eviction_base.put([(_id, cache_data)], model=model)
return cache_data
def update_hit_count(self, primary_id, **kwargs):
+ """Update hit count statistics in SQL storage."""
self.s.update_hit_count_by_id(primary_id)
def hit_cache_callback(self, res_data, **kwargs):
+ """Callback executed on cache hit to update memory cache."""
self.eviction_base.get(res_data[1])
def search(self, embedding_data, **kwargs):
+ """
+ Search for similar vectors in vector storage.
+
+ Applies normalization if configured and delegates to vector backend.
+ """
model = kwargs.pop("model", None)
- embedding_data = normalize(embedding_data)
+ if self.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)
def delete(self, id_list, **kwargs):
- model = kwargs.pop("model", None)
+ """
+ Delete cache entries from all storage backends.
+
+ Removes from memory cache, vector storage, and marks as deleted in SQL.
+ Returns detailed status of deletion operations.
+ """
+ model = kwargs.pop("model")
try:
+ # Remove from memory cache
+ for id in id_list:
+ self.eviction_base.get_cache(model).pop(id, None)
+ # Delete from vector storage
v_delete_count = self.v.delete(ids=id_list, model=model)
except Exception as e:
return {'status': 'failed', 'milvus': 'delete milvus data failed, please check! e: {}'.format(e),
'mysql': 'unexecuted'}
try:
+ # Mark as deleted in SQL storage
s_delete_count = self.s.mark_deleted(id_list)
except Exception as e:
return {'status': 'failed', 'milvus': 'success',
@@ -256,34 +320,44 @@ def delete(self, id_list, **kwargs):
return {'status': 'success', 'milvus': 'delete_count: '+str(v_delete_count),
'mysql': 'delete_count: '+str(s_delete_count)}
- def truncate(self, model_name):
- # model = kwargs.pop("model", None)
- # drop milvus data
+ def create_index(self, model, **kwargs):
+ """Create vector index for a specific model."""
+ return self.v.create(model)
+
+ def truncate(self, model):
+ """
+ Truncate all data for a specific model across all storage backends.
+
+ Clears memory cache, rebuilds vector storage, and deletes SQL data.
+ Returns detailed status of truncation operations.
+ """
+ # Clear memory cache data
+ self.eviction_base.clear(model)
+
+ # Rebuild vector storage (drops and recreates collection)
try:
- resp = self.v.rebuild_col(model_name)
+ vector_resp = self.v.rebuild_col(model)
except Exception as e:
- return {'status': 'failed', 'milvus': 'truncate milvus data failed, please check! e: {}'.format(e),
- 'mysql': 'unexecuted'}
- if resp:
- return {'status': 'failed', 'milvus': resp, 'mysql': 'unexecuted'}
+ return {'status': 'failed', 'VectorDB': 'truncate VectorDB data failed, please check! e: {}'.format(e),
+ 'ScalarDB': 'unexecuted'}
+ if vector_resp:
+ return {'status': 'failed', 'VectorDB': vector_resp, 'ScalarDB': 'unexecuted'}
+
+ # Delete scalar data from SQL storage
try:
- delete_count = self.s.model_deleted(model_name)
+ delete_count = self.s.model_deleted(model)
except Exception as e:
- # return 'truncate milvus data failed, please check!'
- return {'status': 'failed', 'milvus': 'rebuild',
- 'mysql': 'truncate mysql data failed, please check! e: {}'.format(e)}
- return {'status': 'success', 'milvus': 'rebuild', 'mysql': 'delete_count: ' + str(delete_count)}
+ return {'status': 'failed', 'VectorDB': 'rebuild',
+ 'ScalarDB': 'truncate scalar data failed, please check! e: {}'.format(e)}
+ return {'status': 'success', 'VectorDB': 'rebuild', 'ScalarDB': 'delete_count: ' + str(delete_count)}
def flush(self):
+ """Flush all storage backends to ensure data persistence."""
self.s.flush()
self.v.flush()
def close(self):
+ """Close all storage connections and release resources."""
self.s.close()
self.v.close()
-
-# if __name__ == '__main__':
-# from modelcache.manager import CacheBase, VectorBase, get_data_manager
-# data_manager = get_data_manager(CacheBase('mysql'), VectorBase('milvus', dimension=128))
-# data_manager.save('hello', 'hi', np.random.random((128,)).astype('float32'), model='gptcode_6b')
diff --git a/modelcache/manager/eviction/__init__.py b/modelcache/manager/eviction/__init__.py
index 8ca7a3d..633f866 100644
--- a/modelcache/manager/eviction/__init__.py
+++ b/modelcache/manager/eviction/__init__.py
@@ -1,10 +1,2 @@
# -*- coding: utf-8 -*-
-from modelcache.utils.lazy_import import LazyImport
-eviction_manager = LazyImport(
- "eviction_manager", globals(), "modelcache.manager.eviction.manager"
-)
-
-
-def EvictionBase(name: str, **kwargs):
- return eviction_manager.EvictionBase.get(name, **kwargs)
diff --git a/modelcache/manager/eviction/arc_cache.py b/modelcache/manager/eviction/arc_cache.py
new file mode 100644
index 0000000..4f97231
--- /dev/null
+++ b/modelcache/manager/eviction/arc_cache.py
@@ -0,0 +1,133 @@
+from cachetools import Cache
+from collections import OrderedDict
+from readerwriterlock import rwlock
+
+_sentinel = object()
+
+class ARC(Cache):
+ """
+ Adaptive Replacement Cache (ARC) implementation.
+
+ ARC maintains four lists (T1, T2, B1, B2) to adaptively balance
+ between LRU and LFU eviction strategies based on access patterns.
+ """
+
+ def __init__(self, maxsize, getsizeof=None):
+ """Initialize ARC cache with maximum size."""
+ super().__init__(maxsize, getsizeof)
+ self.t1 = OrderedDict() # Recent items
+ self.t2 = OrderedDict() # Frequent items
+ self.b1 = OrderedDict() # Ghost entries for T1
+ self.b2 = OrderedDict() # Ghost entries for T2
+ self.p = 0 # Adaptive parameter
+ self._rw_lock = rwlock.RWLockWrite() # Thread safety
+
+ def __len__(self):
+ """Return total number of cached items."""
+ return len(self.t1) + len(self.t2)
+
+ def __contains__(self, key):
+ """Check if key exists in cache."""
+ return key in self.t1 or key in self.t2
+
+ def _evict_internal(self):
+ """Internal method to evict items when cache is full."""
+ # Evict from cache lists to ghost lists
+ while len(self.t1) + len(self.t2) > self.maxsize:
+ if len(self.t1) > self.p or (len(self.t1) == 0 and len(self.t2) > 0):
+ key, value = self.t1.popitem(last=False)
+ self.b1[key] = value
+ else:
+ key, value = self.t2.popitem(last=False)
+ self.b2[key] = value
+
+ # Maintain ghost list sizes
+ while len(self.b1) > (self.maxsize - self.p):
+ self.b1.popitem(last=False)
+ while len(self.b2) > self.p:
+ self.b2.popitem(last=False)
+
+ def __setitem__(self, key, value):
+ """Insert or update a cache entry."""
+ with self._rw_lock.gen_wlock():
+ # Remove key from all lists first
+ for l in (self.t1, self.t2, self.b1, self.b2):
+ l.pop(key, None)
+ # Add to recent list (T1)
+ self.t1[key] = value
+ self.t1.move_to_end(key)
+ self._evict_internal()
+
+ def __getitem__(self, key):
+ """Retrieve a cache entry and update access pattern."""
+ with self._rw_lock.gen_wlock():
+ if key in self.t1:
+ # Move from recent to frequent list
+ value = self.t1.pop(key)
+ self.t2[key] = value
+ self.t2.move_to_end(key)
+ self.p = max(0, self.p - 1) # Adjust adaptive parameter
+ self._evict_internal()
+ return value
+ if key in self.t2:
+ # Access frequent list
+ value = self.t2.pop(key)
+ self.t2[key] = value
+ self.t2.move_to_end(key)
+ self.p = min(self.maxsize, self.p + 1) # Adjust adaptive parameter
+ self._evict_internal()
+ return value
+ if key in self.b1:
+ # Promote from ghost list B1 to frequent list T2
+ self.b1.pop(key)
+ self.p = min(self.maxsize, self.p + 1) # Adjust adaptive parameter
+ self._evict_internal()
+ value = super().__missing__(key)
+ self.t2[key] = value
+ self.t2.move_to_end(key)
+ return value
+ if key in self.b2:
+ # Promote from ghost list B2 to frequent list T2
+ self.b2.pop(key)
+ self.p = max(0, self.p - 1) # Adjust adaptive parameter
+ self._evict_internal()
+ value = super().__missing__(key)
+ self.t2[key] = value
+ self.t2.move_to_end(key)
+ return value
+ return super().__getitem__(key)
+
+ def __missing__(self, key):
+ """Handle missing keys."""
+ raise KeyError(key)
+
+ def pop(self, key, default=_sentinel):
+ """Remove a cache entry."""
+ with self._rw_lock.gen_wlock():
+ for l in (self.t1, self.t2, self.b1, self.b2):
+ if key in l:
+ return l.pop(key)
+ if default is _sentinel:
+ raise KeyError(key)
+ return default
+
+ def clear(self):
+ """Clear all cache entries."""
+ with self._rw_lock.gen_wlock():
+ self.t1.clear()
+ self.t2.clear()
+ self.b1.clear()
+ self.b2.clear()
+ self.p = 0
+ super().clear()
+
+ def __iter__(self):
+ """Iterate over cache keys."""
+ yield from self.t1
+ yield from self.t2
+
+ def __repr__(self):
+ """Return string representation of the cache."""
+ return (f"ARC(maxsize={self.maxsize}, p={self.p}, len={len(self)}, "
+ f"t1_len={len(self.t1)}, t2_len={len(self.t2)}, "
+ f"b1_len={len(self.b1)}, b2_len={len(self.b2)})")
diff --git a/modelcache/manager/eviction/base.py b/modelcache/manager/eviction/base.py
index 352c9a3..56efa54 100644
--- a/modelcache/manager/eviction/base.py
+++ b/modelcache/manager/eviction/base.py
@@ -9,11 +9,11 @@ class EvictionBase(metaclass=ABCMeta):
"""
@abstractmethod
- def put(self, objs: List[Any]):
+ def put(self, objs: List[Any], model:str):
pass
@abstractmethod
- def get(self, obj: Any):
+ def get(self, obj: Any, model:str):
pass
@property
diff --git a/modelcache/manager/eviction/manager.py b/modelcache/manager/eviction/manager.py
deleted file mode 100644
index 61579f0..0000000
--- a/modelcache/manager/eviction/manager.py
+++ /dev/null
@@ -1,25 +0,0 @@
-# -*- coding: utf-8 -*-
-from typing import Callable, List, Any
-from modelcache.utils.error import NotFoundError
-
-
-class EvictionBase:
- """
- EvictionBase to evict the cache data.
- """
-
- def __init__(self):
- raise EnvironmentError(
- "EvictionBase is designed to be instantiated, "
- "please using the `EvictionBase.get(name, policy, maxsize, clean_size)`."
- )
-
- @staticmethod
- def get(name: str, policy: str, maxsize: int, clean_size: int, on_evict: Callable[[List[Any]], None], **kwargs):
- if name in "memory":
- from modelcache.manager.eviction.memory_cache import MemoryCacheEviction
-
- eviction_base = MemoryCacheEviction(policy, maxsize, clean_size, on_evict, **kwargs)
- else:
- raise NotFoundError("eviction base", name)
- return eviction_base
diff --git a/modelcache/manager/eviction/memory_cache.py b/modelcache/manager/eviction/memory_cache.py
index 3cb487f..838e8cd 100644
--- a/modelcache/manager/eviction/memory_cache.py
+++ b/modelcache/manager/eviction/memory_cache.py
@@ -1,8 +1,10 @@
# -*- coding: utf-8 -*-
-from typing import Any, Callable, List
+from typing import Any, Callable, List, Tuple
import cachetools
from modelcache.manager.eviction.base import EvictionBase
+from .arc_cache import ARC
+from .wtinylfu_cache import W2TinyLFU
def popitem_wrapper(func, wrapper_func, clean_size):
@@ -17,28 +19,52 @@ def wrapper(*args, **kwargs):
class MemoryCacheEviction(EvictionBase):
- def __init__(self, policy: str, maxsize: int, clean_size: int, on_evict: Callable[[List[Any]], None], **kwargs):
+ def __init__(self, policy: str, maxsize: int, clean_size: int, **kwargs):
self._policy = policy.upper()
+ self.model_to_cache = dict()
+ self.maxsize = maxsize
+ self.clean_size = clean_size
+ self.kwargs = kwargs
+
+ def create_cache(self, model: str):
if self._policy == "LRU":
- self._cache = cachetools.LRUCache(maxsize=maxsize, **kwargs)
+ cache = cachetools.LRUCache(maxsize=self.maxsize, **self.kwargs)
elif self._policy == "LFU":
- self._cache = cachetools.LFUCache(maxsize=maxsize, **kwargs)
+ cache = cachetools.LFUCache(maxsize=self.maxsize, **self.kwargs)
elif self._policy == "FIFO":
- self._cache = cachetools.FIFOCache(maxsize=maxsize, **kwargs)
+ cache = cachetools.FIFOCache(maxsize=self.maxsize, **self.kwargs)
elif self._policy == "RR":
- self._cache = cachetools.RRCache(maxsize=maxsize, **kwargs)
+ cache = cachetools.RRCache(maxsize=self.maxsize, **self.kwargs)
+ elif self._policy == "WTINYLFU":
+ cache = W2TinyLFU(maxsize=self.maxsize)
+ elif self._policy == "ARC":
+ cache = ARC(maxsize=self.maxsize)
else:
- raise ValueError(f"Unknown policy {policy}")
+ raise ValueError(f"Unknown policy {self.policy}")
+ return cache
+
+ def put(self, objs: List[Tuple[Any, Any]], model: str):
+ cache = self.get_cache(model)
+ for key, value in objs:
+ cache[key] = value
+
+
+ def get(self, obj: Any, model: str):
+ cache = self.get_cache(model)
+ return cache.get(obj)
- self._cache.popitem = popitem_wrapper(self._cache.popitem, on_evict, clean_size)
- def put(self, objs: List[Any]):
- for obj in objs:
- self._cache[obj] = True
+ def clear(self, model: str):
+ self.model_to_cache.pop(model, None)
+
+
+ def get_cache(self, model: str):
+ if not model in self.model_to_cache:
+ self.model_to_cache[model] = self.create_cache(model)
+ return self.model_to_cache[model]
- def get(self, obj: Any):
- return self._cache.get(obj)
@property
def policy(self) -> str:
return self._policy
+
diff --git a/modelcache/manager/eviction/wtinylfu_cache.py b/modelcache/manager/eviction/wtinylfu_cache.py
new file mode 100644
index 0000000..378b1ed
--- /dev/null
+++ b/modelcache/manager/eviction/wtinylfu_cache.py
@@ -0,0 +1,172 @@
+from cachetools import LRUCache, Cache, LFUCache
+from readerwriterlock import rwlock
+import random
+
+class CountMinSketch:
+ def __init__(self, width=1024, depth=4, decay_interval=10000):
+ """Initialize Count-Min Sketch with specified dimensions."""
+ self.width = width
+ self.depth = depth
+ self.tables = [[0]*width for _ in range(depth)] # Hash tables
+ self.seeds = [random.randrange(1<<30) for _ in range(depth)] # Hash seeds
+ self.ops = 0 # Operation counter for decay trigger
+ self.decay_interval = decay_interval
+
+ def _hash(self, x, seed):
+ """Hash function for mapping items to table positions."""
+ return hash((x, seed)) % self.width
+
+ def add(self, x):
+ """Add an item and increment its frequency estimate."""
+ self.ops += 1
+ est = self.estimate(x) # Get current estimate
+ # Update all hash tables
+ for i, seed in enumerate(self.seeds):
+ idx = self._hash(x, seed)
+ if self.tables[i][idx] <= est:
+ self.tables[i][idx] += 1
+
+ # Periodic decay to handle changing patterns
+ if self.ops >= self.decay_interval:
+ self.decay()
+ self.ops = 0
+
+ def estimate(self, x):
+ """Estimate frequency of an item (minimum across all tables)."""
+ return min(self.tables[i][self._hash(x, seed)]
+ for i, seed in enumerate(self.seeds))
+
+ def decay(self):
+ """Decay all frequency counts by half."""
+ for table in self.tables:
+ for i in range(len(table)):
+ table[i] >>= 1 # Right shift (divide by 2)
+
+class W2TinyLFU(Cache):
+ """
+ Window Tiny LFU cache implementation.
+
+ Combines a small LRU window cache with a main cache divided into
+ probation and protected segments, using frequency estimation for
+ admission control.
+ """
+
+ def __init__(self, maxsize, window_pct=0.01):
+ """
+ Initialize W-TinyLFU cache.
+
+ Args:
+ maxsize: Maximum size of the cache
+ window_pct: Percentage of cache size for the window (default 1%)
+ """
+ super().__init__(maxsize)
+ self.window_size = max(1, int(maxsize * window_pct))
+ rest = maxsize - self.window_size
+ self.probation_size = rest // 2
+ self.protected_size = rest - self.probation_size
+
+ # Three cache segments
+ self.window = LRUCache(maxsize=self.window_size) # Recent items
+ self.probation = LFUCache(maxsize=self.probation_size) # New main cache items
+ self.protected = LFUCache(maxsize=self.protected_size) # Frequently accessed items
+
+ self.cms = CountMinSketch() # Frequency estimator
+ self.data = {} # Cache data storage
+ self._rw_lock = rwlock.RWLockWrite() # Read-write lock for thread safety
+
+ def __setitem__(self, key, value):
+ """Add or update an item in the cache."""
+ with self._rw_lock.gen_wlock():
+ self.data[key] = value
+ self._put(key)
+
+ def __getitem__(self, key):
+ """Retrieve an item from the cache."""
+ val = self.get(key, default=None)
+ if val is None:
+ raise KeyError(key)
+ return val
+
+ def __contains__(self, key):
+ """Check if an item exists in the cache."""
+ return key in self.window or key in self.probation or key in self.protected
+
+ def __delitem__(self, key):
+ """Remove an item from the cache."""
+ with self._rw_lock.gen_wlock():
+ self.data.pop(key, None)
+ self.window.pop(key, None)
+ self.probation.pop(key, None)
+ self.protected.pop(key, None)
+
+ def get(self, key, default=None):
+ """
+ Retrieve an item from the cache, updating its position
+ in the cache hierarchy if necessary.
+ """
+ if key in self.window:
+ self.window[key] = True
+ return self.data.get(key, default)
+ if key in self.protected:
+ self.protected[key] = True
+ return self.data.get(key, default)
+ if key in self.probation:
+ self.probation.pop(key)
+ if len(self.protected) >= self.protected_size:
+ demoted = next(iter(self.protected))
+ self.protected.pop(demoted)
+ self.probation[demoted] = True
+ self.protected[key] = True
+ return self.data.get(key, default)
+ return default
+
+ def _put(self, key):
+ """
+ Add an item to the cache, using frequency-based admission
+ control and eviction policies.
+ """
+ self.cms.add(key)
+ if key in self:
+ return
+
+ if len(self.window) < self.window_size:
+ self.window[key] = True
+ return
+
+ victim = next(iter(self.window))
+ self.window.pop(victim)
+
+ if self.cms.estimate(key) >= self.cms.estimate(victim):
+ self._admit_to_main(victim)
+ self._admit_to_main(key)
+ else:
+ self._admit_to_main(victim)
+ self.data.pop(key, None)
+
+ def _admit_to_main(self, key):
+ """
+ Admit an item to the main cache (probation or protected segment).
+ """
+ if key in self.protected or key in self.probation:
+ return
+ if self.probation_size == 0:
+ self.data.pop(key, None)
+ return
+ if len(self.probation) < self.probation_size:
+ self.probation[key] = True
+ elif self.probation:
+ evicted = next(iter(self.probation))
+ self.probation.pop(evicted)
+ self.probation[key] = True
+ self.data.pop(evicted, None)
+ else:
+ self.data.pop(key, None)
+
+ def clear(self):
+ """Clear all items from the cache."""
+ with self._rw_lock.gen_wlock():
+ self.window.clear()
+ self.probation.clear()
+ self.protected.clear()
+ self.data.clear()
+
diff --git a/modelcache/manager/eviction_manager.py b/modelcache/manager/eviction_manager.py
index 0a6a406..d44996c 100644
--- a/modelcache/manager/eviction_manager.py
+++ b/modelcache/manager/eviction_manager.py
@@ -29,10 +29,10 @@ def check_evict(self):
return True
return False
- def delete(self):
+ def delete(self,model):
mark_ids = self._scalar_storage.get_ids(deleted=True)
self._scalar_storage.clear_deleted_data()
- self._vector_base.delete(mark_ids)
+ self._vector_base.delete(mark_ids,model)
self.delete_count += 1
if self.delete_count >= self.REBUILD_CONDITION:
self.rebuild()
diff --git a/modelcache/manager/object_data/__init__.py b/modelcache/manager/object_data/__init__.py
index f1186c0..633f866 100644
--- a/modelcache/manager/object_data/__init__.py
+++ b/modelcache/manager/object_data/__init__.py
@@ -1,9 +1,2 @@
# -*- coding: utf-8 -*-
-from modelcache.utils.lazy_import import LazyImport
-object_manager = LazyImport(
- "object_manager", globals(), "modelcache.manager.object_data.manager"
-)
-
-def ObjectBase(name: str, **kwargs):
- return object_manager.ObjectBase.get(name, **kwargs)
diff --git a/modelcache/manager/object_data/base.py b/modelcache/manager/object_data/base.py
index 9ed00cf..41466f9 100644
--- a/modelcache/manager/object_data/base.py
+++ b/modelcache/manager/object_data/base.py
@@ -12,10 +12,6 @@ class ObjectBase(ABC):
def put(self, obj: Any) -> str:
pass
- @abstractmethod
- def get(self, obj: str) -> Any:
- pass
-
@abstractmethod
def get_access_link(self, obj: str) -> str:
pass
@@ -23,3 +19,7 @@ def get_access_link(self, obj: str) -> str:
@abstractmethod
def delete(self, to_delete: List[str]):
pass
+
+ @staticmethod
+ def get(name: str) -> Any:
+ pass
diff --git a/modelcache/manager/scalar_data/__init__.py b/modelcache/manager/scalar_data/__init__.py
index b63c430..40a96af 100644
--- a/modelcache/manager/scalar_data/__init__.py
+++ b/modelcache/manager/scalar_data/__init__.py
@@ -1,9 +1 @@
# -*- coding: utf-8 -*-
-from modelcache.utils.lazy_import import LazyImport
-scalar_manager = LazyImport(
- "scalar_manager", globals(), "modelcache.manager.scalar_data.manager"
-)
-
-
-def CacheBase(name: str, **kwargs):
- return scalar_manager.CacheBase.get(name, **kwargs)
diff --git a/modelcache/manager/scalar_data/base.py b/modelcache/manager/scalar_data/base.py
index fd8bb50..5ed898b 100644
--- a/modelcache/manager/scalar_data/base.py
+++ b/modelcache/manager/scalar_data/base.py
@@ -5,6 +5,9 @@
from enum import IntEnum
import numpy as np
+from modelcache.utils import import_sql_client
+from modelcache.utils.error import NotFoundError
+
class DataType(IntEnum):
STR = 0
@@ -94,10 +97,6 @@ class CacheStorage(metaclass=ABCMeta):
def create(self):
pass
- @abstractmethod
- def batch_insert(self, all_data: List[CacheData]):
- pass
-
@abstractmethod
def insert_query_resp(self, query_resp, **kwargs):
pass
@@ -111,7 +110,7 @@ def mark_deleted(self, keys):
pass
@abstractmethod
- def model_deleted(self, model_name):
+ def model_deleted(self, model):
pass
@abstractmethod
@@ -126,9 +125,39 @@ def get_ids(self, deleted=True):
def count(self):
pass
+ @abstractmethod
def flush(self):
pass
@abstractmethod
def close(self):
pass
+
+ @abstractmethod
+ def batch_insert(self, all_data: List[CacheData]):
+ pass
+
+ @abstractmethod
+ def update_hit_count_by_id(self, primary_id):
+ pass
+
+ @staticmethod
+ def get(name, **kwargs):
+ if name in ["mysql", "oceanbase"]:
+ from modelcache.manager.scalar_data.sql_storage import SQLStorage
+ config = kwargs.get("config")
+ import_sql_client(name)
+ cache_base = SQLStorage(db_type=name, config=config)
+ elif name == 'sqlite':
+ SQL_URL = {"sqlite": "./sqlite.db"}
+ from modelcache.manager.scalar_data.sql_storage_sqlite import SQLStorage
+ sql_url = kwargs.get("sql_url", SQL_URL[name])
+ cache_base = SQLStorage(db_type=name, url=sql_url)
+ elif name == 'elasticsearch':
+ from modelcache.manager.scalar_data.sql_storage_es import SQLStorage
+ config = kwargs.get("config")
+ cache_base = SQLStorage(db_type=name, config=config)
+ else:
+ raise NotFoundError("cache store", name)
+ return cache_base
+
diff --git a/modelcache/manager/scalar_data/sql_storage.py b/modelcache/manager/scalar_data/sql_storage.py
index 503217a..20eba9b 100644
--- a/modelcache/manager/scalar_data/sql_storage.py
+++ b/modelcache/manager/scalar_data/sql_storage.py
@@ -1,16 +1,16 @@
# -*- coding: utf-8 -*-
-import os
-import time
+import uuid
import pymysql
import json
-import base64
+import numpy as np
from typing import List
from modelcache.manager.scalar_data.base import CacheStorage, CacheData
from DBUtils.PooledDB import PooledDB
class SQLStorage(CacheStorage):
+
def __init__(
self,
db_type: str = "mysql",
@@ -41,27 +41,60 @@ def _insert(self, data: List):
model = data[3]
answer_type = 0
embedding_data = embedding_data.tobytes()
+ is_deleted = 0
+ _id = str(uuid.uuid4())
- 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 = "modelcache_llm_answer"
+ insert_sql = f"""
+ INSERT INTO {table_name}
+ (id, question, answer, answer_type, model, embedding_data, is_deleted)
+ VALUES (%s, %s, %s, %s, %s, _binary%s, %s)
+ """
conn = self.pool.connection()
try:
with conn.cursor() as cursor:
# 执行插入数据操作
- values = (question, answer, answer_type, model, embedding_data)
+ values = (_id, question, answer, answer_type, model, embedding_data, is_deleted)
cursor.execute(insert_sql, values)
conn.commit()
- id = cursor.lastrowid
finally:
# 关闭连接，将连接返回给连接池
conn.close()
- return id
+ return _id
+
+ def batch_insert(self, all_data: List[List]):
+ table_name = "modelcache_llm_answer"
+ insert_sql = f"""
+ INSERT INTO {table_name}
+ (id, question, answer, answer_type, model, embedding_data, is_deleted)
+ VALUES (%s, %s, %s, %s, %s, %s, %s)
+ """
- def batch_insert(self, all_data: List[CacheData]):
+ values_list = []
ids = []
+
for data in all_data:
- ids.append(self._insert(data))
+ answer = data[0]
+ question = data[1]
+ embedding_data = data[2].tobytes()
+ model = data[3]
+ answer_type = 0
+ is_deleted = 0
+ _id = str(uuid.uuid4())
+ ids.append(_id)
+
+ values_list.append((
+ _id, question, answer, answer_type, model, embedding_data, is_deleted
+ ))
+
+ conn = self.pool.connection()
+ try:
+ with conn.cursor() as cursor:
+ cursor.executemany(insert_sql, values_list)
+ conn.commit()
+ finally:
+ conn.close()
+
return ids
def insert_query_resp(self, query_resp, **kwargs):
@@ -78,7 +111,11 @@ def insert_query_resp(self, query_resp, **kwargs):
hit_query = json.dumps(hit_query, ensure_ascii=False)
table_name = "modelcache_query_log"
- insert_sql = "INSERT INTO {} (error_code, error_desc, cache_hit, model, query, delta_time, hit_query, answer) VALUES (%s, %s, %s, %s, %s, %s, %s, %s)".format(table_name)
+ insert_sql = f"""
+ INSERT INTO {table_name}
+ (error_code, error_desc, cache_hit, model, query, delta_time, hit_query, answer)
+ VALUES (%s, %s, %s, %s, %s, %s, %s, %s)
+ """
conn = self.pool.connection()
try:
with conn.cursor() as cursor:
@@ -91,84 +128,158 @@ def insert_query_resp(self, query_resp, **kwargs):
conn.close()
def get_data_by_id(self, key: int):
- table_name = "cache_codegpt_answer"
- query_sql = "select question, answer, embedding_data, model from {} where id={}".format(table_name, key)
- conn_start = time.time()
+ table_name = "modelcache_llm_answer"
+ query_sql = f"""
+ SELECT answer, question, embedding_data, model
+ FROM {table_name}
+ WHERE id = %s
+ """
conn = self.pool.connection()
-
- search_start = time.time()
try:
with conn.cursor() as cursor:
# 执行数据库操作
- cursor.execute(query_sql)
+ cursor.execute(query_sql, (key,))
resp = cursor.fetchone()
finally:
# 关闭连接，将连接返回给连接池
conn.close()
if resp is not None and len(resp) == 4:
- return resp
+ # parse the numpy array from bytes and return the data
+ return resp[0], resp[1], np.frombuffer(resp[2], dtype=np.float32), resp[3]
else:
return None
def update_hit_count_by_id(self, primary_id: int):
- table_name = "cache_codegpt_answer"
- update_sql = "UPDATE {} SET hit_count = hit_count+1 WHERE id={}".format(table_name, primary_id)
+ table_name = "modelcache_llm_answer"
+ update_sql = f"""
+ UPDATE {table_name}
+ SET hit_count = hit_count+1
+ WHERE id = %s
+ """
conn = self.pool.connection()
# 使用连接执行更新数据操作
try:
with conn.cursor() as cursor:
# 执行更新数据操作
- cursor.execute(update_sql)
+ cursor.execute(update_sql,(primary_id,))
conn.commit()
finally:
# 关闭连接，将连接返回给连接池
conn.close()
def get_ids(self, deleted=True):
- pass
+ table_name = "modelcache_llm_answer"
+ state = 1 if deleted else 0
+ query_sql = f"""
+ SELECT id
+ FROM {table_name}
+ WHERE is_deleted = %s
+ """
+
+ conn = self.pool.connection()
+ try:
+ with conn.cursor() as cursor:
+ cursor.execute(query_sql, (state,))
+ ids = [row[0] for row in cursor.fetchall()]
+ finally:
+ conn.close()
+
+ return ids
def mark_deleted(self, keys):
- table_name = "cache_codegpt_answer"
- delete_sql = "Delete from {} WHERE id in ({})".format(table_name, ",".join([str(i) for i in keys]))
+ table_name = "modelcache_llm_answer"
+ placeholders = ",".join(["%s"] * len(keys))
+ mark_sql = f"""
+ UPDATE {table_name}
+ SET is_deleted=1
+ WHERE id in ({placeholders})
+ """
- # 从连接池中获取连接
conn = self.pool.connection()
try:
with conn.cursor() as cursor:
- # 执行删除数据操作
- cursor.execute(delete_sql)
+ cursor.execute(mark_sql, keys)
delete_count = cursor.rowcount
conn.commit()
finally:
- # 关闭连接，将连接返回给连接池
conn.close()
return delete_count
def model_deleted(self, model_name):
- table_name = "cache_codegpt_answer"
- delete_sql = "Delete from {} WHERE model='{}'".format(table_name, model_name)
+ table_name = "modelcache_llm_answer"
+ delete_sql = f"""
+ Delete from {table_name}
+ WHERE model = %s
+ """
+
+ table_log_name = "modelcache_query_log"
+ delete_log_sql = f"""
+ Delete from {table_log_name}
+ WHERE model = %s
+ """
+
conn = self.pool.connection()
# 使用连接执行删除数据操作
try:
with conn.cursor() as cursor:
# 执行删除数据操作
- resp = cursor.execute(delete_sql)
+ resp = cursor.execute(delete_sql, (model_name,))
conn.commit()
+ # 执行删除该模型对应日志操作 resp_log行数不返回
+ resp_log = cursor.execute(delete_log_sql, (model_name,))
+ conn.commit() # 分别提交事务
finally:
# 关闭连接，将连接返回给连接池
conn.close()
return resp
def clear_deleted_data(self):
- pass
+ table_name = "modelcache_llm_answer"
+ delete_sql = f"""
+ DELETE FROM {table_name}
+ WHERE is_deleted = 1
+ """
+
+ conn = self.pool.connection()
+ try:
+ with conn.cursor() as cursor:
+ cursor.execute(delete_sql)
+ delete_count = cursor.rowcount
+ conn.commit()
+ finally:
+ conn.close()
+
+ return delete_count
def count(self, state: int = 0, is_all: bool = False):
- pass
+ table_name = "modelcache_llm_answer"
+
+ # we're not using prepared statements here, so we need to ensure state is an integer
+ if not isinstance(state, int):
+ raise ValueError("'state' must be an integer.")
+
+ if is_all:
+ count_sql = f"SELECT COUNT(*) FROM {table_name}"
+ else:
+ count_sql = f"SELECT COUNT(*) FROM {table_name} WHERE is_deleted = {state}"
+
+ conn = self.pool.connection()
+ try:
+ with conn.cursor() as cursor:
+ cursor.execute(count_sql)
+ num = cursor.fetchone()[0]
+ finally:
+ conn.close()
+
+ return num
def close(self):
pass
def count_answers(self):
pass
+
+ def flush(self):
+ pass
\ No newline at end of file
diff --git a/modelcache/manager/scalar_data/sql_storage_es.py b/modelcache/manager/scalar_data/sql_storage_es.py
new file mode 100644
index 0000000..3787623
--- /dev/null
+++ b/modelcache/manager/scalar_data/sql_storage_es.py
@@ -0,0 +1,199 @@
+# -*- coding: utf-8 -*-
+import json
+from typing import List
+from elasticsearch import Elasticsearch, helpers
+from modelcache.manager.scalar_data.base import CacheStorage, CacheData
+import time
+from snowflake import SnowflakeGenerator
+
+
+class SQLStorage(CacheStorage):
+
+ def __init__(
+ self,
+ db_type: str = "elasticsearch",
+ config=None
+ ):
+ self.host = config.get('elasticsearch', 'host')
+ self.port = int(config.get('elasticsearch', 'port'))
+ self.client = Elasticsearch(
+ hosts=[{"host": self.host, "port": self.port}],
+ timeout=30,
+ http_auth=('esuser', 'password')
+ )
+
+ self.log_index = "modelcache_query_log"
+ self.ans_index = "modelcache_llm_answer"
+ self.create()
+ self.instance_id = 1 # 雪花算法使用的机器id 使用同一套数据库的分布式系统需要配置不同id
+ # 生成雪花id
+ self.snowflake_id = SnowflakeGenerator(self.instance_id)
+
+ def create(self):
+ answer_index_body = {
+ "mappings": {
+ "properties": {
+ "gmt_create": {"type": "date", "format": "strict_date_optional_time||epoch_millis"},
+ "gmt_modified": {"type": "date", "format": "strict_date_optional_time||epoch_millis"},
+ "question": {"type": "text"},
+ "answer": {"type": "text"},
+ "answer_type": {"type": "integer"},
+ "hit_count": {"type": "integer"},
+ "model": {"type": "keyword"},
+ "embedding_data": {"type": "binary"},
+ "is_deleted": {"type": "integer"},
+ }
+ }
+ }
+
+ log_index_body = {
+ "mappings": {
+ "properties": {
+ "gmt_create": {"type": "date", "format": "strict_date_optional_time||epoch_millis"},
+ "gmt_modified": {"type": "date", "format": "strict_date_optional_time||epoch_millis"},
+ "error_code": {"type": "integer"},
+ "error_desc": {"type": "text"},
+ "cache_hit": {"type": "keyword"},
+ "delta_time": {"type": "float"},
+ "model": {"type": "keyword"},
+ "query": {"type": "text"},
+ "hit_query": {"type": "text"},
+ "answer": {"type": "text"}
+ }
+ }
+ }
+
+ if not self.client.indices.exists(index=self.ans_index):
+ self.client.indices.create(index=self.ans_index, body=answer_index_body)
+
+ if not self.client.indices.exists(index=self.log_index):
+ self.client.indices.create(index=self.log_index, body=log_index_body)
+
+ def _insert(self, data: List) -> str or None:
+ doc = {
+ "answer": data[0],
+ "question": data[1],
+ "embedding_data": data[2].tolist() if hasattr(data[2], "tolist") else data[2],
+ "model": data[3],
+ "answer_type": 0,
+ "hit_count": 0,
+ "is_deleted": 0
+ }
+
+ try:
+
+ response = self.client.index(
+ index=self.ans_index,
+ id=next(self.snowflake_id),
+ body=doc,
+ )
+ return int(response['_id'])
+ except Exception as e:
+
+ print(f"Failed to insert document: {e}")
+ return None
+
+ def batch_insert(self, all_data: List[List]) -> List[str]:
+ successful_ids = []
+ for data in all_data:
+ _id = self._insert(data)
+ if _id is not None:
+ successful_ids.append(_id)
+ self.client.indices.refresh(index=self.ans_index) # 批量插入后手动刷新
+
+ return successful_ids
+
+ def insert_query_resp(self, query_resp, **kwargs):
+ doc = {
+ "error_code": query_resp.get('errorCode'),
+ "error_desc": query_resp.get('errorDesc'),
+ "cache_hit": query_resp.get('cacheHit'),
+ "model": kwargs.get('model'),
+ "query": kwargs.get('query'),
+ "delta_time": kwargs.get('delta_time'),
+ "hit_query": json.dumps(query_resp.get('hit_query'), ensure_ascii=False) if isinstance(
+ query_resp.get('hit_query'), list) else query_resp.get('hit_query'),
+ "answer": query_resp.get('answer'),
+ "hit_count": 0,
+ "is_deleted": 0
+
+ }
+ self.client.index(index=self.log_index, body=doc)
+
+ def get_data_by_id(self, key: int):
+ try:
+ response = self.client.get(index=self.ans_index, id=key, _source=['question', 'answer', 'embedding_data', 'model'])
+ source = response["_source"]
+ result = [
+ source.get('question'),
+ source.get('answer'),
+ source.get('embedding_data'),
+ source.get('model')
+ ]
+ return result
+ except Exception as e:
+ print(e)
+
+ def update_hit_count_by_id(self, primary_id: int):
+ self.client.update(
+ index=self.ans_index,
+ id=primary_id,
+ body={"script": {"source": "ctx._source.hit_count += 1"}}
+ )
+
+ def get_ids(self, deleted=True):
+ query = {
+ "query": {
+ "term": {"is_deleted": 1 if deleted else 0}
+ }
+ }
+ response = self.client.search(index=self.ans_index, body=query)
+ return [hit["_id"] for hit in response["hits"]["hits"]]
+
+ def mark_deleted(self, keys):
+ actions = [
+ {
+ "_op_type": "update",
+ "_index": self.ans_index,
+ "_id": key,
+ "doc": {"is_deleted": 1}
+ }
+ for key in keys
+ ]
+ responses = helpers.bulk(self.client, actions)
+ return responses[0] # 返回更新的文档数
+
+ def model_deleted(self, model_name):
+ query = {
+ "query": {
+ "term": {"model": model_name}
+ }
+ }
+
+ response = self.client.delete_by_query(index=self.ans_index, body=query)
+ return response["deleted"]
+
+ def clear_deleted_data(self):
+ query = {
+ "query": {
+ "term": {"is_deleted": 1}
+ }
+ }
+ response = self.client.delete_by_query(index=self.ans_index, body=query)
+ return response["deleted"]
+
+ def count(self, state: int = 0, is_all: bool = False):
+ query = {"query": {"match_all": {}}} if is_all else {"query": {"term": {"is_deleted": state}}}
+ response = self.client.count(index=self.ans_index, body=query)
+ return response["count"]
+
+ def close(self):
+ self.client.close()
+
+ def count_answers(self):
+ query = {"query": {"match_all": {}}}
+ response = self.client.count(index=self.ans_index, body=query)
+ return response["count"]
+
+ def flush(self):
+ pass
diff --git a/modelcache/manager/scalar_data/sql_storage_sqlite.py b/modelcache/manager/scalar_data/sql_storage_sqlite.py
index 495fbf7..3b69979 100644
--- a/modelcache/manager/scalar_data/sql_storage_sqlite.py
+++ b/modelcache/manager/scalar_data/sql_storage_sqlite.py
@@ -1,74 +1,10 @@
# -*- coding: utf-8 -*-
-import os
-import time
-
-import pymysql
import json
-import base64
from typing import List
from modelcache.manager.scalar_data.base import CacheStorage, CacheData
import sqlite3
-# def insert_single_data(conn, sql, data_tuple):
-# cur = conn.cursor()
-# try:
-# cur.execute(sql, data_tuple)
-# conn.commit()
-# id = cur.lastrowid
-# # print('id: {}'.format(id))
-# return id
-# except Exception as e:
-# print(e)
-# conn.rollback()
-# if cur:
-# cur.close()
-#
-#
-# def excute_sql(conn, sql):
-# cur = conn.cursor()
-# try:
-# cur.execute(sql)
-# conn.commit()
-# except Exception as e:
-# print(e)
-# conn.rollback()
-# if cur:
-# cur.close()
-#
-#
-# def excute_delete_sql(conn, sql):
-# cur = conn.cursor()
-# try:
-# cur.execute(sql)
-# row_count = cur.rowcount
-# conn.commit()
-# except Exception as e:
-# print(e)
-# conn.rollback()
-# if cur:
-# cur.close()
-# return row_count
-#
-#
-# def query_fetch_one_data(conn, sql):
-# cursor = conn.cursor()
-# try:
-# cursor.execute(sql)
-# except Exception as e:
-# print(e)
-# conn.rollback()
-# rows = cursor.fetchone()
-# if cursor:
-# cursor.close()
-# return rows
-#
-#
-# def close(conn):
-# if conn:
-# conn.close()
-
-
class SQLStorage(CacheStorage):
def __init__(
self,
@@ -81,19 +17,6 @@ def __init__(
self.create()
def create(self):
- # answer_table_sql = """CREATE TABLE IF NOT EXISTS `modelcache_llm_answer` (
- # `id` bigint(20) NOT NULL AUTO_INCREMENT comment '主键',
- # `gmt_create` timestamp NOT NULL DEFAULT CURRENT_TIMESTAMP comment '创建时间',
- # `gmt_modified` timestamp NOT NULL DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP comment '修改时间',
- # `question` text NOT NULL comment 'question',
- # `answer` text NOT NULL comment 'answer',
- # `answer_type` int(11) NOT NULL comment 'answer_type',
- # `hit_count` int(11) NOT NULL DEFAULT '0' comment 'hit_count',
- # `model` varchar(1000) NOT NULL comment 'model',
- # `embedding_data` blob NOT NULL comment 'embedding_data',
- # PRIMARY KEY(`id`)
- # ) AUTO_INCREMENT = 1 DEFAULT CHARSET = utf8mb4 COMMENT = 'modelcache_llm_answer';
- # """
answer_table_sql = """CREATE TABLE IF NOT EXISTS modelcache_llm_answer (
id INTEGER PRIMARY KEY AUTOINCREMENT,
gmt_create TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP,
@@ -107,21 +30,6 @@ def create(self):
);
"""
- # log_table_sql = """CREATE TABLE IF NOT EXISTS `modelcache_query_log` (
- # `id` bigint(20) NOT NULL AUTO_INCREMENT comment '主键',
- # `gmt_create` timestamp NOT NULL DEFAULT CURRENT_TIMESTAMP comment '创建时间',
- # `gmt_modified` timestamp NOT NULL DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP comment '修改时间',
- # `error_code` int(11) NOT NULL comment 'errorCode',
- # `error_desc` varchar(1000) NOT NULL comment 'errorDesc',
- # `cache_hit` varchar(100) NOT NULL comment 'cacheHit',
- # `delta_time` float NOT NULL comment 'delta_time',
- # `model` varchar(1000) NOT NULL comment 'model',
- # `query` text NOT NULL comment 'query',
- # `hit_query` text NOT NULL comment 'hitQuery',
- # `answer` text NOT NULL comment 'answer',
- # PRIMARY KEY(`id`)
- # ) AUTO_INCREMENT = 1 DEFAULT CHARSET = utf8mb4 COMMENT = 'modelcache_query_log';
- # """
log_table_sql = """CREATE TABLE IF NOT EXISTS modelcache_query_log (
id INTEGER PRIMARY KEY AUTOINCREMENT,
gmt_create TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP,
@@ -192,8 +100,7 @@ def insert_query_resp(self, query_resp, **kwargs):
hit_query = json.dumps(hit_query, ensure_ascii=False)
table_name = "modelcache_query_log"
- insert_sql = "INSERT INTO {} (error_code, error_desc, cache_hit, model, query, delta_time, hit_query, answer) VALUES (%s, %s, %s, %s, %s, %s, %s, %s)".format(table_name)
-
+ insert_sql = "INSERT INTO {} (error_code, error_desc, cache_hit, model, query, delta_time, hit_query, answer) VALUES (?, ?, ?, ?, ?, ?, ?, ?)".format(table_name)
conn = sqlite3.connect(self._url)
try:
cursor = conn.cursor()
@@ -259,17 +166,27 @@ def mark_deleted(self, keys):
def model_deleted(self, model_name):
table_name = "modelcache_llm_answer"
- delete_sql = "Delete from {} WHERE model='{}'".format(table_name, model_name)
+ delete_sql = "Delete from {} WHERE model=?".format(table_name)
+
+ table_log_name = "modelcache_query_log"
+ delete_log_sql = "Delete from {} WHERE model=?".format(table_log_name)
conn = sqlite3.connect(self._url)
try:
cursor = conn.cursor()
- resp = cursor.execute(delete_sql)
+ cursor.execute(delete_sql, (model_name,))
+ conn.commit()
+ # get delete rows
+ deleted_rows_count = cursor.rowcount
+
+ cursor.execute(delete_log_sql, (model_name,))
conn.commit()
cursor.close()
- conn.close()
+ except sqlite3.Error as e:
+ print(f"SQLite error: {e}")
+ deleted_rows_count = 0 # if except, return 0
finally:
conn.close()
- return resp
+ return deleted_rows_count
def clear_deleted_data(self):
pass
@@ -282,3 +199,6 @@ def close(self):
def count_answers(self):
pass
+
+ def flush(self):
+ pass
diff --git a/modelcache/manager/vector_data/__init__.py b/modelcache/manager/vector_data/__init__.py
index fdb98cb..40a96af 100644
--- a/modelcache/manager/vector_data/__init__.py
+++ b/modelcache/manager/vector_data/__init__.py
@@ -1,10 +1 @@
# -*- coding: utf-8 -*-
-from modelcache.utils.lazy_import import LazyImport
-
-vector_manager = LazyImport(
- "vector_manager", globals(), "modelcache.manager.vector_data.manager"
-)
-
-
-def VectorBase(name: str, **kwargs):
- return vector_manager.VectorBase.get(name, **kwargs)
diff --git a/modelcache/manager/vector_data/base.py b/modelcache/manager/vector_data/base.py
index 7bdf12a..689f833 100644
--- a/modelcache/manager/vector_data/base.py
+++ b/modelcache/manager/vector_data/base.py
@@ -4,6 +4,19 @@
from typing import List
from dataclasses import dataclass
+from modelcache.embedding import MetricType
+from modelcache.utils.error import ParamError, NotFoundError
+
+TOP_K = 1
+FAISS_INDEX_PATH = "faiss.index"
+DIMENSION = 0
+MILVUS_HOST = "localhost"
+MILVUS_PORT = 19530
+MILVUS_USER = ""
+MILVUS_PSW = ""
+MILVUS_SECURE = False
+
+COLLECTION_NAME = "modelcache"
@dataclass
class VectorData:
@@ -11,8 +24,7 @@ class VectorData:
data: np.ndarray
-class VectorBase(ABC):
- """VectorBase: base vector store interface"""
+class VectorStorage(ABC):
@abstractmethod
def mul_add(self, datas: List[VectorData], model=None):
@@ -34,8 +46,108 @@ def delete(self, ids) -> bool:
def rebuild_col(self, model):
pass
+ @abstractmethod
def flush(self):
pass
+ @abstractmethod
def close(self):
pass
+
+ @staticmethod
+ def get(name, **kwargs):
+ top_k = kwargs.get("top_k", TOP_K)
+ if name == "milvus":
+ from modelcache.manager.vector_data.milvus import Milvus
+ dimension = kwargs.get("dimension", DIMENSION)
+ milvus_config = kwargs.get("config")
+ check_dimension(dimension)
+ host = milvus_config.get('milvus', 'host')
+ port = milvus_config.get('milvus', 'port')
+ user = milvus_config.get('milvus', 'user')
+ password = milvus_config.get('milvus', 'password')
+
+ metric_type = kwargs.get("metric_type",MetricType.COSINE)
+ secure = kwargs.get("secure", MILVUS_SECURE)
+ collection_name = kwargs.get("collection_name", COLLECTION_NAME)
+ index_params = kwargs.get("index_params", None)
+ search_params = kwargs.get("search_params", None)
+ local_mode = kwargs.get("local_mode", False)
+ local_data = kwargs.get("local_data", "./milvus_data")
+ vector_base = Milvus(
+ host=host,
+ port=port,
+ user=user,
+ password=password,
+ secure=secure,
+ collection_name=collection_name,
+ dimension=dimension,
+ top_k=top_k,
+ index_params=index_params,
+ search_params=search_params,
+ local_mode=local_mode,
+ local_data=local_data,
+ metric_type=metric_type
+ )
+ elif name == "redis":
+ from modelcache.manager.vector_data.redis import RedisVectorStore
+ dimension = kwargs.get("dimension", DIMENSION)
+ check_dimension(dimension)
+
+ redis_config = kwargs.get("config")
+ host = redis_config.get('redis', 'host')
+ port = redis_config.get('redis', 'port')
+ user = redis_config.get('redis', 'user')
+ password = redis_config.get('redis', 'password')
+ namespace = kwargs.get("namespace", "")
+ # collection_name = kwargs.get("collection_name", COLLECTION_NAME)
+
+ vector_base = RedisVectorStore(
+ host=host,
+ port=port,
+ username=user,
+ password=password,
+ namespace=namespace,
+ top_k=top_k,
+ dimension=dimension,
+ )
+ elif name == "faiss":
+ from modelcache.manager.vector_data.faiss import Faiss
+
+ dimension = kwargs.get("dimension", DIMENSION)
+ index_path = kwargs.pop("index_path", FAISS_INDEX_PATH)
+ check_dimension(dimension)
+ vector_base = Faiss(
+ index_file_path=index_path, dimension=dimension, top_k=top_k
+ )
+ elif name == "chromadb":
+ from modelcache.manager.vector_data.chroma import Chromadb
+
+ chromadb_config = kwargs.get("config", None)
+ persist_directory = chromadb_config.get('chromadb','persist_directory')
+
+ vector_base = Chromadb(
+ persist_directory=persist_directory,
+ top_k=top_k,
+ )
+ elif name == "hnswlib":
+ from modelcache.manager.vector_data.hnswlib_store import Hnswlib
+
+ dimension = kwargs.get("dimension", DIMENSION)
+ index_path = kwargs.pop("index_path", "./hnswlib_index.bin")
+ max_elements = kwargs.pop("max_elements", 100000)
+ VectorStorage.check_dimension(dimension)
+ vector_base = Hnswlib(
+ index_file_path=index_path, dimension=dimension,
+ top_k=top_k, max_elements=max_elements
+ )
+ else:
+ raise NotFoundError("vector store", name)
+ return vector_base
+
+
+def check_dimension(dimension):
+ if dimension <= 0:
+ raise ParamError(f"the dimension should be greater than zero, current value: {dimension}.")
+
+
diff --git a/modelcache/manager/vector_data/chroma.py b/modelcache/manager/vector_data/chroma.py
new file mode 100644
index 0000000..b92ef3c
--- /dev/null
+++ b/modelcache/manager/vector_data/chroma.py
@@ -0,0 +1,92 @@
+from typing import List
+
+import numpy as np
+import logging
+from modelcache.manager.vector_data.base import VectorStorage, VectorData
+from modelcache.utils import import_chromadb, import_torch
+
+import_torch()
+import_chromadb()
+
+import chromadb
+
+
+class Chromadb(VectorStorage):
+
+ def __init__(
+ self,
+ persist_directory="./chromadb",
+ top_k: int = 1,
+ ):
+ self.collection_name = "modelcache"
+ self.top_k = top_k
+
+ self._client = chromadb.PersistentClient(path=persist_directory)
+ self._collection = None
+
+ def mul_add(self, datas: List[VectorData], model=None):
+ collection_name_model = self.collection_name + '_' + model
+ self._collection = self._client.get_or_create_collection(name=collection_name_model)
+
+ data_array, id_array = map(list, zip(*((data.data.tolist(), str(data.id)) for data in datas)))
+ self._collection.add(embeddings=data_array, ids=id_array)
+
+ def search(self, data: np.ndarray, top_k: int = -1, model=None):
+ collection_name_model = self.collection_name + '_' + model
+ self._collection = self._client.get_or_create_collection(name=collection_name_model)
+
+ if self._collection.count() == 0:
+ return []
+ if top_k == -1:
+ top_k = self.top_k
+ results = self._collection.query(
+ query_embeddings=[data.tolist()],
+ n_results=top_k,
+ include=["distances"],
+ )
+ return list(zip(results["distances"][0], [int(x) for x in results["ids"][0]]))
+
+ def rebuild(self, ids=None):
+ pass
+
+ def delete(self, ids, model=None):
+ try:
+ collection_name_model = self.collection_name + '_' + model
+ self._collection = self._client.get_or_create_collection(name=collection_name_model)
+ # 查询集合中实际存在的 ID
+ ids_str = [str(x) for x in ids]
+ existing_ids = set(self._collection.get(ids=ids_str).ids)
+
+ # 删除存在的 ID
+ if existing_ids:
+ self._collection.delete(list(existing_ids))
+
+ # 返回实际删除的条目数量
+ return len(existing_ids)
+
+ except Exception as e:
+ logging.error('Error during deletion: {}'.format(e))
+ raise ValueError(str(e))
+
+ def rebuild_col(self, model):
+ collection_name_model = self.collection_name + '_' + model
+
+ # 检查集合是否存在，如果存在则删除
+ collections = self._client.list_collections()
+ if any(col.name == collection_name_model for col in collections):
+ self._client.delete_collection(collection_name_model)
+ else:
+ return 'model collection not found, please check!'
+
+ try:
+ self._client.create_collection(collection_name_model)
+ except Exception as e:
+ logging.info(f'rebuild_collection: {e}')
+ raise ValueError(str(e))
+
+ def flush(self):
+ # chroma无flush方法
+ pass
+
+ def close(self):
+ pass
diff --git a/modelcache/manager/vector_data/faiss.py b/modelcache/manager/vector_data/faiss.py
index f035ded..77e7617 100644
--- a/modelcache/manager/vector_data/faiss.py
+++ b/modelcache/manager/vector_data/faiss.py
@@ -2,13 +2,13 @@
import os
from typing import List
import numpy as np
-from modelcache.manager.vector_data.base import VectorBase, VectorData
+from modelcache.manager.vector_data.base import VectorStorage, VectorData
from modelcache.utils import import_faiss
import_faiss()
import faiss # pylint: disable=C0413
-class Faiss(VectorBase):
+class Faiss(VectorStorage):
def __init__(self, index_file_path, dimension, top_k):
self._index_file_path = index_file_path
self._dimension = dimension
@@ -34,7 +34,10 @@ def search(self, data: np.ndarray, top_k: int = -1, model=None):
return list(zip(dist[0], ids))
def rebuild_col(self, ids=None):
- return True
+ try:
+ self._index.reset()
+ except Exception as e:
+ return f"An error occurred during index rebuild: {e}"
def rebuild(self, ids=None):
return True
diff --git a/modelcache/manager/vector_data/milvus.py b/modelcache/manager/vector_data/milvus.py
index 50d6ab1..f5c7682 100644
--- a/modelcache/manager/vector_data/milvus.py
+++ b/modelcache/manager/vector_data/milvus.py
@@ -3,9 +3,11 @@
from typing import List
from uuid import uuid4
import numpy as np
+
+from modelcache.embedding import MetricType
from modelcache.utils import import_pymilvus
from modelcache.utils.log import modelcache_log
-from modelcache.manager.vector_data.base import VectorBase, VectorData
+from modelcache.manager.vector_data.base import VectorStorage, VectorData
import_pymilvus()
@@ -21,19 +23,7 @@
)
-class Milvus(VectorBase):
- SEARCH_PARAM = {
- "IVF_FLAT": {"metric_type": "L2", "params": {"nprobe": 10}},
- "IVF_SQ8": {"metric_type": "L2", "params": {"nprobe": 10}},
- "IVF_PQ": {"metric_type": "L2", "params": {"nprobe": 10}},
- "HNSW": {"metric_type": "L2", "params": {"ef": 10}},
- "RHNSW_FLAT": {"metric_type": "L2", "params": {"ef": 10}},
- "RHNSW_SQ": {"metric_type": "L2", "params": {"ef": 10}},
- "RHNSW_PQ": {"metric_type": "L2", "params": {"ef": 10}},
- "IVF_HNSW": {"metric_type": "L2", "params": {"nprobe": 10, "ef": 10}},
- "ANNOY": {"metric_type": "L2", "params": {"search_k": 10}},
- "AUTOINDEX": {"metric_type": "L2", "params": {}},
- }
+class Milvus(VectorStorage):
def __init__(
self,
@@ -48,7 +38,8 @@ def __init__(
index_params: dict = None,
search_params: dict = None,
local_mode: bool = False,
- local_data: str = "./milvus_data"
+ local_data: str = "./milvus_data",
+ metric_type: MetricType = MetricType.COSINE,
):
if dimension <= 0:
raise ValueError(
@@ -58,14 +49,29 @@ def __init__(
self._local_data = local_data
self.dimension = dimension
self.top_k = top_k
- self.index_params = index_params
if self._local_mode:
self._create_local(port, local_data)
self._connect(host, port, user, password, secure)
self.collection_name = collection_name
- self.search_params = (
- search_params or self.SEARCH_PARAM[self.index_params["index_type"]]
- )
+ self.search_params = {
+ "IVF_FLAT": {"metric_type": metric_type.value, "params": {"nprobe": 10}},
+ "IVF_SQ8": {"metric_type": metric_type.value, "params": {"nprobe": 10}},
+ "IVF_PQ": {"metric_type": metric_type.value, "params": {"nprobe": 10}},
+ "HNSW": {"metric_type": metric_type.value, "params": {"ef": 10}},
+ "RHNSW_FLAT": {"metric_type": metric_type.value, "params": {"ef": 10}},
+ "RHNSW_SQ": {"metric_type": metric_type.value, "params": {"ef": 10}},
+ "RHNSW_PQ": {"metric_type": metric_type.value, "params": {"ef": 10}},
+ "IVF_HNSW": {"metric_type": metric_type.value, "params": {"nprobe": 10, "ef": 10}},
+ "ANNOY": {"metric_type": metric_type.value, "params": {"search_k": 10}},
+ "AUTOINDEX": {"metric_type": metric_type.value, "params": {}},
+ }
+ self.index_params ={
+ "metric_type": metric_type.value,
+ "index_type": "HNSW",
+ "params": {"M": 16, "efConstruction": 64},
+ }
+ self.collections = dict()
+
def _connect(self, host, port, user, password, secure):
try:
@@ -87,12 +93,14 @@ def _connect(self, host, port, user, password, secure):
timeout=10
)
+
def _create_collection(self, collection_name):
if not utility.has_collection(collection_name, using=self.alias):
schema = [
FieldSchema(
name="id",
- dtype=DataType.INT64,
+ dtype=DataType.VARCHAR,
+ max_length=36,
is_primary=True,
auto_id=False,
),
@@ -101,7 +109,8 @@ def _create_collection(self, collection_name):
),
]
schema = CollectionSchema(schema)
- self.col = Collection(
+
+ new_collection = Collection(
collection_name,
schema=schema,
consistency_level="Session",
@@ -109,46 +118,48 @@ def _create_collection(self, collection_name):
)
else:
modelcache_log.warning("The %s collection already exists, and it will be used directly.", collection_name)
- self.col = Collection(
+ new_collection = Collection(
collection_name, consistency_level="Session", using=self.alias
)
- if len(self.col.indexes) == 0:
+ self.collections[collection_name] = new_collection
+
+ if len(new_collection.indexes) == 0:
try:
modelcache_log.info("Attempting creation of Milvus index.")
- self.col.create_index("embedding", index_params=self.index_params)
+ new_collection.create_index("embedding", index_params=self.index_params)
modelcache_log.info("Creation of Milvus index successful.")
except MilvusException as e:
modelcache_log.warning("Error with building index: %s, and attempting creation of default index.", e)
i_p = {"metric_type": "L2", "index_type": "AUTOINDEX", "params": {}}
- self.col.create_index("embedding", index_params=i_p)
+ new_collection.create_index("embedding", index_params=i_p)
self.index_params = i_p
else:
- self.index_params = self.col.indexes[0].to_dict()["index_param"]
+ self.index_params = new_collection.indexes[0].to_dict()["index_param"]
+
+ new_collection.load()
- self.col.load()
def _get_collection(self, collection_name):
- self.col = Collection(
- collection_name, consistency_level="Session", using=self.alias
- )
- self.col.load()
+ if collection_name not in self.collections:
+ self._create_collection(collection_name)
+ return self.collections[collection_name]
def mul_add(self, datas: List[VectorData], model=None):
collection_name_model = self.collection_name + '_' + model
- self._create_collection(collection_name_model)
-
+ col = self._get_collection(collection_name_model)
data_array, id_array = map(list, zip(*((data.data, data.id) for data in datas)))
np_data = np.array(data_array).astype("float32")
entities = [id_array, np_data]
- self.col.insert(entities)
+ col.insert(entities)
+
def search(self, data: np.ndarray, top_k: int = -1, model=None):
if top_k == -1:
top_k = self.top_k
collection_name_model = self.collection_name + '_' + model
- self._create_collection(collection_name_model)
- search_result = self.col.search(
+ col = self._get_collection(collection_name_model)
+ search_result = col.search(
data=data.reshape(1, -1).tolist(),
anns_field="embedding",
param=self.search_params,
@@ -156,12 +167,13 @@ def search(self, data: np.ndarray, top_k: int = -1, model=None):
)
return list(zip(search_result[0].distances, search_result[0].ids))
+
def delete(self, ids, model=None):
collection_name_model = self.collection_name + '_' + model
- self._get_collection(collection_name_model)
+ col = self._get_collection(collection_name_model)
- del_ids = ",".join([str(x) for x in ids])
- resp = self.col.delete(f"id in [{del_ids}]")
+ del_ids = ",".join([f'"{x}"' for x in ids])
+ resp = col.delete(f"id in [{del_ids}]")
delete_count = resp.delete_count
return delete_count
@@ -178,10 +190,12 @@ def rebuild_col(self, model):
logging.info('create_collection: {}'.format(e))
def rebuild(self, ids=None): # pylint: disable=unused-argument
- self.col.compact()
+ for col in self.collections.values():
+ col.compact()
def flush(self):
- self.col.flush(_async=True)
+ for col in self.collections.values():
+ col.flush(_async=True)
def close(self):
self.flush()
diff --git a/modelcache/manager/vector_data/redis.py b/modelcache/manager/vector_data/redis.py
new file mode 100644
index 0000000..6e51511
--- /dev/null
+++ b/modelcache/manager/vector_data/redis.py
@@ -0,0 +1,139 @@
+# -*- coding: utf-8 -*-
+from typing import List
+import numpy as np
+from redis.commands.search.indexDefinition import IndexDefinition, IndexType
+from redis.commands.search.query import Query
+from redis.commands.search.field import TagField, VectorField, NumericField
+from redis.client import Redis
+
+from modelcache.manager.vector_data.base import VectorStorage, VectorData
+from modelcache.utils import import_redis
+from modelcache.utils.log import modelcache_log
+from modelcache.utils.index_util import get_index_name
+from modelcache.utils.index_util import get_index_prefix
+import_redis()
+
+
+class RedisVectorStore(VectorStorage):
+ def __init__(
+ self,
+ host: str = "localhost",
+ port: str = "6379",
+ username: str = "",
+ password: str = "",
+ dimension: int = 0,
+ top_k: int = 1,
+ namespace: str = "",
+ ):
+ if dimension <= 0:
+ raise ValueError(
+ f"invalid `dim` param: {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.namespace = namespace
+ self.doc_prefix = f"{self.namespace}doc:"
+
+ def _check_index_exists(self, index_name: str) -> bool:
+ """Check if Redis index exists."""
+ try:
+ self._client.ft(index_name).info()
+ except:
+ modelcache_log.info("Index does not exist")
+ return False
+ modelcache_log.info("Index already exists")
+ return True
+
+ def create_index(self, index_name, index_prefix):
+ dimension = self.dimension
+ if self._check_index_exists(index_name):
+ modelcache_log.info(
+ "The %s already exists, and it will be used directly", index_name
+ )
+ return 'already_exists'
+ else:
+ id_field_name = "data_id"
+ embedding_field_name = "data_vector"
+
+ id = NumericField(name=id_field_name)
+ embedding = VectorField(embedding_field_name,
+ "HNSW", {
+ "TYPE": "FLOAT32",
+ "DIM": dimension,
+ "DISTANCE_METRIC": "L2",
+ "INITIAL_CAP": 1000,
+ }
+ )
+ fields = [id, embedding]
+ definition = IndexDefinition(prefix=[index_prefix], index_type=IndexType.HASH)
+
+ # create Index
+ self._client.ft(index_name).create_index(
+ fields=fields, definition=definition
+ )
+ return 'create_success'
+
+ def mul_add(self, datas: List[VectorData], model=None):
+ # pipe = self._client.pipeline()
+ for data in datas:
+ id: int = data.id
+ embedding = data.data.astype(np.float32).tobytes()
+ 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)
+
+ def search(self, data: np.ndarray, top_k: int = -1, model=None):
+ index_name = get_index_name(model)
+ id_field_name = "data_id"
+ embedding_field_name = "data_vector"
+ base_query = f'*=>[KNN 2 @{embedding_field_name} $vector AS distance]'
+ query = (
+ Query(base_query)
+ .sort_by("distance")
+ .return_fields(id_field_name, "distance")
+ .dialect(2)
+ )
+
+ query_params = {"vector": data.astype(np.float32).tobytes()}
+ results = (
+ self._client.ft(index_name)
+ .search(query, query_params=query_params)
+ .docs
+ )
+ return [(float(result.distance), int(getattr(result, id_field_name))) for result in results]
+
+ def rebuild(self, ids=None) -> bool:
+ pass
+
+ def rebuild_col(self, model):
+ index_name_model = get_index_name(model)
+ if self._check_index_exists(index_name_model):
+ try:
+ self._client.ft(index_name_model).dropindex(delete_documents=True)
+ except Exception as e:
+ raise ValueError(str(e))
+ try:
+ index_prefix = get_index_prefix(model)
+ self.create_index(index_name_model, index_prefix)
+ except Exception as e:
+ raise ValueError(str(e))
+ # return 'rebuild success'
+
+ def delete(self, ids) -> None:
+ pipe = self._client.pipeline()
+ for data_id in ids:
+ pipe.delete(f"{self.doc_prefix}{data_id}")
+ pipe.execute()
+
+ def create(self, model=None):
+ index_name = get_index_name(model)
+ index_prefix = get_index_prefix(model)
+ return self.create_index(index_name, index_prefix)
+
+ def get_index_by_name(self, index_name):
+ pass
diff --git a/modelcache/processor/pre.py b/modelcache/processor/pre.py
index 5875294..6a9d0f6 100644
--- a/modelcache/processor/pre.py
+++ b/modelcache/processor/pre.py
@@ -4,7 +4,7 @@
def insert_last_content(data: Dict[str, Any], **_: Dict[str, Any]) -> Any:
- return data.get("chat_info")[-1]["query"]
+ return data.get("query")[-1]["content"]
def query_last_content(data: Dict[str, Any], **_: Dict[str, Any]) -> Any:
@@ -64,9 +64,14 @@ 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)
+def query_with_role(data: Dict[str, Any], **_: Dict[str, Any]) -> Any:
+ query = data["query"][-1]
+ content = query["content"]
+ role = query["role"]
+ return role+": "+content
def multi_splicing(data_list) -> Any:
result_str = ""
diff --git a/modelcache/similarity_evaluation/__init__.py b/modelcache/similarity_evaluation/__init__.py
index 22dff8f..633f866 100644
--- a/modelcache/similarity_evaluation/__init__.py
+++ b/modelcache/similarity_evaluation/__init__.py
@@ -1,11 +1,2 @@
# -*- coding: utf-8 -*-
-from modelcache.similarity_evaluation.similarity_evaluation import SimilarityEvaluation
-from modelcache.utils.lazy_import import LazyImport
-exact_match = LazyImport(
- "exact_match", globals(), "modelcache.similarity_evaluation.exact_match"
-)
-
-
-def ExactMatchEvaluation():
- return exact_match.ExactMatchEvaluation()
diff --git a/modelcache/similarity_evaluation/similarity_evaluation.py b/modelcache/similarity_evaluation/base.py
similarity index 100%
rename from modelcache/similarity_evaluation/similarity_evaluation.py
rename to modelcache/similarity_evaluation/base.py
diff --git a/modelcache/similarity_evaluation/distance.py b/modelcache/similarity_evaluation/distance.py
index 44ca595..281bcd0 100644
--- a/modelcache/similarity_evaluation/distance.py
+++ b/modelcache/similarity_evaluation/distance.py
@@ -1,7 +1,6 @@
# -*- coding: utf-8 -*-
from typing import Tuple, Dict, Any
-from modelcache.similarity_evaluation import SimilarityEvaluation
-
+from modelcache.similarity_evaluation.base import SimilarityEvaluation
class SearchDistanceEvaluation(SimilarityEvaluation):
def __init__(self, max_distance=4.0, positive=False):
diff --git a/modelcache/similarity_evaluation/exact_match.py b/modelcache/similarity_evaluation/exact_match.py
index 553bd59..e207387 100644
--- a/modelcache/similarity_evaluation/exact_match.py
+++ b/modelcache/similarity_evaluation/exact_match.py
@@ -1,6 +1,6 @@
# -*- coding: utf-8 -*-
from typing import Tuple, Dict, Any
-from modelcache.similarity_evaluation.similarity_evaluation import SimilarityEvaluation
+from modelcache.similarity_evaluation.base import SimilarityEvaluation
class ExactMatchEvaluation(SimilarityEvaluation):
diff --git a/modelcache/utils/__init__.py b/modelcache/utils/__init__.py
index 425b926..15d90ec 100644
--- a/modelcache/utils/__init__.py
+++ b/modelcache/utils/__init__.py
@@ -69,3 +69,11 @@ def import_timm():
def import_pillow():
_check_library("PIL", package="pillow")
+
+
+def import_redis():
+ _check_library("redis")
+
+
+def import_chromadb():
+ _check_library("chromadb", package="chromadb")
\ No newline at end of file
diff --git a/modelcache/utils/index_util.py b/modelcache/utils/index_util.py
new file mode 100644
index 0000000..be6e856
--- /dev/null
+++ b/modelcache/utils/index_util.py
@@ -0,0 +1,9 @@
+# -*- coding: utf-8 -*-
+
+
+def get_index_name(model):
+ return 'modelcache' + '_' + model
+
+
+def get_index_prefix(model):
+ return 'prefix' + '_' + model
diff --git a/modelcache/utils/time.py b/modelcache/utils/time.py
index 7074b58..569af99 100644
--- a/modelcache/utils/time.py
+++ b/modelcache/utils/time.py
@@ -1,15 +1,14 @@
# -*- coding: utf-8 -*-
import time
-from modelcache import cache
-
-def time_cal(func, func_name=None, report_func=None):
+def time_cal(func, func_name=None, report_func=None, **kwargs):
+ cache = kwargs.pop("cache_obj")
def inner(*args, **kwargs):
time_start = time.time()
res = func(*args, **kwargs)
delta_time = time.time() - time_start
- if cache.config.log_time_func:
- cache.config.log_time_func(
+ if cache.log_time_func:
+ cache.log_time_func(
func.__name__ if func_name is None else func_name, delta_time
)
if report_func is not None:
diff --git a/modelcache_mm/__init__.py b/modelcache_mm/__init__.py
new file mode 100644
index 0000000..17c93fb
--- /dev/null
+++ b/modelcache_mm/__init__.py
@@ -0,0 +1,5 @@
+# -*- coding: utf-8 -*-
+from modelcache_mm.core import Cache
+from modelcache_mm.core import cache
+from modelcache_mm.config import Config
+import modelcache_mm.adapter
diff --git a/modelcache_mm/adapter/__init__.py b/modelcache_mm/adapter/__init__.py
new file mode 100644
index 0000000..40a96af
--- /dev/null
+++ b/modelcache_mm/adapter/__init__.py
@@ -0,0 +1 @@
+# -*- coding: utf-8 -*-
diff --git a/modelcache_mm/adapter/adapter.py b/modelcache_mm/adapter/adapter.py
new file mode 100644
index 0000000..44c87a3
--- /dev/null
+++ b/modelcache_mm/adapter/adapter.py
@@ -0,0 +1,64 @@
+# -*- coding: utf-8 -*-
+import logging
+
+from modelcache_mm.adapter.adapter_query import adapt_query
+from modelcache_mm.adapter.adapter_insert import adapt_insert
+from modelcache_mm.adapter.adapter_remove import adapt_remove
+from modelcache_mm.adapter.adapter_register import adapt_register
+
+
+class ChatCompletion(object):
+ """Openai ChatCompletion Wrapper"""
+ @classmethod
+ def create_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)
+ raise e
+
+ @classmethod
+ def create_insert(cls, *args, **kwargs):
+ try:
+ return adapt_insert(
+ *args,
+ **kwargs
+ )
+ except Exception as e:
+ # return str(e)
+ raise e
+
+ @classmethod
+ def create_remove(cls, *args, **kwargs):
+ try:
+ return adapt_remove(
+ *args,
+ **kwargs
+ )
+ except Exception as e:
+ raise e
+
+ @classmethod
+ def create_register(cls, *args, **kwargs):
+ try:
+ return adapt_register(
+ *args,
+ **kwargs
+ )
+ except Exception as e:
+ raise e
+
+
+def construct_resp_from_cache(return_message, return_query):
+ return {
+ "modelcache": True,
+ "hitQuery": return_query,
+ "data": return_message,
+ "errorCode": 0
+ }
diff --git a/modelcache_mm/adapter/adapter_insert.py b/modelcache_mm/adapter/adapter_insert.py
new file mode 100644
index 0000000..89933bc
--- /dev/null
+++ b/modelcache_mm/adapter/adapter_insert.py
@@ -0,0 +1,85 @@
+# -*- coding: utf-8 -*-
+import numpy as np
+from modelcache_mm import cache
+from modelcache_mm.utils.error import NotInitError
+from modelcache_mm.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_dict = 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']
+
+ 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:
+ pre_embedding_image = pre_embedding_image_url
+ 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_resp = time_cal(
+ chat_cache.embedding_func,
+ func_name="image_embedding",
+ report_func=chat_cache.report.embedding,
+ )(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_embeddings, text_embeddings))
+ 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_text,
+ pre_embedding_image_url,
+ pre_embedding_image_id,
+ llm_data,
+ embedding_data,
+ model=model,
+ mm_type=mm_type,
+ extra_param=context.get("mm_save_func", None)
+ )
+ return 'success'
diff --git a/modelcache_mm/adapter/adapter_query.py b/modelcache_mm/adapter/adapter_query.py
new file mode 100644
index 0000000..dc01e62
--- /dev/null
+++ b/modelcache_mm/adapter/adapter_query.py
@@ -0,0 +1,216 @@
+# -*- coding: utf-8 -*-
+import logging
+import numpy as np
+from modelcache_mm import cache
+from modelcache_mm.utils.error import NotInitError
+from modelcache_mm.utils.error import MultiTypeError
+from modelcache_mm.utils.time import time_cal
+
+
+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", {})
+ cache_factor = kwargs.pop("cache_factor", 1.0)
+
+ pre_embedding_data_dict = chat_cache.query_pre_embedding_func(
+ kwargs,
+ extra_param=context.get("pre_embedding_func", None),
+ prompts=chat_cache.config.prompts,
+ )
+
+ pre_embedding_text = '###'.join(pre_embedding_data_dict['text'])
+ pre_embedding_image_raw = pre_embedding_data_dict['imageRaw']
+ pre_embedding_image_url = pre_embedding_data_dict['imageUrl']
+ pre_multi_type = pre_embedding_data_dict['multiType']
+
+ # 判断逻辑
+ if pre_multi_type == 'IMG_TEXT':
+ if pre_embedding_image_raw and pre_embedding_image_url:
+ raise ValueError(
+ "Both pre_embedding_imageUrl and pre_embedding_imageRaw cannot be non-empty at the same time.")
+ if pre_embedding_image_url:
+ pre_embedding_image = pre_embedding_image_url
+ elif pre_embedding_image_raw:
+ pre_embedding_image = pre_embedding_image_raw
+ else:
+ raise ValueError(
+ "Both pre_embedding_imageUrl and pre_embedding_imageRaw are empty. Please provide at least one.")
+ data_dict = {'text': [pre_embedding_text], 'image': pre_embedding_image}
+ elif pre_multi_type == 'TEXT':
+ data_dict = {'text': [pre_embedding_text], 'image': None}
+ else:
+ raise MultiTypeError
+
+ # embedding_data = None
+ # mm_type = None
+ if pre_multi_type == 'IMG_TEXT':
+ embedding_data_resp = time_cal(
+ chat_cache.embedding_func,
+ func_name="mm_embedding",
+ report_func=chat_cache.report.embedding,
+ )(data_dict)
+ else:
+ embedding_data_resp = time_cal(
+ chat_cache.embedding_func,
+ func_name="mm_embedding",
+ report_func=chat_cache.report.embedding,
+ )(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:
+ embedding_data = np.concatenate((image_embeddings, text_embeddings))
+ # 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!')
+
+ if cache_enable:
+ cache_data_list = time_cal(
+ chat_cache.data_manager.search,
+ func_name="vector_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,
+ mm_type=pre_multi_type,
+ )
+
+ cache_answers = []
+ cache_questions = []
+ cache_image_urls = []
+ cache_image_ids = []
+ 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]
+ 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_text,
+ "embedding": embedding_data,
+ }
+
+ eval_cache_data = {
+ "question": ret[0],
+ "image_url": ret[1],
+ "image_raw": ret[2],
+ "answer": ret[3],
+ "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_text) <= 50:
+ if rank_threshold <= rank:
+ cache_answers.append((rank, ret[3]))
+ cache_image_urls.append((rank, ret[1]))
+ cache_image_ids.append((rank, ret[2]))
+ cache_questions.append((rank, ret[0]))
+ cache_ids.append((rank, primary_id))
+ else:
+ if rank_threshold_long <= rank:
+ cache_answers.append((rank, ret[3]))
+ cache_image_urls.append((rank, ret[1]))
+ cache_image_ids.append((rank, ret[2]))
+ 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_image_urls = sorted(cache_image_urls, key=lambda x: x[0], reverse=True)
+ cache_image_ids = sorted(cache_image_ids, 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_image_url = chat_cache.post_process_messages_func(
+ [t[1] for t in cache_image_urls]
+ )
+ return_image_id = chat_cache.post_process_messages_func(
+ [t[1] for t in cache_image_ids]
+ )
+ 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.warning('update_hit_count except, please check!')
+
+ chat_cache.report.hint_cache()
+ return_query_dict = {"image_url": return_image_url, "image_id": return_image_id, "question": return_query}
+ return cache_data_convert(return_message, return_query_dict)
diff --git a/modelcache_mm/adapter/adapter_register.py b/modelcache_mm/adapter/adapter_register.py
new file mode 100644
index 0000000..8983748
--- /dev/null
+++ b/modelcache_mm/adapter/adapter_register.py
@@ -0,0 +1,13 @@
+# -*- coding: utf-8 -*-
+from modelcache_mm import cache
+
+
+def adapt_register(*args, **kwargs):
+ chat_cache = kwargs.pop("cache_obj", cache)
+ model = kwargs.pop("model", None)
+ type = kwargs.pop("type", None)
+ if model is None or len(model) == 0:
+ return ValueError('')
+
+ register_resp = chat_cache.data_manager.create_index(model, type)
+ return register_resp
diff --git a/modelcache_mm/adapter/adapter_remove.py b/modelcache_mm/adapter/adapter_remove.py
new file mode 100644
index 0000000..b6b1f92
--- /dev/null
+++ b/modelcache_mm/adapter/adapter_remove.py
@@ -0,0 +1,24 @@
+# -*- coding: utf-8 -*-
+from modelcache_mm import cache
+from modelcache_mm.utils.error import NotInitError
+
+
+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"
+ return resp
diff --git a/modelcache/config.py b/modelcache_mm/config.py
similarity index 100%
rename from modelcache/config.py
rename to modelcache_mm/config.py
diff --git a/modelcache_mm/config/chromadb_config.ini b/modelcache_mm/config/chromadb_config.ini
new file mode 100644
index 0000000..2a87f3b
--- /dev/null
+++ b/modelcache_mm/config/chromadb_config.ini
@@ -0,0 +1,2 @@
+[chromadb]
+persist_directory=./chromadb
diff --git a/modelcache_mm/config/elasticsearch_config.ini b/modelcache_mm/config/elasticsearch_config.ini
new file mode 100644
index 0000000..ba87a5a
--- /dev/null
+++ b/modelcache_mm/config/elasticsearch_config.ini
@@ -0,0 +1,5 @@
+[elasticsearch]
+host = ''
+port = ''
+user = ''
+password = ''
\ No newline at end of file
diff --git a/modelcache_mm/config/milvus_config.ini b/modelcache_mm/config/milvus_config.ini
new file mode 100644
index 0000000..f5bd532
--- /dev/null
+++ b/modelcache_mm/config/milvus_config.ini
@@ -0,0 +1,5 @@
+[milvus]
+host = ''
+port = ''
+user = ''
+password = ''
\ No newline at end of file
diff --git a/modelcache_mm/config/mysql_config.ini b/modelcache_mm/config/mysql_config.ini
new file mode 100644
index 0000000..2c63f0e
--- /dev/null
+++ b/modelcache_mm/config/mysql_config.ini
@@ -0,0 +1,6 @@
+[mysql]
+host = ''
+port = ''
+username = ''
+password = ''
+database = ''
diff --git a/modelcache_mm/config/redis_config.ini b/modelcache_mm/config/redis_config.ini
new file mode 100644
index 0000000..a1cdb3d
--- /dev/null
+++ b/modelcache_mm/config/redis_config.ini
@@ -0,0 +1,5 @@
+[redis]
+host = ''
+port = ''
+user = ''
+password = ''
diff --git a/modelcache/core.py b/modelcache_mm/core.py
similarity index 79%
rename from modelcache/core.py
rename to modelcache_mm/core.py
index 331fe6b..d6ff6fb 100644
--- a/modelcache/core.py
+++ b/modelcache_mm/core.py
@@ -1,16 +1,16 @@
# -*- coding: utf-8 -*-
import atexit
from typing import Optional, List, Any
-from modelcache.processor.post import first
-from modelcache.similarity_evaluation import ExactMatchEvaluation
-from modelcache.similarity_evaluation import SimilarityEvaluation
-from modelcache.embedding.string import to_embeddings as string_embedding
-from modelcache.report import Report
-from modelcache.config import Config
-from modelcache.utils.cache_func import cache_all
-from modelcache.utils.log import modelcache_log
-from modelcache.manager import get_data_manager
-from modelcache.manager.data_manager import DataManager
+from modelcache_mm.processor.post import first
+from modelcache_mm.similarity_evaluation import ExactMatchEvaluation
+from modelcache_mm.similarity_evaluation import SimilarityEvaluation
+from modelcache_mm.embedding.string import to_embeddings as string_embedding
+from modelcache_mm.report import Report
+from modelcache_mm.config import Config
+from modelcache_mm.utils.cache_func import cache_all
+from modelcache_mm.utils.log import modelcache_log
+from modelcache_mm.manager import get_data_manager
+from modelcache_mm.manager.data_manager import DataManager
class Cache:
diff --git a/modelcache_mm/embedding/__init__.py b/modelcache_mm/embedding/__init__.py
new file mode 100644
index 0000000..1275963
--- /dev/null
+++ b/modelcache_mm/embedding/__init__.py
@@ -0,0 +1,7 @@
+# -*- coding: utf-8 -*-
+from modelcache.utils.lazy_import import LazyImport
+clip = LazyImport("clip", globals(), "modelcache_mm.embedding.clip")
+
+
+def Clip2Vec(model="damo/multi-modal_clip-vit-base-patch16_zh"):
+ return clip.ClipAudio(model)
diff --git a/modelcache_mm/embedding/base.py b/modelcache_mm/embedding/base.py
new file mode 100644
index 0000000..4b316aa
--- /dev/null
+++ b/modelcache_mm/embedding/base.py
@@ -0,0 +1,17 @@
+# -*- coding: utf-8 -*-
+from abc import ABCMeta, abstractmethod
+
+
+class BaseEmbedding(metaclass=ABCMeta):
+ """
+ _Embedding base.
+ """
+
+ @abstractmethod
+ def to_embeddings(self, data, **kwargs):
+ pass
+
+ @property
+ @abstractmethod
+ def dimension(self) -> int:
+ return 0
diff --git a/modelcache_mm/embedding/clip.py b/modelcache_mm/embedding/clip.py
new file mode 100644
index 0000000..e2e2798
--- /dev/null
+++ b/modelcache_mm/embedding/clip.py
@@ -0,0 +1,49 @@
+# -*- coding: utf-8 -*-
+import torch
+from modelcache.embedding.base import BaseEmbedding
+from modelscope.utils.constant import Tasks
+from modelscope.pipelines import pipeline
+from modelscope.preprocessors.image import load_image
+
+
+class ClipAudio(BaseEmbedding):
+ def __init__(self, model: str = 'damo/multi-modal_clip-vit-base-patch16_zh'):
+ self.model = model
+ self.device = 'cuda' if torch.cuda.is_available() else 'cpu'
+ self.clip_pipeline = pipeline(task=Tasks.multi_modal_embedding,
+ model=model, model_revision='v1.0.1')
+ self.__dimension = 1024
+
+ def to_embeddings(self, data_dict, **_):
+ text_list = data_dict['text']
+ image_data = data_dict['image']
+
+ # img_data = None
+ # txt_data = None
+
+ if image_data:
+ input_img = load_image(image_data)
+ img_embedding = self.clip_pipeline.forward({'img': input_img})['img_embedding'].tolist()[0] if input_img else []
+ else:
+ raise ValueError('image_data is None, please check!')
+
+ if text_list and len(text_list) > 0:
+ text_embedding = self.clip_pipeline.forward({'text': text_list})['text_embedding'].tolist()[0] if text_list else []
+ else:
+ raise ValueError('text_list is None, please check!')
+
+ return {'image_embedding': img_embedding, 'text_embeddings': text_embedding}
+
+ def post_proc(self, token_embeddings, inputs):
+ attention_mask = inputs["attention_mask"]
+ input_mask_expanded = (
+ attention_mask.unsqueeze(-1).expand(token_embeddings.size()).float()
+ )
+ sentence_embs = torch.sum(
+ token_embeddings * input_mask_expanded, 1
+ ) / torch.clamp(input_mask_expanded.sum(1), min=1e-9)
+ return sentence_embs
+
+ @property
+ def dimension(self):
+ return self.__dimension
diff --git a/modelcache/embedding/string.py b/modelcache_mm/embedding/string.py
similarity index 100%
rename from modelcache/embedding/string.py
rename to modelcache_mm/embedding/string.py
diff --git a/modelcache_mm/embedding/timm.py b/modelcache_mm/embedding/timm.py
new file mode 100644
index 0000000..e5c9f48
--- /dev/null
+++ b/modelcache_mm/embedding/timm.py
@@ -0,0 +1,71 @@
+# -*- coding: utf-8 -*-
+import numpy as np
+
+from modelcache.utils import import_timm, import_torch, import_pillow
+from modelcache.embedding.base import BaseEmbedding
+
+import_torch()
+import_timm()
+import_pillow()
+
+import torch
+from timm.models import create_model
+from timm.data import create_transform, resolve_data_config
+from PIL import Image
+
+
+class Timm(BaseEmbedding):
+ def __init__(self, model: str = "resnet18", device: str = "default"):
+ if device == "default":
+ self.device = "cuda" if torch.cuda.is_available() else "cpu"
+ else:
+ self.device = device
+ self.model_name = model
+ self.model = create_model(model_name=model, pretrained=True)
+ self.model.eval()
+
+ try:
+ self.__dimension = self.model.embed_dim
+ except Exception:
+ self.__dimension = None
+
+ def to_embeddings(self, data, skip_preprocess: bool = False, **_):
+ if not skip_preprocess:
+ data = self.preprocess(data)
+ if data.dim() == 3:
+ data = data.unsqueeze(0)
+ feats = self.model.forward_features(data)
+ emb = self.post_proc(feats).squeeze(0).detach().numpy()
+
+ return np.array(emb).astype("float32")
+
+ def post_proc(self, features):
+ features = features.to("cpu")
+ if features.dim() == 3:
+ features = features[:, 0]
+ if features.dim() == 4:
+ global_pool = torch.nn.AdaptiveAvgPool2d(1)
+ features = global_pool(features)
+ features = features.flatten(1)
+ assert features.dim() == 2, f"Invalid output dim {features.dim()}"
+ return features
+
+ def preprocess(self, image_path):
+ data_cfg = resolve_data_config(self.model.pretrained_cfg)
+ transform = create_transform(**data_cfg)
+
+ image = Image.open(image_path).convert("RGB")
+ image_tensor = transform(image)
+ return image_tensor
+
+ @property
+ def dimension(self):
+ """Embedding dimension.
+ :return: embedding dimension
+ """
+ if not self.__dimension:
+ input_size = self.model.pretrained_cfg["input_size"]
+ dummy_input = torch.rand((1,) + input_size)
+ feats = self.to_embeddings(dummy_input, skip_preprocess=True)
+ self.__dimension = feats.shape[0]
+ return self.__dimension
diff --git a/modelcache_mm/manager/__init__.py b/modelcache_mm/manager/__init__.py
new file mode 100644
index 0000000..0b2e8be
--- /dev/null
+++ b/modelcache_mm/manager/__init__.py
@@ -0,0 +1,5 @@
+# -*- coding: utf-8 -*-
+from modelcache_mm.manager.scalar_data import CacheBase
+from modelcache_mm.manager.vector_data import VectorBase
+from modelcache_mm.manager.object_data import ObjectBase
+from modelcache_mm.manager.factory import get_data_manager
diff --git a/modelcache_mm/manager/data_manager.py b/modelcache_mm/manager/data_manager.py
new file mode 100644
index 0000000..338ef15
--- /dev/null
+++ b/modelcache_mm/manager/data_manager.py
@@ -0,0 +1,298 @@
+# -*- coding: utf-8 -*-
+import logging
+import time
+import requests
+import pickle
+import numpy as np
+import cachetools
+from abc import abstractmethod, ABCMeta
+from typing import List, Any, Optional, Union
+from modelcache_mm.manager.scalar_data.base import (
+ CacheStorage,
+ CacheData,
+ DataType,
+ Answer,
+ Question
+)
+from modelcache_mm.utils.error import CacheError, ParamError
+from modelcache_mm.manager.vector_data.base import VectorBase, VectorData
+from modelcache_mm.manager.object_data.base import ObjectBase
+from modelcache_mm.utils.log import modelcache_log
+
+
+class DataManager(metaclass=ABCMeta):
+ """DataManager manage the cache data, including save and search"""
+
+ @abstractmethod
+ 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, texts: List[Any], image_urls: List[Any], image_ids: List[Any], answers: List[Answer],
+ embeddings: List[Any], model: Any, mm_type: Any
+ ):
+ pass
+
+ @abstractmethod
+ def get_scalar_data(self, res_data, **kwargs) -> CacheData:
+ pass
+
+ @abstractmethod
+ def update_hit_count(self, primary_id, **kwargs):
+ pass
+
+ def hit_cache_callback(self, res_data, **kwargs):
+ pass
+
+ @abstractmethod
+ def search(self, embedding_data, **kwargs):
+ pass
+
+ @abstractmethod
+ def delete(self, id_list, **kwargs):
+ pass
+
+ def truncate(self, model_name):
+ pass
+
+ def flush(self):
+ pass
+
+ @abstractmethod
+ def close(self):
+ pass
+
+
+class MapDataManager(DataManager):
+ def __init__(self, data_path, max_size, get_data_container=None):
+ if get_data_container is None:
+ self.data = cachetools.LRUCache(max_size)
+ else:
+ self.data = get_data_container(max_size)
+ self.data_path = data_path
+ self.init()
+
+ def init(self):
+ try:
+ with open(self.data_path, "rb") as f:
+ self.data = pickle.load(f)
+ except FileNotFoundError:
+ return
+ except PermissionError:
+ raise CacheError( # pylint: disable=W0707
+ f"You don't have permission to access this file <{self.data_path}>."
+ )
+
+ 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, texts: List[Any], image_urls: List[Any], image_ids: List[Any], answers: List[Answer],
+ embeddings: List[Any], model: Any, mm_type: Any
+ ):
+ pass
+
+ def get_scalar_data(self, res_data, **kwargs) -> CacheData:
+ return CacheData(question=res_data[0], answers=res_data[1])
+
+ def update_hit_count(self, primary_id, **kwargs):
+ pass
+
+ def search(self, embedding_data, **kwargs):
+ try:
+ return [self.data[embedding_data]]
+ except KeyError:
+ return []
+
+ def delete(self, id_list, **kwargs):
+ pass
+
+ def truncate(self, model_name):
+ pass
+
+ def flush(self):
+ try:
+ with open(self.data_path, "wb") as f:
+ pickle.dump(self.data, f)
+ except PermissionError:
+ modelcache_log.error(
+ "You don't have permission to access this file %s.", self.data_path
+ )
+
+ def close(self):
+ self.flush()
+
+
+def normalize(vec):
+ magnitude = np.linalg.norm(vec)
+ normalized_v = vec / magnitude
+ return normalized_v
+
+
+class SSDataManager(DataManager):
+ def __init__(
+ self,
+ s: CacheStorage,
+ v: VectorBase,
+ o: Optional[ObjectBase],
+ max_size,
+ clean_size,
+ policy="LRU",
+ ):
+ self.max_size = max_size
+ self.clean_size = clean_size
+ self.s = s
+ self.v = v
+ self.o = o
+
+ def save(self, text, image_url, image_id, answer, embedding, **kwargs):
+ model = kwargs.pop("model", None)
+ 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):
+ self.s.insert_query_resp(query_resp_dict, **kwargs)
+
+ def _process_answer_data(self, answers: Union[Answer, List[Answer]]):
+ if isinstance(answers, Answer):
+ answers = [answers]
+ new_ans = []
+ for ans in answers:
+ if ans.answer_type != DataType.STR:
+ new_ans.append(Answer(self.o.put(ans.answer), ans.answer_type))
+ else:
+ new_ans.append(ans)
+ return new_ans
+
+ def _process_question_data(self, question: Union[str, Question]):
+ if isinstance(question, Question):
+ if question.deps is None:
+ return question
+
+ for dep in question.deps:
+ if dep.dep_type == DataType.IMAGE_URL:
+ dep.dep_type.data = self.o.put(requests.get(dep.data).content)
+ return question
+
+ return Question(question)
+
+ def import_data(
+ self, texts: List[Any], image_urls: List[Any], image_ids: List[Any], answers: List[Answer],
+ embeddings: List[Any], model: Any, mm_type: Any
+ ):
+ if len(texts) != len(answers):
+ raise ParamError("Make sure that all parameters have the same length")
+ cache_datas = []
+
+ embeddings = [
+ normalize(text_embedding) for text_embedding in embeddings
+ ]
+
+ for i, embedding in enumerate(embeddings):
+ if self.o is not None:
+ ans = self._process_answer_data(answers[i])
+ else:
+ ans = answers[i]
+ text = texts[i]
+ image_url = image_urls[i]
+ image_id = image_ids[i]
+ cache_datas.append([ans, text, image_url, image_id, model])
+
+ ids = self.s.batch_insert(cache_datas)
+ self.v.add(
+ [
+ VectorData(id=ids[i], data=embedding)
+ for i, embedding in enumerate(embeddings)
+ ],
+ model,
+ mm_type
+ )
+
+ def get_scalar_data(self, res_data, **kwargs) -> Optional[CacheData]:
+ cache_data = self.s.get_data_by_id(res_data[1])
+ if cache_data is None:
+ return None
+ return cache_data
+
+ def update_hit_count(self, primary_id, **kwargs):
+ self.s.update_hit_count_by_id(primary_id)
+
+ def hit_cache_callback(self, res_data, **kwargs):
+ self.eviction_base.get(res_data[1])
+
+ def search(self, embedding_data, **kwargs):
+ model = kwargs.pop("model", None)
+ mm_type = kwargs.pop("mm_type", None)
+ embedding_data = normalize(embedding_data)
+ top_k = kwargs.get("top_k", -1)
+ try:
+ search_result = self.v.search(data=embedding_data, top_k=top_k, model=model, mm_type=mm_type)
+ except Exception as e:
+ try:
+ message = str(e)
+ if "no such index" in message:
+ logging.info('no such index except，creating...')
+ self.v.create(model, mm_type)
+ search_result = self.v.search(data=embedding_data, top_k=top_k, model=model, mm_type=mm_type)
+ except Exception as e:
+ raise e
+ return search_result
+
+ def delete(self, id_list, **kwargs):
+ model = kwargs.pop("model", None)
+ try:
+ v_delete_count = self.v.delete(ids=id_list, model=model)
+ except Exception as e:
+ return {'status': 'failed', 'milvus': 'delete milvus data failed, please check! e: {}'.format(e),
+ 'mysql': 'unexecuted'}
+ try:
+ s_delete_count = self.s.mark_deleted(id_list)
+ except Exception as e:
+ return {'status': 'failed', 'milvus': 'success',
+ 'mysql': 'delete mysql data failed, please check! e: {}'.format(e)}
+
+ return {'status': 'success', 'milvus': 'delete_count: '+str(v_delete_count),
+ 'mysql': 'delete_count: '+str(s_delete_count)}
+
+ def create_index(self, model, type, **kwargs):
+ return self.v.create(model, type)
+
+ def truncate(self, model_name):
+ try:
+ resp = self.v.rebuild_idx(model_name)
+ except Exception as e:
+ return {'status': 'failed', 'VectorDB': 'truncate VectorDB failed, please check! e: {}'.format(str(e)),
+ 'ScalarDB': 'unexecuted'}
+
+ if resp:
+ return {'status': 'failed', 'VectorDB': resp, 'ScalarDB': 'unexecuted'}
+ try:
+ delete_count = self.s.model_deleted(model_name)
+ except Exception as e:
+ # return 'truncate milvus data failed, please check!'
+ return {'status': 'failed', 'VectorDB': 'rebuild',
+ 'ScalarDB': 'truncate scalardb data failed, please check! e: {}'.format(e)}
+
+ return {'status': 'success', 'VectorDB': 'rebuild', 'ScalarDB': 'delete_count: ' + str(delete_count)}
+
+ def flush(self):
+ self.s.flush()
+ self.v.flush()
+
+ def close(self):
+ self.s.close()
+ self.v.close()
+
+
+# if __name__ == '__main__':
+# from modelcache.manager import CacheBase, VectorBase, get_data_manager
+# data_manager = get_data_manager(CacheBase('mysql'), VectorBase('milvus', dimension=128))
+# data_manager.save('hello', 'hi', np.random.random((128,)).astype('float32'), model='gptcode_6b')
diff --git a/modelcache_mm/manager/eviction/__init__.py b/modelcache_mm/manager/eviction/__init__.py
new file mode 100644
index 0000000..8ca7a3d
--- /dev/null
+++ b/modelcache_mm/manager/eviction/__init__.py
@@ -0,0 +1,10 @@
+# -*- coding: utf-8 -*-
+from modelcache.utils.lazy_import import LazyImport
+
+eviction_manager = LazyImport(
+ "eviction_manager", globals(), "modelcache.manager.eviction.manager"
+)
+
+
+def EvictionBase(name: str, **kwargs):
+ return eviction_manager.EvictionBase.get(name, **kwargs)
diff --git a/modelcache_mm/manager/eviction/base.py b/modelcache_mm/manager/eviction/base.py
new file mode 100644
index 0000000..352c9a3
--- /dev/null
+++ b/modelcache_mm/manager/eviction/base.py
@@ -0,0 +1,22 @@
+# -*- coding: utf-8 -*-
+from abc import ABCMeta, abstractmethod
+from typing import Any, List
+
+
+class EvictionBase(metaclass=ABCMeta):
+ """
+ Eviction base.
+ """
+
+ @abstractmethod
+ def put(self, objs: List[Any]):
+ pass
+
+ @abstractmethod
+ def get(self, obj: Any):
+ pass
+
+ @property
+ @abstractmethod
+ def policy(self) -> str:
+ pass
diff --git a/modelcache_mm/manager/eviction_manager.py b/modelcache_mm/manager/eviction_manager.py
new file mode 100644
index 0000000..f64ee7b
--- /dev/null
+++ b/modelcache_mm/manager/eviction_manager.py
@@ -0,0 +1,38 @@
+# -*- coding: utf-8 -*-
+class EvictionManager:
+ MAX_MARK_COUNT = 5000
+ MAX_MARK_RATE = 0.1
+ BATCH_SIZE = 100000
+ REBUILD_CONDITION = 5
+
+ def __init__(self, scalar_storage, vector_base):
+ self._scalar_storage = scalar_storage
+ self._vector_base = vector_base
+ self.delete_count = 0
+
+ def check_evict(self):
+ mark_count = self._scalar_storage.count(state=-1)
+ all_count = self._scalar_storage.count(is_all=True)
+ if (
+ mark_count > self.MAX_MARK_COUNT
+ or mark_count / all_count > self.MAX_MARK_RATE
+ ):
+ return True
+ return False
+
+ def delete(self):
+ mark_ids = self._scalar_storage.get_ids(deleted=True)
+ self._scalar_storage.clear_deleted_data()
+ self._vector_base.delete(mark_ids)
+ self.delete_count += 1
+ if self.delete_count >= self.REBUILD_CONDITION:
+ self.rebuild()
+
+ def rebuild(self):
+ self._scalar_storage.clear_deleted_data()
+ ids = self._scalar_storage.get_ids(deleted=False)
+ self._vector_base.rebuild(ids)
+ self.delete_count = 0
+
+ def soft_evict(self, marked_keys):
+ self._scalar_storage.mark_deleted(marked_keys)
diff --git a/modelcache/manager/factory.py b/modelcache_mm/manager/factory.py
similarity index 86%
rename from modelcache/manager/factory.py
rename to modelcache_mm/manager/factory.py
index 7246cc6..1f43b4a 100644
--- a/modelcache/manager/factory.py
+++ b/modelcache_mm/manager/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_mm.manager import CacheBase, VectorBase, ObjectBase
+from modelcache_mm.manager.data_manager import SSDataManager, MapDataManager
def get_data_manager(
@@ -16,7 +16,6 @@ def get_data_manager(
):
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):
diff --git a/modelcache_mm/manager/object_data/__init__.py b/modelcache_mm/manager/object_data/__init__.py
new file mode 100644
index 0000000..f1186c0
--- /dev/null
+++ b/modelcache_mm/manager/object_data/__init__.py
@@ -0,0 +1,9 @@
+# -*- coding: utf-8 -*-
+from modelcache.utils.lazy_import import LazyImport
+object_manager = LazyImport(
+ "object_manager", globals(), "modelcache.manager.object_data.manager"
+)
+
+
+def ObjectBase(name: str, **kwargs):
+ return object_manager.ObjectBase.get(name, **kwargs)
diff --git a/modelcache_mm/manager/object_data/base.py b/modelcache_mm/manager/object_data/base.py
new file mode 100644
index 0000000..9ed00cf
--- /dev/null
+++ b/modelcache_mm/manager/object_data/base.py
@@ -0,0 +1,25 @@
+# -*- coding: utf-8 -*-
+from abc import ABC, abstractmethod
+from typing import Any, List
+
+
+class ObjectBase(ABC):
+ """
+ Object storage base.
+ """
+
+ @abstractmethod
+ def put(self, obj: Any) -> str:
+ pass
+
+ @abstractmethod
+ def get(self, obj: str) -> Any:
+ pass
+
+ @abstractmethod
+ def get_access_link(self, obj: str) -> str:
+ pass
+
+ @abstractmethod
+ def delete(self, to_delete: List[str]):
+ pass
diff --git a/modelcache_mm/manager/scalar_data/__init__.py b/modelcache_mm/manager/scalar_data/__init__.py
new file mode 100644
index 0000000..e499db4
--- /dev/null
+++ b/modelcache_mm/manager/scalar_data/__init__.py
@@ -0,0 +1,9 @@
+# -*- coding: utf-8 -*-
+from modelcache_mm.utils.lazy_import import LazyImport
+scalar_manager = LazyImport(
+ "scalar_manager", globals(), "modelcache_mm.manager.scalar_data.manager"
+)
+
+
+def CacheBase(name: str, **kwargs):
+ return scalar_manager.CacheBase.get(name, **kwargs)
diff --git a/modelcache_mm/manager/scalar_data/base.py b/modelcache_mm/manager/scalar_data/base.py
new file mode 100644
index 0000000..fd8bb50
--- /dev/null
+++ b/modelcache_mm/manager/scalar_data/base.py
@@ -0,0 +1,134 @@
+# -*- coding: utf-8 -*-
+from abc import ABCMeta, abstractmethod
+from dataclasses import dataclass
+from typing import Union, Dict, List, Optional, Any
+from enum import IntEnum
+import numpy as np
+
+
+class DataType(IntEnum):
+ STR = 0
+ IMAGE_BASE64 = 1
+ IMAGE_URL = 2
+
+
+@dataclass
+class QuestionDep:
+ """
+ QuestionDep
+ """
+
+ name: str
+ data: str
+ dep_type: int = DataType.STR
+
+ @classmethod
+ def from_dict(cls, d: Dict):
+ return cls(
+ name=d["name"],
+ data=d["data"],
+ dep_type=d["dep_type"]
+ )
+
+
+@dataclass
+class Question:
+ """
+ Question
+ """
+
+ content: str
+ deps: Optional[List[QuestionDep]] = None
+
+ @classmethod
+ def from_dict(cls, d: Dict):
+ deps = []
+ for dep in d["deps"]:
+ deps.append(QuestionDep.from_dict(dep))
+ return cls(d["content"], deps)
+
+
+@dataclass
+class Answer:
+ """
+ data_type:
+ 0: str
+ 1: base64 image
+ """
+
+ answer: Any
+ answer_type: int = DataType.STR
+
+
+@dataclass
+class CacheData:
+ """
+ CacheData
+ """
+
+ question: Union[str, Question]
+ answers: List[Answer]
+ embedding_data: Optional[np.ndarray] = None
+
+ def __init__(self, question, answers, embedding_data=None):
+ self.question = question
+ self.answers = []
+ if isinstance(answers, (str, Answer)):
+ answers = [answers]
+ for data in answers:
+ if isinstance(data, (list, tuple)):
+ self.answers.append(Answer(*data))
+ elif isinstance(data, Answer):
+ self.answers.append(data)
+ else:
+ self.answers.append(Answer(answer=data))
+ self.embedding_data = embedding_data
+
+
+class CacheStorage(metaclass=ABCMeta):
+ """
+ BaseStorage for scalar data.
+ """
+
+ @abstractmethod
+ def create(self):
+ pass
+
+ @abstractmethod
+ def batch_insert(self, all_data: List[CacheData]):
+ pass
+
+ @abstractmethod
+ def insert_query_resp(self, query_resp, **kwargs):
+ pass
+
+ @abstractmethod
+ def get_data_by_id(self, key):
+ pass
+
+ @abstractmethod
+ def mark_deleted(self, keys):
+ pass
+
+ @abstractmethod
+ def model_deleted(self, model_name):
+ pass
+
+ @abstractmethod
+ def clear_deleted_data(self):
+ pass
+
+ @abstractmethod
+ def get_ids(self, deleted=True):
+ pass
+
+ @abstractmethod
+ def count(self):
+ pass
+
+ def flush(self):
+ pass
+
+ @abstractmethod
+ def close(self):
+ pass
diff --git a/modelcache/manager/scalar_data/manager.py b/modelcache_mm/manager/scalar_data/manager.py
similarity index 74%
rename from modelcache/manager/scalar_data/manager.py
rename to modelcache_mm/manager/scalar_data/manager.py
index 4c02c45..3781787 100644
--- a/modelcache/manager/scalar_data/manager.py
+++ b/modelcache_mm/manager/scalar_data/manager.py
@@ -1,6 +1,6 @@
# -*- coding: utf-8 -*-
-from modelcache.utils import import_sql_client
-from modelcache.utils.error import NotFoundError
+from modelcache_mm.utils import import_sql_client
+from modelcache_mm.utils.error import NotFoundError
SQL_URL = {"sqlite": "./sqlite.db"}
@@ -19,12 +19,12 @@ def __init__(self):
def get(name, **kwargs):
if name in ["mysql", "oceanbase"]:
- from modelcache.manager.scalar_data.sql_storage import SQLStorage
+ from modelcache_mm.manager.scalar_data.sql_storage import SQLStorage
config = kwargs.get("config")
import_sql_client(name)
cache_base = SQLStorage(db_type=name, config=config)
elif name == 'sqlite':
- from modelcache.manager.scalar_data.sql_storage_sqlite import SQLStorage
+ from modelcache_mm.manager.scalar_data.sql_storage_sqlite import SQLStorage
sql_url = kwargs.get("sql_url", SQL_URL[name])
cache_base = SQLStorage(db_type=name, url=sql_url)
else:
diff --git a/modelcache_mm/manager/scalar_data/sql_storage.py b/modelcache_mm/manager/scalar_data/sql_storage.py
new file mode 100644
index 0000000..f80fb1b
--- /dev/null
+++ b/modelcache_mm/manager/scalar_data/sql_storage.py
@@ -0,0 +1,203 @@
+# -*- coding: utf-8 -*-
+import os
+import time
+
+import pymysql
+import json
+import base64
+from typing import List
+from modelcache_mm.manager.scalar_data.base import CacheStorage, CacheData
+from DBUtils.PooledDB import PooledDB
+
+
+class SQLStorage(CacheStorage):
+ def __init__(
+ self,
+ db_type: str = "mysql",
+ config=None
+ ):
+ self.host = config.get('mysql', 'host')
+ self.port = int(config.get('mysql', 'port'))
+ self.username = config.get('mysql', 'username')
+ self.password = config.get('mysql', 'password')
+ self.database = config.get('mysql', 'database')
+ self.pool = PooledDB(
+ creator=pymysql,
+ host=self.host,
+ user=self.username,
+ password=self.password,
+ port=self.port,
+ database=self.database
+ )
+
+ def create(self):
+ pass
+
+ # def _insert(self, data: List):
+ # answer = data[0]
+ # text = data[1]
+ # image_url = data[2]
+ # image_id = data[3]
+ # model = data[4]
+ # answer_type = 0
+ #
+ # 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:
+ # # 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
+
+ def _insert(self, data: List):
+ answer = data[0]
+ text = data[1]
+ image_url = data[2]
+ image_id = data[3]
+ model = data[4]
+ answer_type = 0
+
+ table_name = "open_cache_mm_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:
+ # insert data 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 to the connection pool.
+ conn.close()
+ return id
+
+ def batch_insert(self, all_data: List[CacheData]):
+ ids = []
+ for data in all_data:
+ ids.append(self._insert(data))
+ return ids
+
+ def insert_query_resp(self, query_resp, **kwargs):
+ error_code = query_resp.get('errorCode')
+ error_desc = query_resp.get('errorDesc')
+ cache_hit = query_resp.get('cacheHit')
+ model = kwargs.get('model')
+ query = kwargs.get('query')
+ delta_time = kwargs.get('delta_time')
+ hit_query = query_resp.get('hit_query')
+ answer = query_resp.get('answer')
+
+ if isinstance(hit_query, list):
+ hit_query = json.dumps(hit_query, ensure_ascii=False)
+
+ table_name = "open_cache_mm_query_log"
+ insert_sql = "INSERT INTO {} (error_code, error_desc, cache_hit, model, query, delta_time, hit_query, answer) VALUES (%s, %s, %s, %s, %s, %s, %s, %s)".format(table_name)
+ conn = self.pool.connection()
+ try:
+ with conn.cursor() as cursor:
+ # 执行插入数据操作
+ values = (error_code, error_desc, cache_hit, model, query, delta_time, hit_query, answer)
+ cursor.execute(insert_sql, values)
+ conn.commit()
+ finally:
+ # 关闭连接，将连接返回给连接池
+ conn.close()
+
+ # def get_data_by_id(self, key: int):
+ # table_name = "cache_codegpt_answer"
+ # query_sql = "select question, answer, embedding_data, model from {} where id={}".format(table_name, key)
+ # conn_start = time.time()
+ # conn = self.pool.connection()
+ #
+ # search_start = time.time()
+ # try:
+ # with conn.cursor() as cursor:
+ # # 执行数据库操作
+ # cursor.execute(query_sql)
+ # resp = cursor.fetchone()
+ # finally:
+ # # 关闭连接，将连接返回给连接池
+ # conn.close()
+ #
+ # if resp is not None and len(resp) == 4:
+ # return resp
+ # else:
+ # return None
+
+ def get_data_by_id(self, key: int):
+ table_name = "open_cache_mm_answer"
+ query_sql = "select question_text, image_url, image_id, answer, model from {} where id={}".format(table_name, key)
+ conn = self.pool.connection()
+ try:
+ with conn.cursor() as cursor:
+ cursor.execute(query_sql)
+ resp = cursor.fetchone()
+ finally:
+ conn.close()
+
+ if resp is not None and len(resp) == 5:
+ return resp
+ else:
+ return None
+
+ def update_hit_count_by_id(self, primary_id: int):
+ table_name = "open_cache_mm_answer"
+ update_sql = "UPDATE {} SET hit_count = hit_count+1 WHERE id={}".format(table_name, primary_id)
+ conn = self.pool.connection()
+
+ try:
+ with conn.cursor() as cursor:
+ cursor.execute(update_sql)
+ conn.commit()
+ finally:
+ conn.close()
+
+ def get_ids(self, deleted=True):
+ pass
+
+ def mark_deleted(self, keys):
+ table_name = "open_cache_mm_answer"
+ delete_sql = "Delete from {} WHERE id in ({})".format(table_name, ",".join([str(i) for i in keys]))
+
+ conn = self.pool.connection()
+ try:
+ with conn.cursor() as cursor:
+ cursor.execute(delete_sql)
+ delete_count = cursor.rowcount
+ conn.commit()
+ finally:
+ conn.close()
+ return delete_count
+
+ def model_deleted(self, model_name):
+ table_name = "open_cache_mm_answer"
+ delete_sql = "Delete from {} WHERE model='{}'".format(table_name, model_name)
+ conn = self.pool.connection()
+ try:
+ with conn.cursor() as cursor:
+ resp = cursor.execute(delete_sql)
+ conn.commit()
+ finally:
+ conn.close()
+ return resp
+
+ def clear_deleted_data(self):
+ pass
+
+ def count(self, state: int = 0, is_all: bool = False):
+ pass
+
+ def close(self):
+ pass
+
+ def count_answers(self):
+ pass
diff --git a/modelcache_mm/manager/scalar_data/sql_storage_es.py b/modelcache_mm/manager/scalar_data/sql_storage_es.py
new file mode 100644
index 0000000..562b349
--- /dev/null
+++ b/modelcache_mm/manager/scalar_data/sql_storage_es.py
@@ -0,0 +1,198 @@
+# -*- coding: utf-8 -*-
+import json
+from typing import List
+from elasticsearch import Elasticsearch, helpers
+from modelcache.manager.scalar_data.base import CacheStorage, CacheData
+import time
+from snowflake import SnowflakeGenerator
+
+
+class SQLStorage(CacheStorage):
+ def __init__(
+ self,
+ db_type: str = "elasticsearch",
+ config=None
+ ):
+ self.host = config.get('elasticsearch', 'host')
+ self.port = int(config.get('elasticsearch', 'port'))
+ self.client = Elasticsearch(
+ hosts=[{"host": self.host, "port": self.port}],
+ timeout=30,
+ http_auth=('esuser', 'password')
+ )
+
+ self.log_index = "open_cache_mm_query_log"
+ self.ans_index = "open_cache_mm_answer"
+ self.create()
+ self.instance_id = 1 # 雪花算法使用的机器id 使用同一套数据库的分布式系统需要配置不同id
+ # 生成雪花id
+ self.snowflake_id = SnowflakeGenerator(self.instance_id)
+
+ def create(self):
+ answer_index_body = {
+ "mappings": {
+ "properties": {
+ "gmt_create": {"type": "date", "format": "strict_date_optional_time||epoch_millis"},
+ "gmt_modified": {"type": "date", "format": "strict_date_optional_time||epoch_millis"},
+ "question": {"type": "text"},
+ "answer": {"type": "text"},
+ "answer_type": {"type": "integer"},
+ "hit_count": {"type": "integer"},
+ "model": {"type": "keyword"},
+ "image_url": {"type": "text"},
+ "image_id": {"type": "text"},
+ "is_deleted": {"type": "integer"},
+ }
+ }
+ }
+
+ log_index_body = {
+ "mappings": {
+ "properties": {
+ "gmt_create": {"type": "date", "format": "strict_date_optional_time||epoch_millis"},
+ "gmt_modified": {"type": "date", "format": "strict_date_optional_time||epoch_millis"},
+ "error_code": {"type": "integer"},
+ "error_desc": {"type": "text"},
+ "cache_hit": {"type": "keyword"},
+ "delta_time": {"type": "float"},
+ "model": {"type": "keyword"},
+ "query": {"type": "text"},
+ "hit_query": {"type": "text"},
+ "answer": {"type": "text"}
+ }
+ }
+ }
+
+ if not self.client.indices.exists(index=self.ans_index):
+ self.client.indices.create(index=self.ans_index, body=answer_index_body)
+
+ if not self.client.indices.exists(index=self.log_index):
+ self.client.indices.create(index=self.log_index, body=log_index_body)
+
+ def _insert(self, data: List) -> str or None:
+ doc = {
+ "answer": data[0],
+ "question": data[1],
+ "image_url": data[2],
+ "image_id": data[3],
+ "model": data[4],
+ "answer_type": 0,
+ "hit_count": 0,
+ "is_deleted": 0
+ }
+
+ try:
+
+ response = self.client.index(
+ index=self.ans_index,
+ id=next(self.snowflake_id),
+ body=doc,
+ )
+ return int(response['_id'])
+ except Exception as e:
+
+ print(f"Failed to insert document: {e}")
+ return None
+
+ def batch_insert(self, all_data: List[List]) -> List[str]:
+ successful_ids = []
+ for data in all_data:
+ _id = self._insert(data)
+ if _id is not None:
+ successful_ids.append(_id)
+ self.client.indices.refresh(index=self.ans_index) # 批量插入后手动刷新
+
+ return successful_ids
+
+ def insert_query_resp(self, query_resp, **kwargs):
+ doc = {
+ "error_code": query_resp.get('errorCode'),
+ "error_desc": query_resp.get('errorDesc'),
+ "cache_hit": query_resp.get('cacheHit'),
+ "model": kwargs.get('model'),
+ "query": kwargs.get('query'),
+ "delta_time": kwargs.get('delta_time'),
+ "hit_query": json.dumps(query_resp.get('hit_query'), ensure_ascii=False) if isinstance(
+ query_resp.get('hit_query'), list) else query_resp.get('hit_query'),
+ "answer": query_resp.get('answer'),
+ "hit_count": 0,
+ "is_deleted": 0
+
+ }
+ self.client.index(index=self.log_index, body=doc)
+
+ def get_data_by_id(self, key: int):
+ try:
+ response = self.client.get(index=self.ans_index, id=key, _source=['question', 'image_url','image_id', 'answer', 'model'])
+ source = response["_source"]
+ result = [
+ source.get('question'),
+ source.get('image_url'),
+ source.get('image_id'),
+ source.get('answer'),
+ source.get('model')
+ ]
+ return result
+ except Exception as e:
+ print(e)
+
+ def update_hit_count_by_id(self, primary_id: int):
+ self.client.update(
+ index=self.ans_index,
+ id=primary_id,
+ body={"script": {"source": "ctx._source.hit_count += 1"}}
+ )
+
+ def get_ids(self, deleted=True):
+ query = {
+ "query": {
+ "term": {"is_deleted": 1 if deleted else 0}
+ }
+ }
+ response = self.client.search(index=self.ans_index, body=query)
+ return [hit["_id"] for hit in response["hits"]["hits"]]
+
+ def mark_deleted(self, keys):
+ actions = [
+ {
+ "_op_type": "update",
+ "_index": self.ans_index,
+ "_id": key,
+ "doc": {"is_deleted": 1}
+ }
+ for key in keys
+ ]
+ responses = helpers.bulk(self.client, actions)
+ return responses[0] # 返回更新的文档数
+
+ def model_deleted(self, model_name):
+ query = {
+ "query": {
+ "term": {"model": model_name}
+ }
+ }
+
+ response = self.client.delete_by_query(index=self.ans_index, body=query)
+ return response["deleted"]
+
+ def clear_deleted_data(self):
+ query = {
+ "query": {
+ "term": {"is_deleted": 1}
+ }
+ }
+ response = self.client.delete_by_query(index=self.ans_index, body=query)
+ return response["deleted"]
+
+ def count(self, state: int = 0, is_all: bool = False):
+ query = {"query": {"match_all": {}}} if is_all else {"query": {"term": {"is_deleted": state}}}
+ response = self.client.count(index=self.ans_index, body=query)
+ return response["count"]
+
+ def close(self):
+ self.client.close()
+
+ def count_answers(self):
+ query = {"query": {"match_all": {}}}
+ response = self.client.count(index=self.ans_index, body=query)
+ return response["count"]
diff --git a/modelcache_mm/manager/scalar_data/sql_storage_sqlite.py b/modelcache_mm/manager/scalar_data/sql_storage_sqlite.py
new file mode 100644
index 0000000..9d16c00
--- /dev/null
+++ b/modelcache_mm/manager/scalar_data/sql_storage_sqlite.py
@@ -0,0 +1,206 @@
+# -*- coding: utf-8 -*-
+import json
+from typing import List
+from modelcache.manager.scalar_data.base import CacheStorage, CacheData
+import sqlite3
+
+
+class SQLStorage(CacheStorage):
+ def __init__(
+ self,
+ db_type: str = "mysql",
+ config=None,
+ url="./sqlite.db"
+ ):
+ self._url = url
+ self.create()
+
+ def create(self):
+ # answer_table_sql = """CREATE TABLE IF NOT EXISTS modelcache_llm_answer (
+ # id INTEGER PRIMARY KEY AUTOINCREMENT,
+ # gmt_create TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP,
+ # gmt_modified TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP,
+ # question TEXT NOT NULL,
+ # answer TEXT NOT NULL,
+ # answer_type INTEGER NOT NULL,
+ # hit_count INTEGER NOT NULL DEFAULT 0,
+ # model VARCHAR(1000) NOT NULL,
+ # embedding_data BLOB NOT NULL
+ # );
+ # """
+
+ answer_table_sql = """CREATE TABLE IF NOT EXISTS `open_cache_mm_answer` (
+ `id` INTEGER PRIMARY KEY AUTOINCREMENT,
+ `gmt_create` TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP,
+ `gmt_modified` TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP,
+ `question_text` TEXT NOT NULL,
+ `image_url` VARCHAR(2048) NOT NULL,
+ `answer` TEXT NOT NULL,
+ `answer_type` INTEGER NOT NULL,
+ `hit_count` INTEGER NOT NULL DEFAULT 0,
+ `model` VARCHAR(1000) NOT NULL,
+ `image_raw` BLOB DEFAULT NULL,
+ `image_id` VARCHAR(1000) DEFAULT NULL
+);
+ """
+
+ log_table_sql = """CREATE TABLE IF NOT EXISTS modelcache_query_log (
+ id INTEGER PRIMARY KEY AUTOINCREMENT,
+ gmt_create TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP,
+ gmt_modified TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP,
+ error_code INTEGER NOT NULL,
+ error_desc VARCHAR(1000) NOT NULL,
+ cache_hit VARCHAR(100) NOT NULL,
+ delta_time REAL NOT NULL,
+ model VARCHAR(1000) NOT NULL,
+ query TEXT NOT NULL,
+ hit_query TEXT NOT NULL,
+ answer TEXT NOT NULL
+ );
+ """
+
+ conn = sqlite3.connect(self._url)
+ try:
+ cursor = conn.cursor()
+ cursor.execute(answer_table_sql)
+ cursor.execute(log_table_sql)
+ conn.commit()
+ cursor.close()
+ conn.close()
+ finally:
+ conn.close()
+
+ def _insert(self, data: List):
+ answer = data[0]
+ text = data[1]
+ image_url = data[2]
+ image_id = data[3]
+ model = data[4]
+ answer_type = 0
+
+ table_name = "open_cache_mm_answer"
+ insert_sql = "INSERT INTO {} (question_text, image_url, image_id, answer, answer_type, model) VALUES (?, ?, ?, ?, ?, ?)".format(table_name)
+
+ conn = sqlite3.connect(self._url)
+ try:
+ cursor = conn.cursor()
+ values = (text, image_url, image_id, answer, answer_type, model)
+ cursor.execute(insert_sql, values)
+ conn.commit()
+ id = cursor.lastrowid
+ cursor.close()
+ conn.close()
+ finally:
+ conn.close()
+ return id
+
+ def batch_insert(self, all_data: List[CacheData]):
+ ids = []
+ for data in all_data:
+ ids.append(self._insert(data))
+ return ids
+
+ def insert_query_resp(self, query_resp, **kwargs):
+ error_code = query_resp.get('errorCode')
+ error_desc = query_resp.get('errorDesc')
+ cache_hit = query_resp.get('cacheHit')
+ model = kwargs.get('model')
+ query = kwargs.get('query')
+ delta_time = kwargs.get('delta_time')
+ hit_query = query_resp.get('hit_query')
+ answer = query_resp.get('answer')
+
+ if isinstance(hit_query, list):
+ hit_query = json.dumps(hit_query, ensure_ascii=False)
+
+ table_name = "modelcache_query_log"
+ insert_sql = "INSERT INTO {} (error_code, error_desc, cache_hit, model, query, delta_time, hit_query, answer) VALUES (%s, %s, %s, %s, %s, %s, %s, %s)".format(table_name)
+
+ conn = sqlite3.connect(self._url)
+ try:
+ cursor = conn.cursor()
+ values = (error_code, error_desc, cache_hit, model, query, delta_time, hit_query, answer)
+ cursor.execute(insert_sql, values)
+ conn.commit()
+ cursor.close()
+ conn.close()
+ finally:
+ conn.close()
+
+ def get_data_by_id(self, key: int):
+ table_name = "open_cache_mm_answer"
+ query_sql = "select question, answer, embedding_data, model from {} where id={}".format(table_name, key)
+ conn = sqlite3.connect(self._url)
+ try:
+ cursor = conn.cursor()
+ cursor.execute(query_sql)
+ resp = cursor.fetchone()
+ conn.commit()
+ cursor.close()
+ conn.close()
+ finally:
+ conn.close()
+
+ if resp is not None and len(resp) == 4:
+ return resp
+ else:
+ return None
+
+ def update_hit_count_by_id(self, primary_id: int):
+ table_name = "open_cache_mm_answer"
+ update_sql = "UPDATE {} SET hit_count = hit_count+1 WHERE id={}".format(table_name, primary_id)
+
+ conn = sqlite3.connect(self._url)
+ try:
+ cursor = conn.cursor()
+ cursor.execute(update_sql)
+ conn.commit()
+ cursor.close()
+ conn.close()
+ finally:
+ # 关闭连接，将连接返回给连接池
+ conn.close()
+
+ def get_ids(self, deleted=True):
+ pass
+
+ def mark_deleted(self, keys):
+ table_name = "open_cache_mm_answer"
+ delete_sql = "Delete from {} WHERE id in ({})".format(table_name, ",".join([str(i) for i in keys]))
+ conn = sqlite3.connect(self._url)
+ try:
+ cursor = conn.cursor()
+ cursor.execute(delete_sql)
+ delete_count = cursor.rowcount
+ conn.commit()
+ cursor.close()
+ conn.close()
+ finally:
+ conn.close()
+ return delete_count
+
+ def model_deleted(self, model_name):
+ table_name = "open_cache_mm_answer"
+ delete_sql = "Delete from {} WHERE model='{}'".format(table_name, model_name)
+ conn = sqlite3.connect(self._url)
+ try:
+ cursor = conn.cursor()
+ resp = cursor.execute(delete_sql)
+ conn.commit()
+ cursor.close()
+ conn.close()
+ finally:
+ conn.close()
+ return resp
+
+ def clear_deleted_data(self):
+ pass
+
+ def count(self, state: int = 0, is_all: bool = False):
+ pass
+
+ def close(self):
+ pass
+
+ def count_answers(self):
+ pass
diff --git a/modelcache_mm/manager/vector_data/__init__.py b/modelcache_mm/manager/vector_data/__init__.py
new file mode 100644
index 0000000..44c30ee
--- /dev/null
+++ b/modelcache_mm/manager/vector_data/__init__.py
@@ -0,0 +1,10 @@
+# -*- coding: utf-8 -*-
+from modelcache_mm.utils.lazy_import import LazyImport
+
+vector_manager = LazyImport(
+ "vector_manager", globals(), "modelcache_mm.manager.vector_data.manager"
+)
+
+
+def VectorBase(name: str, **kwargs):
+ return vector_manager.VectorBase.get(name, **kwargs)
diff --git a/modelcache_mm/manager/vector_data/base.py b/modelcache_mm/manager/vector_data/base.py
new file mode 100644
index 0000000..98151a8
--- /dev/null
+++ b/modelcache_mm/manager/vector_data/base.py
@@ -0,0 +1,49 @@
+# -*- coding: utf-8 -*-
+from abc import ABC, abstractmethod
+import numpy as np
+from typing import List
+from dataclasses import dataclass
+
+
+@dataclass
+class VectorData:
+ id: int
+ data: np.ndarray
+
+
+class VectorBase(ABC):
+ """VectorBase: base vector store interface"""
+
+ @abstractmethod
+ def add(self, datas: List[VectorData], model=None, mm_type=None):
+ pass
+
+ # @abstractmethod
+ # def search(self, data: np.ndarray, top_k: int, model):
+ # pass
+
+ @abstractmethod
+ def search(self, data: np.ndarray, top_k: int, model, mm_type):
+ pass
+
+ @abstractmethod
+ def create(self, model=None, mm_type=None):
+ pass
+
+ @abstractmethod
+ def rebuild(self, ids=None) -> bool:
+ pass
+
+ @abstractmethod
+ def delete(self, ids) -> bool:
+ pass
+
+ @abstractmethod
+ def rebuild_idx(self, model):
+ pass
+
+ def flush(self):
+ pass
+
+ def close(self):
+ pass
diff --git a/modelcache_mm/manager/vector_data/chroma.py b/modelcache_mm/manager/vector_data/chroma.py
new file mode 100644
index 0000000..7f5b78c
--- /dev/null
+++ b/modelcache_mm/manager/vector_data/chroma.py
@@ -0,0 +1,99 @@
+from typing import List
+
+import numpy as np
+import logging
+from modelcache_mm.manager.vector_data.base import VectorBase, VectorData
+from modelcache_mm.utils import import_chromadb, import_torch
+from modelcache_mm.utils.index_util import get_mm_index_name
+
+import_torch()
+import_chromadb()
+
+import chromadb
+
+
+class Chromadb(VectorBase):
+
+ def __init__(
+ self,
+ persist_directory="./chromadb",
+ top_k: int = 1,
+ ):
+ self.top_k = top_k
+
+ self._client = chromadb.PersistentClient(path=persist_directory)
+ self._collection = None
+
+ def create(self, model=None, mm_type=None):
+ try:
+ collection_name_model = get_mm_index_name(model, mm_type)
+ # collection_name_model = self.collection_name + '_' + model
+ self._client.get_or_create_collection(name=collection_name_model)
+ except Exception as e:
+ raise ValueError(str(e))
+
+ def add(self, datas: List[VectorData], model=None, mm_type=None):
+ collection_name_model = get_mm_index_name(model, mm_type)
+ self._collection = self._client.get_or_create_collection(name=collection_name_model)
+
+ data_array, id_array = map(list, zip(*((data.data.tolist(), str(data.id)) for data in datas)))
+ self._collection.add(embeddings=data_array, ids=id_array)
+
+ def search(self, data: np.ndarray, top_k: int = -1, model=None, mm_type='mm'):
+ collection_name_model = get_mm_index_name(model, mm_type)
+ self._collection = self._client.get_or_create_collection(name=collection_name_model)
+
+ if self._collection.count() == 0:
+ return []
+ if top_k == -1:
+ top_k = self.top_k
+ results = self._collection.query(
+ query_embeddings=[data.tolist()],
+ n_results=top_k,
+ include=["distances"],
+ )
+ return list(zip(results["distances"][0], [int(x) for x in results["ids"][0]]))
+
+ def delete(self, ids, model=None, mm_type=None):
+ try:
+ collection_name_model = get_mm_index_name(model, mm_type)
+ self._collection = self._client.get_or_create_collection(name=collection_name_model)
+ # 查询集合中实际存在的 ID
+ ids_str = [str(x) for x in ids]
+ existing_ids = set(self._collection.get(ids=ids_str).ids)
+
+ # 删除存在的 ID
+ if existing_ids:
+ self._collection.delete(list(existing_ids))
+
+ # 返回实际删除的条目数量
+ return len(existing_ids)
+
+ except Exception as e:
+ logging.error('Error during deletion: {}'.format(e))
+ raise ValueError(str(e))
+
+ def rebuild_idx(self, model, mm_type=None):
+ collection_name_model = get_mm_index_name(model, mm_type)
+
+ # 检查集合是否存在，如果存在则删除
+ collections = self._client.list_collections()
+ if any(col.name == collection_name_model for col in collections):
+ self._client.delete_collection(collection_name_model)
+ else:
+ return 'model collection not found, please check!'
+
+ try:
+ self._client.create_collection(collection_name_model)
+ except Exception as e:
+ logging.info(f'rebuild_collection: {e}')
+ raise ValueError(str(e))
+
+ def rebuild(self, ids=None):
+ pass
+
+ def flush(self):
+ pass
+
+ def close(self):
+ pass
diff --git a/modelcache_mm/manager/vector_data/faiss.py b/modelcache_mm/manager/vector_data/faiss.py
new file mode 100644
index 0000000..75b2d3a
--- /dev/null
+++ b/modelcache_mm/manager/vector_data/faiss.py
@@ -0,0 +1,73 @@
+# -*- coding: utf-8 -*-
+import os
+from typing import List
+import numpy as np
+from modelcache_mm.manager.vector_data.base import VectorBase, VectorData
+from modelcache_mm.utils import import_faiss
+import_faiss()
+import faiss # pylint: disable=C0413
+
+
+class Faiss(VectorBase):
+ def __init__(self,
+ index_file_path,
+ dimension: int = 0,
+ top_k: int = 1
+ ):
+ self._dimension = dimension
+ self._index_file_path = index_file_path
+ self._index = faiss.index_factory(self._dimension, "IDMap,Flat", faiss.METRIC_L2)
+ self._top_k = top_k
+ if os.path.isfile(index_file_path):
+ self._index = faiss.read_index(index_file_path)
+
+ def add(self, datas: List[VectorData], model=None, mm_type=None):
+ data_array, id_array = map(list, zip(*((data.data, data.id) for data in datas)))
+ np_data = np.array(data_array).astype("float32")
+ ids = np.array(id_array)
+ self._index.add_with_ids(np_data, ids)
+
+ def search(self, data: np.ndarray, top_k: int, model, mm_type='mm'):
+ if self._index.ntotal == 0:
+ return None
+ if top_k == -1:
+ top_k = self._top_k
+ np_data = np.array(data).astype("float32").reshape(1, -1)
+ dist, ids = self._index.search(np_data, top_k)
+ ids = [int(i) for i in ids[0]]
+ return list(zip(dist[0], ids))
+
+ def rebuild_col(self, ids=None):
+ try:
+ self._index.reset()
+ except Exception as e:
+ return f"An error occurred during index rebuild: {e}"
+
+ def rebuild(self, ids=None):
+ return True
+
+ def delete(self, ids):
+ ids_to_remove = np.array(ids)
+ self._index.remove_ids(faiss.IDSelectorBatch(ids_to_remove.size, faiss.swig_ptr(ids_to_remove)))
+
+ def create(self, model=None, mm_type=None):
+ pass
+ # collection_name_model = get_mm_index_name(model, mm_type)
+ # try:
+ # index_prefix = get_mm_index_prefix(model, mm_type)
+ # self.create_index(collection_name_model, mm_type, index_prefix)
+ # except Exception as e:
+ # raise ValueError(str(e))
+ # return 'success'
+
+ def flush(self):
+ faiss.write_index(self._index, self._index_file_path)
+
+ def close(self):
+ self.flush()
+
+ def rebuild_idx(self, model):
+ pass
+
+ def count(self):
+ return self._index.ntotal
diff --git a/modelcache/manager/vector_data/manager.py b/modelcache_mm/manager/vector_data/manager.py
similarity index 61%
rename from modelcache/manager/vector_data/manager.py
rename to modelcache_mm/manager/vector_data/manager.py
index 54f7c55..64cca45 100644
--- a/modelcache/manager/vector_data/manager.py
+++ b/modelcache_mm/manager/vector_data/manager.py
@@ -1,8 +1,8 @@
# -*- coding: utf-8 -*-
-from modelcache.utils.error import NotFoundError, ParamError
+from modelcache_mm.utils.error import NotFoundError, ParamError
TOP_K = 1
-FAISS_INDEX_PATH = "faiss.index"
+FAISS_INDEX_PATH = "mm_faiss.index"
DIMENSION = 0
MILVUS_HOST = "localhost"
MILVUS_PORT = 19530
@@ -40,8 +40,8 @@ def get(name, **kwargs):
top_k = kwargs.get("top_k", TOP_K)
if name == "milvus":
from modelcache.manager.vector_data.milvus import Milvus
- dimension = kwargs.get("dimension", DIMENSION)
milvus_config = kwargs.get("milvus_config")
+ dimension = kwargs.get("dimension", DIMENSION)
VectorBase.check_dimension(dimension)
host = milvus_config.get('milvus', 'host')
port = milvus_config.get('milvus', 'port')
@@ -68,38 +68,55 @@ def get(name, **kwargs):
local_mode=local_mode,
local_data=local_data
)
- elif name == "faiss":
- from modelcache.manager.vector_data.faiss import Faiss
+ elif name == "redis":
+ from modelcache_mm.manager.vector_data.redis import RedisVectorStore
+ redis_config = kwargs.get("redis_config")
+ mm_dimension = kwargs.get("mm_dimension", DIMENSION)
+ i_dimension = kwargs.get("i_dimension", DIMENSION)
+ t_dimension = kwargs.get("t_dimension", DIMENSION)
+ VectorBase.check_dimension(mm_dimension)
+ VectorBase.check_dimension(i_dimension)
+ VectorBase.check_dimension(t_dimension)
+
+ host = redis_config.get('redis', 'host')
+ port = redis_config.get('redis', 'port')
+ user = redis_config.get('redis', 'user')
+ password = redis_config.get('redis', 'password')
+ namespace = kwargs.get("namespace", "")
+ # collection_name = kwargs.get("collection_name", COLLECTION_NAME)
+
+ vector_base = RedisVectorStore(
+ host=host,
+ port=port,
+ username=user,
+ password=password,
+ namespace=namespace,
+ top_k=top_k,
+ mm_dimension=mm_dimension,
+ i_dimension=i_dimension,
+ t_dimension=t_dimension,
+ )
+ elif name == "faiss":
+ from modelcache_mm.manager.vector_data.faiss import Faiss
dimension = kwargs.get("dimension", DIMENSION)
- index_path = kwargs.pop("index_path", FAISS_INDEX_PATH)
VectorBase.check_dimension(dimension)
+
+ index_path = kwargs.pop("index_path", FAISS_INDEX_PATH)
vector_base = Faiss(
- index_file_path=index_path, dimension=dimension, top_k=top_k
+ index_file_path=index_path,
+ dimension=dimension,
+ top_k=top_k
)
elif name == "chromadb":
- from modelcache.manager.vector_data.chroma import Chromadb
+ from modelcache_mm.manager.vector_data.chroma import Chromadb
- client_settings = kwargs.get("client_settings", None)
- persist_directory = kwargs.get("persist_directory", None)
- collection_name = kwargs.get("collection_name", COLLECTION_NAME)
+ chromadb_config = kwargs.get("chromadb_config", None)
+ persist_directory = chromadb_config.get('chromadb', 'persist_directory')
vector_base = Chromadb(
- client_settings=client_settings,
persist_directory=persist_directory,
- collection_name=collection_name,
top_k=top_k,
)
- elif name == "hnswlib":
- from modelcache.manager.vector_data.hnswlib_store import Hnswlib
-
- dimension = kwargs.get("dimension", DIMENSION)
- index_path = kwargs.pop("index_path", "./hnswlib_index.bin")
- max_elements = kwargs.pop("max_elements", 100000)
- VectorBase.check_dimension(dimension)
- vector_base = Hnswlib(
- index_file_path=index_path, dimension=dimension,
- top_k=top_k, max_elements=max_elements
- )
else:
raise NotFoundError("vector store", name)
return vector_base
diff --git a/modelcache_mm/manager/vector_data/redis.py b/modelcache_mm/manager/vector_data/redis.py
new file mode 100644
index 0000000..38ccb88
--- /dev/null
+++ b/modelcache_mm/manager/vector_data/redis.py
@@ -0,0 +1,159 @@
+# -*- coding: utf-8 -*-
+from typing import List
+import numpy as np
+from redis.commands.search.indexDefinition import IndexDefinition, IndexType
+from redis.commands.search.query import Query
+from redis.commands.search.field import VectorField, NumericField
+from redis.client import Redis
+
+from modelcache_mm.manager.vector_data.base import VectorBase, VectorData
+from modelcache_mm.utils import import_redis
+from modelcache_mm.utils.log import modelcache_log
+from modelcache_mm.utils.index_util import get_mm_index_name
+from modelcache_mm.utils.index_util import get_mm_index_prefix
+import_redis()
+
+
+class RedisVectorStore(VectorBase):
+ def __init__(
+ self,
+ host: str = "localhost",
+ port: str = "6379",
+ username: str = "",
+ password: str = "",
+ mm_dimension: int = 0,
+ i_dimension: int = 0,
+ t_dimension: int = 0,
+ top_k: int = 1,
+ namespace: str = "",
+ ):
+ if mm_dimension <= 0:
+ raise ValueError(
+ 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.mm_dimension = mm_dimension
+ self.i_dimension = i_dimension
+ self.t_dimension = t_dimension
+ self.namespace = namespace
+ self.doc_prefix = f"{self.namespace}doc:"
+
+ def _check_index_exists(self, index_name: str) -> bool:
+ """Check if Redis index exists."""
+ try:
+ self._client.ft(index_name).info()
+ except:
+ modelcache_log.info("Index does not exist")
+ return False
+ modelcache_log.info("Index already exists")
+ return True
+
+ def create_index(self, index_name, type, index_prefix):
+ # dimension = self.dimension
+ if type == 'IMG_TEXT':
+ dimension = self.mm_dimension
+ elif type == 'IMG':
+ dimension = self.i_dimension
+ elif type == 'TEXT':
+ dimension = self.t_dimension
+ else:
+ raise ValueError('dimension type exception')
+ if self._check_index_exists(index_name):
+ modelcache_log.info(
+ "The %s already exists, and it will be used directly", index_name
+ )
+ return 'already_exists'
+ else:
+ id_field_name = "data_id"
+ embedding_field_name = "data_vector"
+
+ id = NumericField(name=id_field_name)
+ embedding = VectorField(embedding_field_name,
+ "HNSW", {
+ "TYPE": "FLOAT32",
+ "DIM": dimension,
+ "DISTANCE_METRIC": "L2",
+ "INITIAL_CAP": 1000,
+ }
+ )
+ fields = [id, embedding]
+ definition = IndexDefinition(prefix=[index_prefix], index_type=IndexType.HASH)
+
+ # create Index
+ self._client.ft(index_name).create_index(
+ fields=fields, definition=definition
+ )
+ return 'create_success'
+
+ def 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()
+ index_prefix = get_mm_index_prefix(model, mm_type)
+ id_field_name = "data_id"
+ embedding_field_name = "data_vector"
+ obj = {id_field_name: id, embedding_field_name: embedding}
+ self._client.hset(f"{index_prefix}{id}", mapping=obj)
+
+ def search(self, data: np.ndarray, top_k: int = -1, model=None, mm_type=None):
+ index_name = get_mm_index_name(model, mm_type)
+ id_field_name = "data_id"
+ embedding_field_name = "data_vector"
+
+ base_query = f'*=>[KNN 2 @{embedding_field_name} $vector AS distance]'
+ query = (
+ Query(base_query)
+ .sort_by("distance")
+ .return_fields(id_field_name, "distance")
+ .dialect(2)
+ )
+ query_params = {"vector": data.astype(np.float32).tobytes()}
+ results = (
+ self._client.ft(index_name)
+ .search(query, query_params=query_params)
+ .docs
+ )
+ return [(float(result.distance), int(getattr(result, id_field_name))) for result in results]
+
+ def create(self, model=None, mm_type=None):
+ collection_name_model = get_mm_index_name(model, mm_type)
+ try:
+ index_prefix = get_mm_index_prefix(model, mm_type)
+ self.create_index(collection_name_model, mm_type, index_prefix)
+ except Exception as e:
+ raise ValueError(str(e))
+ return 'success'
+
+ def rebuild(self, ids=None) -> bool:
+ pass
+
+ def rebuild_idx(self, model, mm_type=None):
+ for mm_type in ['IMG_TEXT', 'TEXT']:
+ index_name = get_mm_index_name(model, mm_type)
+ if self._check_index_exists(index_name):
+ try:
+ self._client.ft(index_name).dropindex(delete_documents=True)
+ except Exception as e:
+ raise ValueError(str(e))
+ try:
+ index_prefix = get_mm_index_prefix(model, mm_type)
+ self.create_index(index_name, mm_type, index_prefix)
+ except Exception as e:
+ raise ValueError(str(e))
+
+ def delete(self, ids) -> None:
+ pipe = self._client.pipeline()
+ for data_id in ids:
+ pipe.delete(f"{self.doc_prefix}{data_id}")
+ pipe.execute()
+
+ def create(self, model=None, type=None):
+ index_name = get_mm_index_name(model, type)
+ index_prefix = get_mm_index_prefix(model, type)
+ return self.create_index(index_name, type, index_prefix)
+
+ def get_index_by_name(self, index_name):
+ pass
diff --git a/modelcache_mm/processor/__init__.py b/modelcache_mm/processor/__init__.py
new file mode 100644
index 0000000..40a96af
--- /dev/null
+++ b/modelcache_mm/processor/__init__.py
@@ -0,0 +1 @@
+# -*- coding: utf-8 -*-
diff --git a/modelcache_mm/processor/post.py b/modelcache_mm/processor/post.py
new file mode 100644
index 0000000..f92ddc5
--- /dev/null
+++ b/modelcache_mm/processor/post.py
@@ -0,0 +1,15 @@
+# -*- coding: utf-8 -*-
+import random
+from typing import List, Any
+
+
+def random_one(messages: List[Any]) -> Any:
+ return random.choice(messages)
+
+
+def first(messages: List[Any]) -> Any:
+ return messages[0]
+
+
+def nop(messages: List[Any]) -> Any:
+ return messages
diff --git a/modelcache_mm/processor/pre.py b/modelcache_mm/processor/pre.py
new file mode 100644
index 0000000..f65515d
--- /dev/null
+++ b/modelcache_mm/processor/pre.py
@@ -0,0 +1,115 @@
+# -*- coding: utf-8 -*-
+import re
+from typing import Dict, Any
+
+
+def insert_last_content(data: Dict[str, Any], **_: Dict[str, Any]) -> Any:
+ return data.get("chat_info")[-1]["query"]
+
+
+def query_last_content(data: Dict[str, Any], **_: Dict[str, Any]) -> Any:
+ return data.get("query")[-1]["content"]
+
+
+def last_content_without_prompt(data: Dict[str, Any], **params: Dict[str, Any]) -> Any:
+ last_content_str = data.get("messages")[-1]["content"]
+ prompts = params.get("prompts", [])
+ if prompts is None:
+ return last_content_str
+ pattern = "|".join(prompts)
+ new_content_str = re.sub(pattern, "", last_content_str)
+ return new_content_str
+
+
+def all_content(data: Dict[str, Any], **_: Dict[str, Any]) -> Any:
+ s = ""
+ messages = data.get("messages")
+ for i, message in enumerate(messages):
+ if i == len(messages) - 1:
+ s += message["content"]
+ else:
+ s += message["content"] + "\n"
+ return s
+
+
+def nop(data: Dict[str, Any], **_: Dict[str, Any]) -> Any:
+ return data
+
+
+def get_prompt(data: Dict[str, Any], **_: Dict[str, Any]) -> Any:
+ return data.get("prompt")
+
+
+def get_file_name(data: Dict[str, Any], **_: Dict[str, Any]) -> str:
+ return data.get("file").name
+
+
+def get_file_bytes(data: Dict[str, Any], **_: Dict[str, Any]) -> bytes:
+ return data.get("file").peek()
+
+
+def get_input_str(data: Dict[str, Any], **_: Dict[str, Any]) -> str:
+ input_data = data.get("input")
+ return str(input_data["image"].peek()) + input_data["question"]
+
+
+def get_input_image_file_name(data: Dict[str, Any], **_: Dict[str, Any]) -> str:
+ input_data = data.get("input")
+ return input_data["image"].name
+
+
+def query_multi_splicing(data: Dict[str, Any], **_: Dict[str, Any]) -> Any:
+ query_list = data.get("query")
+ return multi_splicing(query_list)
+
+
+def insert_multi_splicing(data: Dict[str, Any], **_: Dict[str, Any]) -> Any:
+ insert_query_list = data.get("chat_info")[-1]['query']
+ return multi_splicing(insert_query_list)
+
+
+def multi_splicing(data_list) -> Any:
+ result_str = ""
+ for d in data_list:
+ role = d.get('role', '')
+ content = d.get('content', '')
+ result_str += role + "###" + content + "|||"
+
+ # 去掉最后一个"|||"
+ result_str = result_str[:-3]
+
+ return result_str
+
+
+def multi_analysis(dialog_str):
+ sub_strings = dialog_str.split('|||')
+ dict_list = []
+ for s in sub_strings:
+ parts = s.split('###')
+ if len(parts) == 2:
+ role = parts[0]
+ content = parts[1]
+ elif len(parts) > 2:
+ role = parts[0]
+ content = '###'.join(parts[1:])
+ else:
+ content = 'exception'
+ if content == '':
+ d = {"role": role}
+ else:
+ d = {"role": role, "content": content}
+ dict_list.append(d)
+ result_list = dict_list
+
+ # 输出结果
+ return result_list
+
+
+def mm_insert_dict(data: Dict[str, Any], **_: Dict[str, Any]) -> Any:
+ query_dict = data.get("chat_info")[-1]['query']
+ return query_dict
+
+
+def mm_query_dict(data: Dict[str, Any], **_: Dict[str, Any]) -> Any:
+ query_dict = data.get("query")
+ return query_dict
diff --git a/modelcache_mm/report.py b/modelcache_mm/report.py
new file mode 100644
index 0000000..5ac4ce8
--- /dev/null
+++ b/modelcache_mm/report.py
@@ -0,0 +1,44 @@
+# -*- coding: utf-8 -*-
+class Report:
+ def __init__(self):
+ self.embedding_all_time = 0
+ self.embedding_count = 0
+ self.search_all_time = 0
+ self.search_count = 0
+ self.hint_cache_count = 0
+
+ def embedding(self, delta_time):
+ """Embedding counts and time.
+
+ :param delta_time: additional runtime.
+ """
+ self.embedding_all_time += delta_time
+ self.embedding_count += 1
+
+ def search(self, delta_time):
+ """Search counts and time.
+
+ :param delta_time: additional runtime.
+ """
+ self.search_all_time += delta_time
+ self.search_count += 1
+
+ def average_embedding_time(self):
+ """Average embedding time."""
+ return round(
+ self.embedding_all_time / self.embedding_count
+ if self.embedding_count != 0
+ else 0,
+ 4,
+ )
+
+ def average_search_time(self):
+ return round(
+ self.search_all_time / self.search_count
+ if self.embedding_count != 0
+ else 0,
+ 4,
+ )
+
+ def hint_cache(self):
+ self.hint_cache_count += 1
diff --git a/modelcache_mm/similarity_evaluation/__init__.py b/modelcache_mm/similarity_evaluation/__init__.py
new file mode 100644
index 0000000..22dff8f
--- /dev/null
+++ b/modelcache_mm/similarity_evaluation/__init__.py
@@ -0,0 +1,11 @@
+# -*- coding: utf-8 -*-
+from modelcache.similarity_evaluation.similarity_evaluation import SimilarityEvaluation
+from modelcache.utils.lazy_import import LazyImport
+
+exact_match = LazyImport(
+ "exact_match", globals(), "modelcache.similarity_evaluation.exact_match"
+)
+
+
+def ExactMatchEvaluation():
+ return exact_match.ExactMatchEvaluation()
diff --git a/modelcache_mm/similarity_evaluation/distance.py b/modelcache_mm/similarity_evaluation/distance.py
new file mode 100644
index 0000000..44ca595
--- /dev/null
+++ b/modelcache_mm/similarity_evaluation/distance.py
@@ -0,0 +1,24 @@
+# -*- coding: utf-8 -*-
+from typing import Tuple, Dict, Any
+from modelcache.similarity_evaluation import SimilarityEvaluation
+
+
+class SearchDistanceEvaluation(SimilarityEvaluation):
+ def __init__(self, max_distance=4.0, positive=False):
+ self.max_distance = max_distance
+ self.positive = positive
+
+ def evaluation(
+ self, src_dict: Dict[str, Any], cache_dict: Dict[str, Any], **_
+ ) -> float:
+ distance, _ = cache_dict["search_result"]
+ if distance < 0:
+ distance = 0
+ elif distance > self.max_distance:
+ distance = self.max_distance
+ if self.positive:
+ return distance
+ return self.max_distance - distance
+
+ def range(self) -> Tuple[float, float]:
+ return 0.0, self.max_distance
diff --git a/modelcache_mm/similarity_evaluation/exact_match.py b/modelcache_mm/similarity_evaluation/exact_match.py
new file mode 100644
index 0000000..2dcb554
--- /dev/null
+++ b/modelcache_mm/similarity_evaluation/exact_match.py
@@ -0,0 +1,16 @@
+# -*- coding: utf-8 -*-
+from typing import Tuple, Dict, Any
+from modelcache.similarity_evaluation.similarity_evaluation import SimilarityEvaluation
+
+
+class ExactMatchEvaluation(SimilarityEvaluation):
+ def __init__(self):
+ pass
+
+ def evaluation(
+ self, src_dict: Dict[str, Any], cache_dict: Dict[str, Any], **_
+ ) -> float:
+ return 1 if cache_dict["question"] == src_dict["question"] else 0
+
+ def range(self) -> Tuple[float, float]:
+ return 0, 1
diff --git a/modelcache_mm/similarity_evaluation/similarity_evaluation.py b/modelcache_mm/similarity_evaluation/similarity_evaluation.py
new file mode 100644
index 0000000..79f4294
--- /dev/null
+++ b/modelcache_mm/similarity_evaluation/similarity_evaluation.py
@@ -0,0 +1,15 @@
+# -*- coding: utf-8 -*-
+from abc import ABCMeta, abstractmethod
+from typing import Tuple, Dict, Any
+
+
+class SimilarityEvaluation(metaclass=ABCMeta):
+ @abstractmethod
+ def evaluation(
+ self, src_dict: Dict[str, Any], cache_dict: Dict[str, Any], **kwargs
+ ) -> float:
+ pass
+
+ @abstractmethod
+ def range(self) -> Tuple[float, float]:
+ pass
diff --git a/modelcache_mm/utils/__init__.py b/modelcache_mm/utils/__init__.py
new file mode 100644
index 0000000..15d90ec
--- /dev/null
+++ b/modelcache_mm/utils/__init__.py
@@ -0,0 +1,79 @@
+# -*- coding: utf-8 -*-
+import importlib.util
+from typing import Optional
+from modelcache.utils.dependency_control import prompt_install
+
+
+def _check_library(libname: str, prompt: bool = True, package: Optional[str] = None):
+ is_avail = False
+ if importlib.util.find_spec(libname):
+ is_avail = True
+ if not is_avail and prompt:
+ prompt_install(package if package else libname)
+ return is_avail
+
+
+def import_onnxruntime():
+ _check_library("onnxruntime")
+
+
+def import_huggingface():
+ _check_library("transformers")
+
+
+def import_huggingface_hub():
+ _check_library("huggingface_hub", package="huggingface-hub")
+
+
+def import_pymysql():
+ _check_library("pymysql")
+
+
+def import_sql_client(db_name):
+ if db_name in ["mysql"]:
+ import_pymysql()
+
+
+def import_pymilvus():
+ _check_library("pymilvus")
+
+
+def import_milvus_lite():
+ _check_library("milvus")
+
+
+def import_faiss():
+ _check_library("faiss", package="faiss-cpu")
+
+
+def import_torch():
+ _check_library("torch")
+
+
+def import_fasttext():
+ _check_library("fasttext")
+
+
+def import_paddle():
+ prompt_install("protobuf==3.20.0")
+ _check_library("paddlepaddle")
+
+
+def import_paddlenlp():
+ _check_library("paddlenlp")
+
+
+def import_timm():
+ _check_library("timm", package="timm")
+
+
+def import_pillow():
+ _check_library("PIL", package="pillow")
+
+
+def import_redis():
+ _check_library("redis")
+
+
+def import_chromadb():
+ _check_library("chromadb", package="chromadb")
\ No newline at end of file
diff --git a/modelcache/utils/cache_func.py b/modelcache_mm/utils/cache_func.py
similarity index 100%
rename from modelcache/utils/cache_func.py
rename to modelcache_mm/utils/cache_func.py
diff --git a/modelcache_mm/utils/dependency_control.py b/modelcache_mm/utils/dependency_control.py
new file mode 100644
index 0000000..0f94355
--- /dev/null
+++ b/modelcache_mm/utils/dependency_control.py
@@ -0,0 +1,18 @@
+# -*- coding: utf-8 -*-
+import subprocess
+from modelcache.utils.error import PipInstallError
+from modelcache.utils.log import modelcache_log
+
+
+def prompt_install(package: str, warn: bool = False): # pragma: no cover
+ """
+ Function used to prompt user to install a package.
+ """
+ cmd = f"pip install {package}"
+ try:
+ if warn and input(f"Install {package}? Y/n: ") != "Y":
+ raise ModuleNotFoundError(f"No module named {package}")
+ subprocess.check_call(cmd, shell=True)
+ modelcache_log.info("%s installed successfully!", package)
+ except subprocess.CalledProcessError as e:
+ raise PipInstallError(package) from e
diff --git a/modelcache_mm/utils/env_config.py b/modelcache_mm/utils/env_config.py
new file mode 100644
index 0000000..40a96af
--- /dev/null
+++ b/modelcache_mm/utils/env_config.py
@@ -0,0 +1 @@
+# -*- coding: utf-8 -*-
diff --git a/modelcache_mm/utils/error.py b/modelcache_mm/utils/error.py
new file mode 100644
index 0000000..9022f00
--- /dev/null
+++ b/modelcache_mm/utils/error.py
@@ -0,0 +1,35 @@
+# -*- coding: utf-8 -*-
+class CacheError(Exception):
+ """ModelCache base error"""
+
+
+class NotInitError(CacheError):
+ """Raise when the cache has been used before it's inited"""
+ def __init__(self):
+ super().__init__("The cache should be inited before using")
+
+
+class RemoveError(CacheError):
+ """Raise when the cache has been used before it's inited"""
+ def __init__(self):
+ super().__init__("The cache remove error")
+
+class NotFoundError(CacheError):
+ """Raise when getting an unsupported store."""
+ def __init__(self, store_type, current_type_name):
+ super().__init__(f"Unsupported ${store_type}: {current_type_name}")
+
+
+class ParamError(CacheError):
+ """Raise when receiving an invalid param."""
+
+
+class PipInstallError(CacheError):
+ """Raise when failed to install package."""
+ def __init__(self, package):
+ super().__init__(f"Ran into error installing {package}.")
+
+
+class MultiTypeError(CacheError):
+ def __init__(self):
+ super().__init__("multichat type error, please check")
diff --git a/modelcache_mm/utils/index_util.py b/modelcache_mm/utils/index_util.py
new file mode 100644
index 0000000..efea863
--- /dev/null
+++ b/modelcache_mm/utils/index_util.py
@@ -0,0 +1,33 @@
+# -*- coding: utf-8 -*-
+
+
+def get_index_name(model):
+ return 'multicache' + '_' + model
+
+
+def get_index_prefix(model):
+ return 'prefix' + '_' + model
+
+
+def get_mm_index_name(model, mm_type):
+ if mm_type not in ['IMG_TEXT', 'mm', 'IMG', 'image', 'TEXT', 'text']:
+ raise ValueError('mm_type is not normal!')
+ if mm_type == 'IMG_TEXT':
+ mm_type = 'mm'
+ elif mm_type == 'IMG':
+ mm_type = 'image'
+ elif mm_type == 'TEXT':
+ mm_type = 'text'
+ return 'multi_cache' + '_' + model + '_' + mm_type
+
+
+def get_mm_index_prefix(model, mm_type):
+ if mm_type not in ['IMG_TEXT', 'mm', 'IMG', 'image', 'TEXT', 'text']:
+ raise ValueError('mm_type is not normal!')
+ if mm_type == 'IMG_TEXT':
+ mm_type = 'mm'
+ elif mm_type == 'IMG':
+ mm_type = 'image'
+ elif mm_type == 'TEXT':
+ mm_type = 'text'
+ return 'prefix' + '_' + model + '_' + mm_type
diff --git a/modelcache_mm/utils/lazy_import.py b/modelcache_mm/utils/lazy_import.py
new file mode 100644
index 0000000..3f24eec
--- /dev/null
+++ b/modelcache_mm/utils/lazy_import.py
@@ -0,0 +1,27 @@
+# -*- coding: utf-8 -*-
+import importlib
+from types import ModuleType
+
+
+class LazyImport(ModuleType):
+ """
+ Lazily import a module.
+ """
+ def __init__(self, local_name, parent_module_globals, name):
+ self._local_name = local_name
+ self._parent_module_globals = parent_module_globals
+ super().__init__(name)
+
+ def _load(self):
+ module = importlib.import_module(self.__name__)
+ self._parent_module_globals[self._local_name] = module
+ self.__dict__.update(module.__dict__)
+ return module
+
+ def __getattr__(self, item):
+ module = self._load()
+ return getattr(module, item)
+
+ def __dir__(self):
+ module = self._load()
+ return dir(module)
diff --git a/modelcache_mm/utils/log.py b/modelcache_mm/utils/log.py
new file mode 100644
index 0000000..68f4bd4
--- /dev/null
+++ b/modelcache_mm/utils/log.py
@@ -0,0 +1,7 @@
+# -*- coding: utf-8 -*-
+import logging
+
+FORMAT = '%(asctime)s - %(thread)d - %(filename)s-%(module)s:%(lineno)s - %(levelname)s: %(message)s'
+logging.basicConfig(format=FORMAT)
+
+modelcache_log = logging.getLogger('modelcache')
diff --git a/modelcache_mm/utils/time.py b/modelcache_mm/utils/time.py
new file mode 100644
index 0000000..7074b58
--- /dev/null
+++ b/modelcache_mm/utils/time.py
@@ -0,0 +1,20 @@
+# -*- coding: utf-8 -*-
+import time
+from modelcache import cache
+
+
+def time_cal(func, func_name=None, report_func=None):
+ def inner(*args, **kwargs):
+ time_start = time.time()
+ res = func(*args, **kwargs)
+ delta_time = time.time() - time_start
+ if cache.config.log_time_func:
+ cache.config.log_time_func(
+ func.__name__ if func_name is None else func_name, delta_time
+ )
+ if report_func is not None:
+ report_func(delta_time)
+ return res
+
+ return inner
+
diff --git a/mulicache-readme-cn.md b/mulicache-readme-cn.md
new file mode 100644
index 0000000..9d21d83
--- /dev/null
+++ b/mulicache-readme-cn.md
@@ -0,0 +1,127 @@
+# MultiModal Cache
+
+为满足多模态的性能要求，我们在 LLModel Cache 的基础上，开发了 MultiModal Cache 系统。MultiModal Cache 增强了 ModelCache 功能，架优化架构，适应多种应用场景。
+
+- [MultiModal Cache](#multimodal-cache)
+ - [最新动态](#最新动态)
+ - [特性](#特性)
+ - [性能](#性能)
+ - [效果评估](#效果评估)
+ - [参与贡献](#参与贡献)
+
+## 最新动态
+
+- [2024.12.12] MultiModal Cache 系统正式发布。
+
+## 特性
+
+| 场景 | 数据类型 | 图像格式 | 数据隔离 |
+|------|----------|----------|----------|
+| 文本对话 | 文本 | 不适用 | 支持 |
+| 图文理解 | 文本+图像 | image_url/image_base64 | 支持 |
+
+- **兼容性**：支持文本和图片链接（image_url）和图片 Base64 编码三种数据格式及其组合。
+- **数据隔离**：支持多模型数据隔离，允许不同数据模型在同一系统中独立运行。
+- **模态隔离**：支持同一模型下不同模态数据（如文本和图像）的隔离处理。
+
+## 性能
+
+我们在生产环境中使用企业级数据库对 MultiModal Cache 进行了全面的性能评估。以下是详细的性能数据：
+
+