From 5000afc3cacc635da9b0c0d62d142f89bc5b05dd Mon Sep 17 00:00:00 2001
From: Wenjie Zhang <xerrors@163.com>
Date: Wed, 25 Jun 2025 02:12:19 +0800
Subject: [PATCH] =?UTF-8?q?feat:=20=E6=9B=B4=E6=96=B0=E6=A8=A1=E5=9E=8B?=
 =?UTF-8?q?=E6=9C=8D=E5=8A=A1=E8=84=9A=E6=9C=AC=E5=92=8C=E9=85=8D=E7=BD=AE?=
 =?UTF-8?q?=EF=BC=8C=E6=94=AF=E6=8C=81=20Qwen3=20=E7=B3=BB=E5=88=97?=
 =?UTF-8?q?=E6=A8=A1=E5=9E=8B?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

- 修改 run.sh 脚本，新增对 Qwen3-32B、Qwen3-Embedding-0.6B 和 Qwen3-Reranker-0.6B 模型的支持。
- 更新 rerank_model.py，优化 API URL 获取逻辑，使用 get_docker_safe_url 函数。
- 在 models.yaml 中添加 Qwen3 系列模型的配置信息，包括名称、维度和 API 密钥。
---
 scripts/vllm/run.sh        | 45 +++++++++++++++++++++++++++++---------
 src/models/rerank_model.py | 11 +++++-----
 src/static/models.yaml     | 13 +++++++++++
 3 files changed, 54 insertions(+), 15 deletions(-)

diff --git a/scripts/vllm/run.sh b/scripts/vllm/run.sh
index 91fe6ee4..5bc4fdb8 100644
--- a/scripts/vllm/run.sh
+++ b/scripts/vllm/run.sh
@@ -1,23 +1,48 @@
-MODEL=Meta-Llama-3-8B-Instruct
+MODEL_DIR=/data/public/models
+PORT=8081
+
+TENSOR_PARALLEL_SIZE=1
+export CUDA_VISIBLE_DEVICES="0"
+
+source .venv/bin/activate
 
 if [ -z "$1" ]; then
     echo "Error: No argument provided. Please specify a model name."
     exit 1
 fi
 
-if [ "$1" = "llama" ]; then
-    python -m vllm.entrypoints.openai.api_server \
-        --model="/hdd/zwj/models/meta-llama/Meta-Llama-3-8B-Instruct" \
-        --tensor-parallel-size 2 \
+if [ "$1" = "qwen3:32b" ]; then
+    vllm serve "$MODEL_DIR/Qwen/Qwen3-32B" \
         --trust-remote-code \
-        --device auto \
-        --gpu-memory-utilization 0.8 \
-        --dtype half \
+        --device cuda --dtype auto --tensor-parallel-size $TENSOR_PARALLEL_SIZE \
+        --max_model_len 16384 \
         --served-model-name "$1" \
-        --host 0.0.0.0 \
-        --port 8080
+        --enable-auto-tool-choice \
+        --tool-call-parser hermes \
+        --host 0.0.0.0 --port $PORT
 fi
 
+# Qwen/Qwen3-Embedding-0.6B
+if [ "$1" = "Qwen3-Embedding-0.6B" ]; then
+    vllm serve "$MODEL_DIR/Qwen/Qwen3-Embedding-0.6B"  --task embed \
+        --trust-remote-code --max_model_len 4096 \
+        --device cuda --dtype auto --tensor-parallel-size $TENSOR_PARALLEL_SIZE \
+        --served-model-name "$1" --host 0.0.0.0 --port $PORT
+fi
+
+if [ "$1" = "Qwen3-Reranker-0.6B" ]; then
+    vllm serve "$MODEL_DIR/Qwen/Qwen3-Reranker-0.6B" --task rerank \
+        --trust-remote-code \
+        --device cuda --dtype auto --tensor-parallel-size $TENSOR_PARALLEL_SIZE \
+        --max_model_len 4096 \
+        --served-model-name "$1" --host 0.0.0.0 --port $PORT
+fi
+
+
+
+
+
+
 # https://docs.vllm.ai/en/latest/serving/openai_compatible_server.html#named-arguments
 # model	模型路径，以文件夹结尾
 # tensor-parallel-size	张量并行副本数，即GPU的数量，咱这儿只有2张卡
diff --git a/src/models/rerank_model.py b/src/models/rerank_model.py
index 60f6948b..f8d00968 100644
--- a/src/models/rerank_model.py
+++ b/src/models/rerank_model.py
@@ -5,7 +5,7 @@ import numpy as np
 from FlagEmbedding import FlagReranker
 
 from src import config
-from src.utils.logging_config import logger
+from src.utils import logger, get_docker_safe_url
 
 
 class LocalReranker(FlagReranker):
@@ -24,11 +24,12 @@ def sigmoid(x):
 
 class SiliconFlowReranker:
     def __init__(self, **kwargs):
-        self.url = "https://api.siliconflow.cn/v1/rerank"
-        self.model = config.reranker_names[config.reranker]["name"]
+        model_info = config.reranker_names[config.reranker]
+        self.url = get_docker_safe_url(model_info["url"])
+        self.model = model_info["name"]
 
-        api_key = os.getenv("SILICONFLOW_API_KEY")
-        assert api_key, "SILICONFLOW_API_KEY is required"
+        api_key = os.getenv(model_info["api_key"], model_info["api_key"])
+        assert api_key, f"{model_info['name']} api_key is required"
         self.headers = {
             "Authorization": f"Bearer {api_key}",
             "Content-Type": "application/json"
diff --git a/src/static/models.yaml b/src/static/models.yaml
index b0b0d960..61ff1074 100644
--- a/src/static/models.yaml
+++ b/src/static/models.yaml
@@ -130,6 +130,12 @@ EMBED_MODEL_INFO:
     url: https://api.siliconflow.cn/v1/embeddings
     api_key: SILICONFLOW_API_KEY
 
+  vllm/Qwen/Qwen3-Embedding-0.6B:
+    name: qwen3-embedding-0.6b
+    dimension: 1024
+    url: http://172.19.13.6:8081/v1/embeddings
+    api_key: no_api_key
+
   ollama/nomic-embed-text:
     name: nomic-embed-text
     dimension: 768
@@ -146,3 +152,10 @@ RERANKER_LIST:
 
   siliconflow/BAAI/bge-reranker-v2-m3:
     name: BAAI/bge-reranker-v2-m3
+    url: https://api.siliconflow.cn/v1/rerank
+    api_key: SILICONFLOW_API_KEY
+
+  vllm/Qwen/Qwen3-Reranker-0.6B:
+    name: Qwen/Qwen3-Reranker-0.6B
+    url: http://172.19.13.6:8081/v1/rerank
+    api_key: no_api_key