169 lines
6.9 KiB
Python
169 lines
6.9 KiB
Python
import os
|
||
import torch
|
||
import torch.nn.functional as F
|
||
from fastapi import FastAPI
|
||
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")
|
||
|
||
# ===================== 初始化 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"
|
||
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"
|
||
query: str
|
||
documents: List[str]
|
||
|
||
class RerankResponse(BaseModel):
|
||
scores: List[float]
|
||
|
||
# ===================== Embedding 端点 =====================
|
||
@app.post("/embed", response_model=EmbedResponse)
|
||
async def get_embeddings(req: EmbedRequest):
|
||
# 为 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):
|
||
# 构造带格式的输入对
|
||
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"}
|