diff --git a/.env.example b/.env.example index 768382a..a01fc65 100644 --- a/.env.example +++ b/.env.example @@ -5,6 +5,16 @@ OPENAI_BASE_URL=https://api.openai.com/v1 # 向量数据库配置 CHROMA_PERSIST_DIRECTORY=./chroma_db +# 模型配置 + +# RERANK_MODEL_PATH=/Volumes/LRW/Model/Qwen3-Reranker-0.6B +# RERANK_MODEL_TYPE=Qwen3-Reranker-0.6B +# RERANK_MODEL_DEVICE=cpu + +# EMBEDDING_MODEL_PATH=your_embedding_model_path_here +# EMBEDDING_MODEL_TYPE=your_embedding_model_type_here +# EMBEDDING_MODEL_DEVICE=cpu + # 应用配置 APP_NAME=Easy RAG Service APP_VERSION=1.0.0 diff --git a/config.py b/config.py index 257b1d6..0b2172f 100644 --- a/config.py +++ b/config.py @@ -21,6 +21,14 @@ class Config: OPENAI_API_KEY = os.getenv("OPENAI_API_KEY") OPENAI_BASE_URL = os.getenv("OPENAI_BASE_URL", "https://api.openai.com/v1") + # 模型配置 + EMBEDDING_MODEL_PATH = os.getenv("EMBEDDING_MODEL_PATH", "") + EMBEDDING_MODEL_TYPE = os.getenv("EMBEDDING_MODEL_TYPE", "") + EMBEDDING_MODEL_DEVICE = os.getenv("EMBEDDING_MODEL_DEVICE", "") + RERANK_MODEL_PATH = os.getenv("RERANK_MODEL_PATH", "") + RERANK_MODEL_TYPE = os.getenv("RERANK_MODEL_TYPE", "") + RERANK_MODEL_DEVICE = os.getenv("RERANK_MODEL_DEVICE", "") + # 向量数据库配置 CHROMA_PERSIST_DIRECTORY = os.getenv("CHROMA_PERSIST_DIRECTORY", "./chroma_db") diff --git a/services/rag_service.py b/services/rag_service.py index 8c4873e..2251540 100644 --- a/services/rag_service.py +++ b/services/rag_service.py @@ -21,6 +21,9 @@ class AsyncRAGService: self.vector_store = AsyncVectorStore() self.openai_api_base = os.getenv("OPENAI_BASE_URL") self.openai_api_key = os.getenv("OPENAI_API_KEY") + self.rerank_model_path = os.getenv("RERANK_MODEL_PATH", "") + self.open_rerank = bool(self.rerank_model_path) + self.llm = ChatOpenAI( model="deepseek-r1:8b", temperature=0.7, @@ -39,17 +42,23 @@ class AsyncRAGService: ), ) - self.tokenizer = AutoTokenizer.from_pretrained("/Volumes/LRW/Model/Qwen3-Embedding-0.6B", padding_side="left") - # 强制设置 padding token - if self.tokenizer.pad_token is None: - self.tokenizer.pad_token = self.tokenizer.eos_token - self.tokenizer.pad_token_id = self.tokenizer.eos_token_id - - self.rerank_model = AutoModelForSequenceClassification.from_pretrained("/Volumes/LRW/Model/Qwen3-Embedding-0.6B").eval() - - # 确保模型配置与 tokenizer 一致 - if hasattr(self.rerank_model.config, "pad_token_id"): - self.rerank_model.config.pad_token_id = self.tokenizer.pad_token_id + # 如果需要重排 + if self.open_rerank: + self.logger.info("初始化 Ranker 模型...") + self.tokenizer = AutoTokenizer.from_pretrained( + self.rerank_model_path, padding_side="left" + ) + # 强制设置 padding token + if self.tokenizer.pad_token is None: + self.tokenizer.pad_token = self.tokenizer.eos_token + self.tokenizer.pad_token_id = self.tokenizer.eos_token_id + self.rerank_model = AutoModelForSequenceClassification.from_pretrained( + self.rerank_model_path + ).eval() + # 确保模型配置与 tokenizer 一致 + if hasattr(self.rerank_model.config, "pad_token_id"): + self.rerank_model.config.pad_token_id = self.tokenizer.pad_token_id + self.logger.info("✓ 初始化 Ranker 模型成功") self.logger.info("RAG服务初始化完成") @@ -91,7 +100,7 @@ class AsyncRAGService: } # rerank - reranked_results = await self._rerank_results(question, search_results) + reranked_results = await self._rerank_results(question, search_results, skip_rerank=not self.open_rerank) # 并行执行上下文构建和 LLM 调用准备 context_task = asyncio.create_task( @@ -146,7 +155,7 @@ class AsyncRAGService: return # rerank - reranked_results = await self._rerank_results(question, search_results) + reranked_results = await self._rerank_results(question, search_results, skip_rerank=not self.open_rerank) # 构建上下文和源信息 context_task = self._build_context_async(reranked_results) @@ -218,7 +227,7 @@ class AsyncRAGService: return await asyncio.to_thread(_format_sources) async def _rerank_results( - self, question: str, search_results: List[Dict[str, Any]], skip_rerank: bool = True + self, question: str, search_results: List[Dict[str, Any]], skip_rerank: bool = False ) -> List[Dict[str, Any]]: """使用 Qwen3-Reranker 对搜索结果批量重排序""" @@ -229,19 +238,22 @@ class AsyncRAGService: if not search_results: return [] - # 模型相关常量(可初始化时提前保存) - instruction = ( - "Given a web search query, retrieve relevant passages that answer the query" + # ==== Prompt 设置 ==== + prefix = ( + "<|im_start|>system\n" + "You are a helpful assistant that determines whether a document answers a given query. " + 'Respond only with "yes" if the document is helpful, otherwise "no".\n' + "<|im_end|>\n" + "<|im_start|>user\n" ) - 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\n\n\n\n" prefix_tokens = self.tokenizer.encode(prefix, add_special_tokens=False) suffix_tokens = self.tokenizer.encode(suffix, add_special_tokens=False) # 构造符合格式的输入 - def format_pair(query, doc): - return f": {instruction}\n: {query}\n: {doc}" + def format_pair(query: str, doc: str) -> str: + return f": {query}\n: {doc}" pairs = [ format_pair(question, r["content"][:1000]) # 文本截断,避免超长 @@ -292,31 +304,29 @@ class AsyncRAGService: token_false_id = self.tokenizer.convert_tokens_to_ids("no") # 推理评分 + self.logger.info("模型准备输入完毕,开始推理...") with torch.no_grad(): outputs = self.rerank_model(**inputs) logits = outputs.logits # 检查 logits 的维度 - if logits.dim() == 3: - # 如果是3维,取最后一个token的logits + if logits.dim() == 3: # 如果是3维,取最后一个token的logits logits = logits[:, -1, :] - elif logits.dim() == 2: - # 如果是2维,直接使用 - pass - else: - raise ValueError(f"Unexpected logits dimension: {logits.dim()}") + elif logits.dim() != 2: # 如果是2维,直接使用 + raise ValueError(f"Unexpected logits shape: {logits.shape}") # 提取 yes/no token 的 logits true_logits = logits[:, token_true_id] false_logits = logits[:, token_false_id] - stacked = torch.stack([false_logits, true_logits], dim=1) - probs = torch.nn.functional.softmax(stacked, dim=1) - scores = probs[:, 1].tolist() # 取 "yes" 的概率值 + # 推荐用 logits 差值作为分数 + scores = (true_logits - false_logits).tolist() + self.logger.info("模型推理完成") # 写入每条结果 for r, score in zip(search_results, scores): r["rerank_score"] = round(float(score), 4) + self.logger.info(f"重排序完成,得分范围: {min(scores)} - {max(scores)} \n\n {json.dumps(search_results, indent=4)}") return search_results