From 6ab4ec77b9e609e7c7ba9d815ae8dfa9a9a37bd8 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E6=9D=8E=E5=A6=82=E5=A8=81?= Date: Sat, 12 Jul 2025 00:13:51 +0800 Subject: [PATCH] =?UTF-8?q?feat:=20=E5=90=91=E9=87=8F=E6=A8=A1=E5=9E=8B?= =?UTF-8?q?=E6=94=B9=E7=94=A8=E6=9C=AC=E5=9C=B0?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- config.py | 1 + services/rag_service.py | 4 ++-- services/vector_store.py | 2 +- 3 files changed, 4 insertions(+), 3 deletions(-) diff --git a/config.py b/config.py index 0b2172f..1d41d36 100644 --- a/config.py +++ b/config.py @@ -25,6 +25,7 @@ class Config: 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", "") diff --git a/services/rag_service.py b/services/rag_service.py index 2251540..2f1b3e0 100644 --- a/services/rag_service.py +++ b/services/rag_service.py @@ -5,7 +5,7 @@ from langchain.prompts import PromptTemplate from langchain.callbacks import AsyncIteratorCallbackHandler from services.vector_store import AsyncVectorStore from utils.logger import get_logger -from transformers import AutoTokenizer, AutoModelForSequenceClassification +from transformers import AutoTokenizer, AutoModelForCausalLM import os import time import torch @@ -52,7 +52,7 @@ class AsyncRAGService: 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 = AutoModelForCausalLM.from_pretrained( self.rerank_model_path ).eval() # 确保模型配置与 tokenizer 一致 diff --git a/services/vector_store.py b/services/vector_store.py index a645467..0cb90a3 100644 --- a/services/vector_store.py +++ b/services/vector_store.py @@ -31,7 +31,7 @@ class AsyncVectorStore: # 尝试初始化向量编码器,如果网络失败则使用本地方案 try: self.logger.info("正在加载向量编码模型...") - self.encoder = SentenceTransformer("all-MiniLM-L6-v2") + self.encoder = SentenceTransformer(os.getenv("EMBEDDING_MODEL_PATH")) self.logger.info("✓ 向量编码模型加载成功") except Exception as e: self.logger.error(f"⚠️ 向量编码模型加载失败: {e}")