优化 项目架构 以更好的支持 更多的 embedding 后端

Author: shell-nlp

gpt_server/model_worker/embedding.py renamed to gpt_server/model_worker/embedding_sentence_transformers.py

@@ -0,0 +1,82 @@
+import os
+from typing import List
+import asyncio
+from loguru import logger
+
+from infinity_emb import AsyncEngineArray, EngineArgs, AsyncEmbeddingEngine
+from infinity_emb.inference.select_model import get_engine_type_from_config
+from gpt_server.model_worker.base.model_worker_base import ModelWorkerBase
+from gpt_server.model_worker.utils import get_embedding_mode
+import torch
+import vllm
+from vllm import LLM
+
+label_to_category = {
+    "S": "sexual",
+    "H": "hate",
+    "HR": "harassment",
+    "SH": "self-harm",
+    "S3": "sexual/minors",
+    "H2": "hate/threatening",
+    "V2": "violence/graphic",
+    "V": "violence",
+    "OK": "OK",
+}
+
+
+class EmbeddingWorker(ModelWorkerBase):
+    def __init__(
+        self,
+        controller_addr: str,
+        worker_addr: str,
+        worker_id: str,
+        model_path: str,
+        model_names: List[str],
+        limit_worker_concurrency: int,
+        conv_template: str = None,  # type: ignore
+    ):
+        super().__init__(
+            controller_addr,
+            worker_addr,
+            worker_id,
+            model_path,
+            model_names,
+            limit_worker_concurrency,
+            conv_template,
+            model_type="embedding",
+        )
+        tensor_parallel_size = int(os.getenv("num_gpus", "1"))
+        max_model_len = os.getenv("max_model_len", None)
+        gpu_memory_utilization = float(os.getenv("gpu_memory_utilization", 0.8))
+        enable_prefix_caching = bool(os.getenv("enable_prefix_caching", False))
+
+        self.mode = get_embedding_mode(model_path=model_path)
+        self.engine = LLM(
+            model=model_path,
+            tensor_parallel_size=tensor_parallel_size,
+            max_model_len=max_model_len,
+            gpu_memory_utilization=gpu_memory_utilization,
+            enable_prefix_caching=enable_prefix_caching,
+        )
+
+        logger.warning(f"模型：{model_names[0]}")
+        logger.warning(f"正在使用 {self.mode} 模型...")
+
+    async def get_embeddings(self, params):
+        self.call_ct += 1
+        ret = {"embedding": [], "token_num": 0}
+        texts: list = params["input"]
+        if self.mode == "embedding":
+            usage = None
+            texts = list(map(lambda x: x.replace("\n", " "), texts))
+            # ----------
+            outputs = self.engine.embed(prompts=texts)
+            embedding = [o.outputs.embedding for o in outputs]
+
+        ret["embedding"] = embedding
+        ret["token_num"] = usage
+        return ret
+
+
+if __name__ == "__main__":
+    EmbeddingWorker.run()

gpt_server/model_worker/embedding_vllm.py

@@ -87,8 +87,8 @@ models:
     enable: true # false true
     model_config:
       model_name_or_path: /home/dev/model/Xorbits/bge-reranker-base/
-    model_type: embedding_infinity # embedding_infinity/embedding 
-    work_mode: hf
+    model_type: embedding
+    work_mode: infinity # 可选 ["vllm", "infinity", "sentence_transformers"]，但并不是所有后端都支持
     device: gpu # gpu / cpu
     workers:
     - gpus:
@@ -101,8 +101,8 @@ models:
     model_config:
       model_name_or_path: /home/dev/model/jinaai/jina-reranker-m0/
       task_type: auto # auto 、embedding 、 reranker 或者 classify 不设置这个参数，默认为 auto,自动识别可能会识别错误
-    model_type: embedding # 这里仅支持 embedding
-    work_mode: hf
+    model_type: embedding 
+    work_mode: sentence_transformers # 可选 ["vllm", "infinity", "sentence_transformers"]，但并不是所有后端都支持
     device: gpu
     workers:
     - gpus:
@@ -115,8 +115,8 @@ models:
     model_config:
       model_name_or_path: /home/dev/model/aspire/acge_text_embedding
       task_type: auto # auto 、embedding 、 reranker 或者 classify 不设置这个参数，默认为 auto,自动识别可能会识别错误
-    model_type: embedding_infinity # embedding_infinity/embedding 
-    work_mode: hf
+    model_type: embedding
+    work_mode: infinity # 可选 ["vllm", "infinity", "sentence_transformers"]，但并不是所有后端都支持
     device: gpu # gpu / cpu
     workers:
     - gpus:
@@ -128,8 +128,8 @@ models:
     enable: true
     model_config:
       model_name_or_path: /home/dev/model/BAAI/BGE-VL-base/
-    model_type: embedding # 这里仅支持 embedding
-    work_mode: hf
+    model_type: embedding
+    work_mode: sentence_transformers # 可选 ["vllm", "infinity", "sentence_transformers"]，但并不是所有后端都支持
     device: gpu
     workers:
     - gpus:
@@ -141,8 +141,8 @@ models:
     enable: true
     model_config:
       model_name_or_path: /home/dev/model/KoalaAI/Text-Moderation
-    model_type: embedding_infinity # embedding_infinity
-    work_mode: hf
+    model_type: embedding
+    work_mode: infinity # 可选 ["vllm", "infinity", "sentence_transformers"]，但并不是所有后端都支持
     device: gpu
     workers:
     - gpus:

gpt_server/script/config_example.yaml

@@ -127,7 +127,6 @@ def parse_env_var(cls, field_name: str, raw_val: str):
 async def timing_tasks():
     """定时任务"""
     global model_address_map, models_
-    logger.info("定时任务已启动！")
     controller_address = app_settings.controller_address
 
     while True:

gpt_server/serving/openai_api_server.py

@@ -127,7 +127,8 @@ def get_model_types():
     return model_types
 
 
-model_types = get_model_types()
+model_types = get_model_types() + ["embedding"]
+embedding_backend_type = ["vllm", "infinity", "sentence_transformers"]
 
 
 def start_model_worker(config: dict):
@@ -201,7 +202,6 @@ def start_model_worker(config: dict):
                         f"不支持model_type: {model_type},仅支持{model_types}模型之一！"
                     )
                     sys.exit()
-                py_path = f"-m gpt_server.model_worker.{model_type}"
 
                 model_names = model_name
                 if model_config["alias"]:
@@ -240,7 +240,11 @@ def start_model_worker(config: dict):
                     else:
                         raise Exception("目前仅支持 CPU/GPU设备!")
                     backend = model_config["work_mode"]
+                    if model_type == "embedding":
+                        assert backend in embedding_backend_type
+                        model_type = f"embedding_{backend}"
 
+                    py_path = f"-m gpt_server.model_worker.{model_type}"
                     cmd = (
                         CUDA_VISIBLE_DEVICES
                         + run_mode