添加可以指定 任务类型 task_type

Author: shell-nlp

gpt_server/model_worker/base/model_worker_base.py

@@ -249,6 +249,8 @@ def run(cls):
         parser.add_argument("--punc_model", type=str, default="")
         # log_level
         parser.add_argument("--log_level", type=str, default="WARNING")
+        # task_type
+        parser.add_argument("--task_type", type=str, default="auto")
         args = parser.parse_args()
         os.environ["num_gpus"] = str(args.num_gpus)
         if args.backend == "vllm":
@@ -276,6 +278,7 @@ def run(cls):
         os.environ["kv_cache_quant_policy"] = args.kv_cache_quant_policy
         os.environ["dtype"] = args.dtype
         os.environ["log_level"] = args.log_level
+        os.environ["task_type"] = args.task_type
         logger.remove(0)
         log_level = os.getenv("log_level", "WARNING")
         logger.add(sys.stderr, level=log_level)

gpt_server/model_worker/utils.py

@@ -3,7 +3,7 @@
 from fastapi import HTTPException
 import base64
 import io
-
+import os
 from PIL.Image import Image
 
 
@@ -53,6 +53,14 @@ def get_embedding_mode(model_path: str):
     from transformers import AutoConfig
     from infinity_emb.inference.select_model import get_engine_type_from_config
 
+    task_type = os.environ.get("task_type", "auto")
+    if task_type == "embedding":
+        return "embedding"
+    elif task_type == "reranker":
+        return "rerank"
+    elif task_type == "classify":
+        return "classify"
+
     model_config = AutoConfig.from_pretrained(model_path, trust_remote_code=True)
     architectures = getattr(model_config, "architectures", [])
     if "JinaVLForRanking" in architectures:

gpt_server/script/config_example.yaml

@@ -94,11 +94,12 @@ models:
       - 2
 
 - jina-reranker:
-    # 多模态多语言的重排模型
+    # 多模态多语言的重排模型，这个模型task_type 只能是 auto
     alias: null
     enable: true
     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
     device: gpu
@@ -112,6 +113,7 @@ models:
     enable: true # false true
     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
     device: gpu # gpu / cpu

gpt_server/utils.py

@@ -167,6 +167,7 @@ def start_model_worker(config: dict):
                     )
                     vad_model = engine_config.get("vad_model", "")
                     punc_model = engine_config.get("punc_model", "")
+                    task_type = engine_config.get("task_type", "auto")
 
                 else:
                     logger.error(
@@ -252,6 +253,7 @@ def start_model_worker(config: dict):
                         + f" --gpu_memory_utilization {gpu_memory_utilization}"  # 占用GPU比例
                         + f" --kv_cache_quant_policy {kv_cache_quant_policy}"  # kv cache 量化策略
                         + f" --log_level {log_level}"  # 日志水平
+                        + f" --task_type {task_type}"  # 日志水平
                     )
                     # 处理为 None的情况
                     if lora: