feat: 向量模型改用本地

This commit is contained in:
李如威 2025-07-12 00:13:51 +08:00
parent 7fb2b9b876
commit 6ab4ec77b9
3 changed files with 4 additions and 3 deletions

View File

@ -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", "")

View File

@ -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 一致

View File

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