feat: 引入本地模型
This commit is contained in:
parent
7807866a37
commit
e98c2feb36
|
@ -5,8 +5,10 @@ from langchain.prompts import PromptTemplate
|
||||||
from langchain.callbacks import AsyncIteratorCallbackHandler
|
from langchain.callbacks import AsyncIteratorCallbackHandler
|
||||||
from services.vector_store import AsyncVectorStore
|
from services.vector_store import AsyncVectorStore
|
||||||
from utils.logger import get_logger
|
from utils.logger import get_logger
|
||||||
|
from transformers import AutoTokenizer, AutoModelForSequenceClassification
|
||||||
import os
|
import os
|
||||||
import time
|
import time
|
||||||
|
import torch
|
||||||
|
|
||||||
|
|
||||||
class AsyncRAGService:
|
class AsyncRAGService:
|
||||||
|
@ -25,25 +27,6 @@ class AsyncRAGService:
|
||||||
openai_api_base=self.openai_api_base,
|
openai_api_base=self.openai_api_base,
|
||||||
)
|
)
|
||||||
|
|
||||||
self.rerank_llm = ChatOpenAI(
|
|
||||||
model="dengcao/Qwen3-Reranker-8B:Q3_K_M",
|
|
||||||
temperature=0.7,
|
|
||||||
openai_api_key=self.openai_api_key,
|
|
||||||
openai_api_base=self.openai_api_base,
|
|
||||||
)
|
|
||||||
|
|
||||||
self.rerank_prompt_template = PromptTemplate(
|
|
||||||
input_variables=["question", "content"],
|
|
||||||
template=(
|
|
||||||
"你是一个智能评分助手,请判断以下“文档片段”与“用户问题”的相关程度。\n"
|
|
||||||
"请只输出一个介于 0 到 1 之间的分数,数值越高表示相关性越强。\n\n"
|
|
||||||
"用户问题:\n"
|
|
||||||
"{question}\n\n"
|
|
||||||
"文档片段:\n"
|
|
||||||
"{content}"
|
|
||||||
),
|
|
||||||
)
|
|
||||||
|
|
||||||
self.prompt_template = PromptTemplate(
|
self.prompt_template = PromptTemplate(
|
||||||
input_variables=["context", "question"],
|
input_variables=["context", "question"],
|
||||||
template=(
|
template=(
|
||||||
|
@ -55,6 +38,15 @@ class AsyncRAGService:
|
||||||
),
|
),
|
||||||
)
|
)
|
||||||
|
|
||||||
|
self.tokenizer = AutoTokenizer.from_pretrained(
|
||||||
|
"/Volumes/LRW/Model/Qwen3-Embedding-0.6B", trust_remote_code=True
|
||||||
|
)
|
||||||
|
self.rerank_model = AutoModelForSequenceClassification.from_pretrained(
|
||||||
|
"/Volumes/LRW/Model/Qwen3-Embedding-0.6B",
|
||||||
|
trust_remote_code=True,
|
||||||
|
device_map="auto", # 或 "cuda"
|
||||||
|
)
|
||||||
|
|
||||||
self.logger.info("RAG服务初始化完成")
|
self.logger.info("RAG服务初始化完成")
|
||||||
|
|
||||||
async def add_document_async(self, content: str, filename: str) -> str:
|
async def add_document_async(self, content: str, filename: str) -> str:
|
||||||
|
@ -224,27 +216,36 @@ class AsyncRAGService:
|
||||||
async def _rerank_results(
|
async def _rerank_results(
|
||||||
self, question: str, search_results: List[Dict[str, Any]]
|
self, question: str, search_results: List[Dict[str, Any]]
|
||||||
) -> List[Dict[str, Any]]:
|
) -> List[Dict[str, Any]]:
|
||||||
"""使用 rerank LLM 对搜索结果重新排序"""
|
"""使用 Qwen3-Reranker 对搜索结果批量重排序"""
|
||||||
|
|
||||||
async def score_result(result: Dict[str, Any]) -> float:
|
# 准备 batch 输入:格式必须是 Query: xxx\nDocument: yyy
|
||||||
prompt = self.rerank_prompt_template.format(
|
batch_texts = [
|
||||||
content=result["content"][:1000],
|
f"Query: {question}\nDocument: {r['content'][:1000]}" # 可以根据显存调整截断长度
|
||||||
question=question
|
for r in search_results
|
||||||
)
|
]
|
||||||
try:
|
|
||||||
response = await asyncio.to_thread(self.rerank_llm.invoke, prompt)
|
|
||||||
self.logger.info(f"rerank 评分: {response.content.strip()} for {result['metadata']['filename']}")
|
|
||||||
score = float(response.content.strip())
|
|
||||||
return max(0.0, min(score, 1.0))
|
|
||||||
except Exception as e:
|
|
||||||
self.logger.warning(f"rerank 评分失败,fallback 使用向量相似度: {e}")
|
|
||||||
return 1 - result["distance"]
|
|
||||||
|
|
||||||
scores = await asyncio.gather(*[score_result(r) for r in search_results])
|
# 使用 tokenizer 构建 batch 输入
|
||||||
|
inputs = self.tokenizer(
|
||||||
|
batch_texts,
|
||||||
|
return_tensors="pt",
|
||||||
|
padding=True,
|
||||||
|
truncation=True,
|
||||||
|
max_length=1024, # Qwen3 的最大上下文长度,建议限制
|
||||||
|
).to(self.rerank_model.device)
|
||||||
|
|
||||||
|
# 推理打分(关闭梯度计算)
|
||||||
|
with torch.no_grad():
|
||||||
|
outputs = self.rerank_model(**inputs)
|
||||||
|
logits = outputs.logits.squeeze(-1)
|
||||||
|
|
||||||
|
# 如果是二分类模型,通常需要做 sigmoid 激活
|
||||||
|
scores = torch.sigmoid(logits).tolist()
|
||||||
|
|
||||||
|
# 写入到每个 search_result 中
|
||||||
for r, score in zip(search_results, scores):
|
for r, score in zip(search_results, scores):
|
||||||
r["rerank_score"] = score
|
r["rerank_score"] = max(0.0, min(score, 1.0)) # 保证分数在 0-1 范围
|
||||||
|
|
||||||
return sorted(search_results, key=lambda r: r["rerank_score"], reverse=True)
|
return search_results
|
||||||
|
|
||||||
async def _build_context_async(self, search_results: List[Dict[str, Any]]) -> str:
|
async def _build_context_async(self, search_results: List[Dict[str, Any]]) -> str:
|
||||||
"""异步构建上下文"""
|
"""异步构建上下文"""
|
||||||
|
|
Loading…
Reference in New Issue