New feature: Embedding Dispatcher for parallel embedding

Based on the multiprocessing module, provides true parallel embedding.
The number of workers can be adjusted in the cache initialization.

Access to the caching logic is now done asynchronously using asyncio module

Author: Yuval-Roth

fastapi4modelcache.py

@@ -1,61 +1,46 @@
 # -*- coding: utf-8 -*-
+import asyncio
+from contextlib import asynccontextmanager
 import uvicorn
 import json
-from fastapi import FastAPI, Request, HTTPException
+from fastapi.responses import JSONResponse
+from fastapi import FastAPI, Request
 from modelcache.cache import Cache
-
-#创建一个FastAPI实例
-app = FastAPI()
-
-cache = Cache.init("mysql", "milvus")
+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:
-        raw_body = await request.body()
-        # 解析字符串为JSON对象
-        if isinstance(raw_body, bytes):
-            raw_body = raw_body.decode("utf-8")
-        if isinstance(raw_body, str):
-            try:
-                # 尝试将字符串解析为JSON对象
-                request_data = json.loads(raw_body)
-            except json.JSONDecodeError as e:
-                # 如果无法解析，返回格式错误
-                result = {"errorCode": 101, "errorDesc": str(e), "cacheHit": False, "delta_time": 0, "hit_query": '',
-                  "answer": ''}
-                cache.save_query_info(result, model='', query='', delta_time_log=0)
-                raise HTTPException(status_code=101, detail="Invalid JSON format")
-        else:
-            request_data = raw_body
 
-        # 确保request_data是字典对象
-        if isinstance(request_data, str):
-            try:
-                request_data = json.loads(request_data)
-            except json.JSONDecodeError:
-                raise HTTPException(status_code=101, detail="Invalid JSON format")
-
-        return cache.handle_request(request_data)
+    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:
-        request_data = raw_body if 'raw_body' in locals() else None
-        result = {
-            "errorCode": 103,
-            "errorDesc": str(e),
-            "cacheHit": False,
-            "delta_time": 0,
-            "hit_query": '',
-            "answer": '',
-            "para_dict": request_data
-        }
-        return 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 JSONResponse(status_code=500, content=result)
 
-# TODO: 可以修改为在命令行中使用`uvicorn your_module_name:app --host 0.0.0.0 --port 5000 --reload`的命令启动
 if __name__ == '__main__':
-    uvicorn.run(app, host='0.0.0.0', port=5000)
+    uvicorn.run(app, host='0.0.0.0', port=5000, loop="asyncio", http="httptools")

fastapi4modelcache_demo.py

@@ -1,61 +1,46 @@
 # -*- coding: utf-8 -*-
+import asyncio
+from contextlib import asynccontextmanager
 import uvicorn
 import json
-from fastapi import FastAPI, Request, HTTPException
-
+from fastapi.responses import JSONResponse
+from fastapi import FastAPI, Request
 from modelcache.cache import Cache
-
-# 创建一个FastAPI实例
-app = FastAPI()
-
-cache = Cache.init("sqlite", "faiss")
+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:
-        raw_body = await request.body()
-        # 解析字符串为JSON对象
-        if isinstance(raw_body, bytes):
-            raw_body = raw_body.decode("utf-8")
-        if isinstance(raw_body, str):
-            try:
-                # 尝试将字符串解析为JSON对象
-                request_data = json.loads(raw_body)
-            except json.JSONDecodeError as e:
-                # 如果无法解析，返回格式错误
-                result = {"errorCode": 101, "errorDesc": str(e), "cacheHit": False, "delta_time": 0, "hit_query": '',
-                          "answer": ''}
-                cache.save_query_info(result, model='', query='', delta_time_log=0)
-                raise HTTPException(status_code=101, detail="Invalid JSON format")
-        else:
-            request_data = raw_body
 
-        # 确保request_data是字典对象
-        if isinstance(request_data, str):
-            try:
-                request_data = json.loads(request_data)
-            except json.JSONDecodeError:
-                raise HTTPException(status_code=101, detail="Invalid JSON format")
-
-        return cache.handle_request(request_data)
+    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:
-        request_data = raw_body if 'raw_body' in locals() else None
-        result = {
-            "errorCode": 103,
-            "errorDesc": str(e),
-            "cacheHit": False,
-            "delta_time": 0,
-            "hit_query": '',
-            "answer": '',
-            "para_dict": request_data
-        }
-        return 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 JSONResponse(status_code=500, content=result)
 
-# TODO: 可以修改为在命令行中使用`uvicorn your_module_name:app --host 0.0.0.0 --port 5000 --reload`的命令启动
 if __name__ == '__main__':
-    uvicorn.run(app, host='0.0.0.0', port=5000)
+    uvicorn.run(app, host='0.0.0.0', port=5000, loop="asyncio", http="httptools")

