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}")