支持 qwen3 reranker 的 vllm 后端

Author: shell-nlp

gpt_server/model_worker/base/model_worker_base.py

@@ -311,6 +311,8 @@ def run(cls):
         parser.add_argument("--port", type=int, default=None)
         # model_type
         parser.add_argument("--model_type", type=str, default="auto")
+        # hf_overrides
+        parser.add_argument("--hf_overrides", type=str, default="")
         args = parser.parse_args()
         os.environ["num_gpus"] = str(args.num_gpus)
         if args.backend == "vllm":
@@ -332,6 +334,8 @@ def run(cls):
             os.environ["vad_model"] = args.vad_model
         if args.punc_model:
             os.environ["punc_model"] = args.punc_model
+        if args.hf_overrides:
+            os.environ["hf_overrides"] = args.hf_overrides
 
         os.environ["model_type"] = args.model_type
         os.environ["enable_prefix_caching"] = args.enable_prefix_caching

gpt_server/model_worker/embedding_vllm.py

@@ -1,14 +1,12 @@
 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 numpy as np
-from vllm import LLM
+from vllm import LLM, EmbeddingRequestOutput, ScoringRequestOutput
+from gpt_server.settings import get_model_config
 
 label_to_category = {
     "S": "sexual",
@@ -23,6 +21,24 @@
 }
 
 
+def template_format(queries: List[str], documents: List[str]):
+    model_config = get_model_config()
+    hf_overrides = model_config.hf_overrides
+    if hf_overrides:
+        if hf_overrides["architectures"][0] == "Qwen3ForSequenceClassification":
+            logger.info("使用 Qwen3ForSequenceClassification 模板格式化...")
+            prefix = '<|im_start|>system\nJudge whether the Document meets the requirements based on the Query and the Instruct provided. Note that the answer can only be "yes" or "no".<|im_end|>\n<|im_start|>user\n'
+            suffix = "<|im_end|>\n<|im_start|>assistant\n<think>\n\n</think>\n\n"
+            instruction = "Given a web search query, retrieve relevant passages that answer the query"
+
+            query_template = f"{prefix}<Instruct>: {instruction}\n<Query>: {{query}}\n"
+            document_template = f"<Document>: {{doc}}{suffix}"
+            queries = [query_template.format(query=query) for query in queries]
+            documents = [document_template.format(doc=doc) for doc in documents]
+            return queries, documents
+    return queries, documents
+
+
 class EmbeddingWorker(ModelWorkerBase):
     def __init__(
         self,
@@ -44,18 +60,20 @@ def __init__(
             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.6))
-        enable_prefix_caching = bool(os.getenv("enable_prefix_caching", False))
-
+        model_config = get_model_config()
+        hf_overrides = model_config.hf_overrides
         self.mode = get_embedding_mode(model_path=model_path)
+        runner = "auto"
+        if self.model == "rerank":
+            runner = "pooling"
         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,
+            tensor_parallel_size=model_config.num_gpus,
+            max_model_len=model_config.max_model_len,
+            gpu_memory_utilization=model_config.gpu_memory_utilization,
+            enable_prefix_caching=model_config.enable_prefix_caching,
+            runner=runner,
+            hf_overrides=hf_overrides,
         )
 
         logger.warning(f"模型：{model_names[0]}")
@@ -69,13 +87,20 @@ async def get_embeddings(self, params):
         if self.mode == "embedding":
             texts = list(map(lambda x: x.replace("\n", " "), texts))
             # ----------
-            outputs = self.engine.embed(texts)
+            outputs: list[EmbeddingRequestOutput] = self.engine.embed(texts)
             embedding = [o.outputs.embedding for o in outputs]
             embeddings_np = np.array(embedding)
             # ------ L2归一化（沿axis=1，即对每一行进行归一化）-------
             norm = np.linalg.norm(embeddings_np, ord=2, axis=1, keepdims=True)
             normalized_embeddings_np = embeddings_np / norm
             embedding = normalized_embeddings_np.tolist()
+        elif self.mode == "rerank":
+            query = params.get("query", None)
+            data_1 = [query] * len(texts)
+            data_2 = texts
+            data_1, data_2 = template_format(queries=data_1, documents=data_2)
+            scores: list[ScoringRequestOutput] = self.engine.score(data_1, data_2)
+            embedding = [[score.outputs.score] for score in scores]
 
         ret["embedding"] = embedding
         return ret

gpt_server/script/config_example.yaml

@@ -93,6 +93,21 @@ models:
     workers:
     - gpus:
       - 2
+# 部署 qwen3-reranker 样例
+- qwen3-reranker:
+    alias: null
+    enable: true
+    model_config:
+      model_name_or_path: /home/dev/model/Qwen/Qwen3-Reranker-0___6B/
+      dtype: auto
+      task_type: reranker
+      hf_overrides: { "architectures": [ "Qwen3ForSequenceClassification" ], "classifier_from_token": [ "no", "yes" ], "is_original_qwen3_reranker": True }
+    model_type: embedding
+    work_mode: vllm
+    device: gpu
+    workers:
+    - gpus:
+      - 6
 
 - jina-reranker:
     # 多模态多语言的重排模型，这个模型task_type 只能是 auto

gpt_server/settings.py

@@ -12,6 +12,8 @@ class ModelConfig(BaseSettings):
     dtype: str = "auto"
     num_gpus: int = 1
     lora: str | None = None
+    hf_overrides: dict | None = None
+    """HuggingFace 配置覆盖参数"""
 
 
 def get_model_config() -> ModelConfig:

gpt_server/utils.py

@@ -211,6 +211,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")
+                    hf_overrides = engine_config.get("hf_overrides", "")
 
                 else:
                     logger.error(
@@ -315,6 +316,8 @@ def start_model_worker(config: dict):
                         cmd += f" --vad_model '{vad_model}'"
                     if punc_model:
                         cmd += f" --vad_model '{punc_model}'"
+                    if hf_overrides:
+                        cmd += f" --hf_overrides '{json.dumps(hf_overrides)}'"
                     p = Process(target=run_cmd, args=(cmd,))
                     # p.start()
                     process.append(p)