flask4modelcache.py

@@ -1,34 +1,48 @@
 # -*- coding: utf-8 -*-
-from flask import Flask, request
-import json
+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():
 
-cache = Cache.init("mysql","milvus")
+    # 创建一个Flask实例
+    app = Flask(__name__)
 
-@app.route('/welcome')
-def first_flask():  # 视图函数
-    return 'hello, modelcache!'
+    cache,loop = await Cache.init(
+        sql_storage="mysql",
+        vector_storage="milvus",
+        embedding_model=EmbeddingModel.HUGGINGFACE_ALL_MPNET_BASE_V2,
+        embedding_workers_num=2
+    )
 
+    @app.route('/welcome')
+    def first_flask():  # 视图函数
+        return 'hello, modelcache!'
 
-@app.route('/modelcache', methods=['GET', 'POST'])
-def user_backend():
-    param_dict = {}
-    try:
-        if request.method == 'POST':
+
+    @app.post('/modelcache')
+    def user_backend():
+        try:
             param_dict = request.json
-        elif request.method == 'GET':
-            param_dict = request.args
+        except Exception:
+            result = {"errorCode": 400, "errorDesc": "bad request", "cacheHit": False, "delta_time": 0, "hit_query": '',"answer": ''}
+            return jsonify(result), 400
+
+        try:
+            result = asyncio.run_coroutine_threadsafe(
+                cache.handle_request(param_dict), loop
+            ).result()
+            return jsonify(result), 200
+        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 jsonify(result), 500
 
-        return json.dumps(cache.handle_request(param_dict))
-    except Exception as e:
-        result = {"errorCode": 101, "errorDesc": str(e), "cacheHit": False, "delta_time": 0, "hit_query": '',
-                  "answer": ''}
-        cache.save_query_resp(result, model='', query='', delta_time=0)
-        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)
+    asyncio.run(main())

flask4modelcache_demo.py

@@ -1,34 +1,48 @@
 # -*- coding: utf-8 -*-
-from flask import Flask, request
-import json
+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():
 
-cache = Cache.init("sqlite","faiss")
+    # 创建一个Flask实例
+    app = Flask(__name__)
 
-@app.route('/welcome')
-def first_flask():  # 视图函数
-    return 'hello, modelcache!'
+    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!'
 
-@app.route('/modelcache', methods=['GET', 'POST'])
-def user_backend():
-    param_dict = {}
-    try:
-        if request.method == 'POST':
+
+    @app.post('/modelcache')
+    def user_backend():
+        try:
             param_dict = request.json
-        elif request.method == 'GET':
-            param_dict = request.args
+        except Exception:
+            result = {"errorCode": 400, "errorDesc": "bad request", "cacheHit": False, "delta_time": 0, "hit_query": '',"answer": ''}
+            return jsonify(result), 400
+
+        try:
+            result = asyncio.run_coroutine_threadsafe(
+                cache.handle_request(param_dict), loop
+            ).result()
+            return jsonify(result), 200
+        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 jsonify(result), 500
 
-        return json.dumps(cache.handle_request(param_dict))
-    except Exception as e:
-        result = {"errorCode": 101, "errorDesc": str(e), "cacheHit": False, "delta_time": 0, "hit_query": '',
-                  "answer": ''}
-        cache.save_query_resp(result, model='', query='', delta_time=0)
-        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)
+    asyncio.run(main())

model/text2vec-base-chinese/logs.txt

@@ -1 +1,19 @@
-
+Epoch:0 Valid| corr: 0.794410
+Epoch:0 Valid| corr: 0.691819
+Epoch:1 Valid| corr: 0.722749
+Epoch:2 Valid| corr: 0.735054
+Epoch:3 Valid| corr: 0.738295
+Epoch:4 Valid| corr: 0.739411
+Test | corr: 0.679971
+Epoch:0 Valid| corr: 0.817416
+Epoch:1 Valid| corr: 0.832376
+Epoch:2 Valid| corr: 0.842308
+Epoch:3 Valid| corr: 0.843520
+Epoch:4 Valid| corr: 0.841837
+Test | corr: 0.793495
+Epoch:0 Valid| corr: 0.814648
+Epoch:1 Valid| corr: 0.831609
+Epoch:2 Valid| corr: 0.841678
+Epoch:3 Valid| corr: 0.842387
+Epoch:4 Valid| corr: 0.841435
+Test | corr: 0.794840

modelcache/adapter/adapter.py

@@ -10,11 +10,11 @@ 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
@@ -24,9 +24,9 @@ def cache_data_convert(cache_data, cache_query):
             return str(e)
 
     @classmethod
-    def create_insert(cls, *args, **kwargs):
+    async def create_insert(cls, *args, **kwargs):
         try:
-            return adapt_insert(
+            return await adapt_insert(
                 *args,
                 **kwargs
             )