225 lines
9.4 KiB
Python
225 lines
9.4 KiB
Python
import os
|
||
import argparse
|
||
import torch
|
||
import torch.nn.functional as F
|
||
from fastapi import FastAPI, Request, HTTPException
|
||
from fastapi.responses import JSONResponse
|
||
from pydantic import BaseModel
|
||
from typing import List
|
||
from modelscope import AutoTokenizer, AutoModel, AutoModelForCausalLM
|
||
from transformers import BitsAndBytesConfig
|
||
|
||
# 自动获取脚本所在目录作为基础路径
|
||
BASE_DIR = os.path.dirname(os.path.abspath(__file__))
|
||
|
||
# 在 load_models() 里添加量化配置
|
||
quantization_config = BitsAndBytesConfig(load_in_8bit=True)
|
||
# ===================== 模型路径 =====================
|
||
EMBED_MODEL_PATH = os.path.join(BASE_DIR, "models", "Qwen", "Qwen3-Embedding-0___6B")
|
||
RERANK_MODEL_PATH = os.path.join(BASE_DIR, "models", "Qwen", "Qwen3-Reranker-4B")
|
||
|
||
# ===================== 服务参数（--api-key / --model-name）=====================
|
||
# api_key 为空：不校验（任意值放行）；model_name 为空：不校验请求 model 字段
|
||
_model_args: dict = {"api_key": "", "model_name": ""}
|
||
|
||
def check_api_key(request: Request):
|
||
"""校验请求 API Key：服务未配置 key（空）时放行任意值；否则校验 Authorization / X-API-Key。
|
||
|
||
校验失败抛出 HTTP 401。
|
||
"""
|
||
api_key = _model_args.get("api_key", "")
|
||
if not api_key:
|
||
return
|
||
header = request.headers.get("Authorization", "")
|
||
if header.startswith("Bearer "):
|
||
provided = header[len("Bearer "):].strip()
|
||
else:
|
||
provided = request.headers.get("X-API-Key", "").strip()
|
||
if provided != api_key:
|
||
raise HTTPException(status_code=401, detail="API Key 无效或缺失")
|
||
|
||
def check_model_name(model: str | None):
|
||
"""校验请求 modelname：服务未配置模型名（空）时放行任意值；否则要求请求 model 匹配。
|
||
|
||
校验失败抛出 HTTP 400。
|
||
"""
|
||
model_name = _model_args.get("model_name", "")
|
||
if not model_name:
|
||
return
|
||
if str(model or "").strip() != model_name:
|
||
raise HTTPException(status_code=400, detail=f"模型名不匹配: 期望 '{model_name}'，收到 '{model or ''}'")
|
||
|
||
# ===================== 初始化 FastAPI =====================
|
||
app = FastAPI(title="RAG Backend with 8-bit Qwen3 Models")
|
||
|
||
# ===================== 全局变量（模型、分词器） =====================
|
||
embed_tokenizer = None
|
||
embed_model = None
|
||
rerank_tokenizer = None
|
||
rerank_model = None
|
||
|
||
# ===================== 工具函数 =====================
|
||
def last_token_pool(last_hidden_states, attention_mask):
|
||
left_padding = (attention_mask[:, -1].sum() == attention_mask.shape[0])
|
||
if left_padding:
|
||
return last_hidden_states[:, -1]
|
||
else:
|
||
sequence_lengths = attention_mask.sum(dim=1) - 1
|
||
batch_size = last_hidden_states.shape[0]
|
||
return last_hidden_states[torch.arange(batch_size, device=last_hidden_states.device), sequence_lengths]
|
||
|
||
def get_detailed_instruct(task_description: str, query: str) -> str:
|
||
return f'Instruct: {task_description}\nQuery:{query}'
|
||
|
||
def format_rerank_instruction(instruction, query, doc):
|
||
if instruction is None:
|
||
instruction = 'Given a web search query, retrieve relevant passages that answer the query'
|
||
return f"<Instruct>: {instruction}\n<Query>: {query}\n<Document>: {doc}"
|
||
|
||
# ===================== 启动时加载模型（8‑bit） =====================
|
||
@app.on_event("startup")
|
||
def load_models():
|
||
global embed_tokenizer, embed_model, rerank_tokenizer, rerank_model
|
||
|
||
# ---------- 加载 Embedding 模型 ----------
|
||
print("Loading Embedding model with 16-bit quantization...")
|
||
embed_tokenizer = AutoTokenizer.from_pretrained(EMBED_MODEL_PATH, padding_side='left')
|
||
embed_model = AutoModel.from_pretrained(
|
||
EMBED_MODEL_PATH,
|
||
#load_in_8bit=True, # 启用 8-bit 量化
|
||
device_map="auto", # 自动分配到 GPU
|
||
dtype=torch.float16 # 计算精度保持 fp16
|
||
)
|
||
embed_model.eval()
|
||
print("Embedding model loaded.")
|
||
|
||
# ---------- 加载 Reranker 模型 ----------
|
||
print("Loading Reranker model with 8-bit quantization...")
|
||
rerank_tokenizer = AutoTokenizer.from_pretrained(RERANK_MODEL_PATH, padding_side='left')
|
||
rerank_model = AutoModelForCausalLM.from_pretrained(
|
||
RERANK_MODEL_PATH,
|
||
quantization_config=quantization_config,
|
||
device_map="auto",
|
||
dtype=torch.float16
|
||
)
|
||
rerank_model.eval()
|
||
|
||
# 预计算 Reranker 所需的固定 tokens
|
||
global token_true_id, token_false_id, prefix_tokens, suffix_tokens, max_rerank_length
|
||
token_false_id = rerank_tokenizer.convert_tokens_to_ids("no")
|
||
token_true_id = rerank_tokenizer.convert_tokens_to_ids("yes")
|
||
max_rerank_length = 32768
|
||
|
||
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"
|
||
prefix_tokens = rerank_tokenizer.encode(prefix, add_special_tokens=False)
|
||
suffix_tokens = rerank_tokenizer.encode(suffix, add_special_tokens=False)
|
||
|
||
print("Reranker model loaded.")
|
||
|
||
# ===================== API 请求体定义 =====================
|
||
class EmbedRequest(BaseModel):
|
||
task_description: str = "Given a web search query, retrieve relevant passages that answer the query"
|
||
model: str | None = None # 可选，服务配置了 --model-name 时须匹配
|
||
queries: List[str]
|
||
documents: List[str] # 可选，若只查询文本可只传 queries
|
||
|
||
class EmbedResponse(BaseModel):
|
||
query_embeddings: List[List[float]]
|
||
doc_embeddings: List[List[float]] # 当传入 documents 时返回
|
||
|
||
class RerankRequest(BaseModel):
|
||
task_description: str = "Given a web search query, retrieve relevant passages that answer the query"
|
||
model: str | None = None # 可选，服务配置了 --model-name 时须匹配
|
||
query: str
|
||
documents: List[str]
|
||
|
||
class RerankResponse(BaseModel):
|
||
scores: List[float]
|
||
|
||
# ===================== Embedding 端点 =====================
|
||
@app.post("/embed", response_model=EmbedResponse)
|
||
async def get_embeddings(req: EmbedRequest, request: Request):
|
||
check_api_key(request)
|
||
check_model_name(req.model)
|
||
# 为 query 加上指令，document 不加指令
|
||
instructed_queries = [get_detailed_instruct(req.task_description, q) for q in req.queries]
|
||
# 所有文本拼成一个 batch：先是 queries，后是 documents（如果有）
|
||
input_texts = instructed_queries + req.documents
|
||
|
||
# 分词
|
||
batch_dict = embed_tokenizer(
|
||
input_texts,
|
||
padding=True,
|
||
truncation=True,
|
||
return_tensors="pt",
|
||
).to(embed_model.device)
|
||
|
||
with torch.no_grad():
|
||
outputs = embed_model(**batch_dict)
|
||
embeddings = last_token_pool(outputs.last_hidden_state, batch_dict['attention_mask'])
|
||
embeddings = F.normalize(embeddings, p=2, dim=1)
|
||
|
||
# 拆分 query 和 document 的嵌入
|
||
nq = len(instructed_queries)
|
||
query_embs = embeddings[:nq].cpu().tolist()
|
||
doc_embs = embeddings[nq:].cpu().tolist() if req.documents else []
|
||
|
||
return EmbedResponse(query_embeddings=query_embs, doc_embeddings=doc_embs)
|
||
|
||
# ===================== Reranker 端点 =====================
|
||
@app.post("/rerank", response_model=RerankResponse)
|
||
async def rerank_documents(req: RerankRequest, request: Request):
|
||
check_api_key(request)
|
||
check_model_name(req.model)
|
||
# 构造带格式的输入对
|
||
pairs = [
|
||
format_rerank_instruction(req.task_description, req.query, doc)
|
||
for doc in req.documents
|
||
]
|
||
|
||
# 处理输入，添加 prefix/suffix tokens 并 pad
|
||
inputs = rerank_tokenizer(
|
||
pairs,
|
||
padding=False,
|
||
truncation='longest_first',
|
||
return_attention_mask=False,
|
||
max_length=max_rerank_length - len(prefix_tokens) - len(suffix_tokens)
|
||
)
|
||
for i, ele in enumerate(inputs['input_ids']):
|
||
inputs['input_ids'][i] = prefix_tokens + ele + suffix_tokens
|
||
inputs = rerank_tokenizer.pad(inputs, padding=True, return_tensors="pt", max_length=max_rerank_length)
|
||
inputs = {k: v.to(rerank_model.device) for k, v in inputs.items()}
|
||
|
||
with torch.no_grad():
|
||
batch_scores = rerank_model(**inputs).logits[:, -1, :]
|
||
true_vector = batch_scores[:, token_true_id]
|
||
false_vector = batch_scores[:, token_false_id]
|
||
scores = torch.stack([false_vector, true_vector], dim=1)
|
||
scores = torch.nn.functional.log_softmax(scores, dim=1)
|
||
scores = scores[:, 1].exp().cpu().tolist()
|
||
|
||
return RerankResponse(scores=scores)
|
||
|
||
# ===================== 健康检查 =====================
|
||
@app.get("/health")
|
||
def health():
|
||
return {"status": "ok"}
|
||
|
||
|
||
# ======================== 启动入口 ========================
|
||
if __name__ == "__main__":
|
||
parser = argparse.ArgumentParser(description="Qwen3 RAG Server (Embedding + Reranker)")
|
||
parser.add_argument("--port", type=int, default=8903, help="服务端口（默认 8903）")
|
||
parser.add_argument("--model-name", type=str, default="Qwen3-Embedding-0.6B",
|
||
help="服务模型名（校验请求 body.model；默认 Qwen3-Embedding-0.6B，留空则不校验）")
|
||
parser.add_argument("--api-key", type=str, default="",
|
||
help="API Key（校验 Authorization: Bearer <key> 或 X-API-Key；默认空=任意值放行）")
|
||
args = parser.parse_args()
|
||
|
||
_model_args["api_key"] = args.api_key
|
||
_model_args["model_name"] = args.model_name
|
||
|
||
import uvicorn
|
||
uvicorn.run(app, host="0.0.0.0", port=args.port)
|