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