feat: 向量模型改用本地
This commit is contained in:
parent
7fb2b9b876
commit
6ab4ec77b9
|
@ -25,6 +25,7 @@ class Config:
|
||||||
EMBEDDING_MODEL_PATH = os.getenv("EMBEDDING_MODEL_PATH", "")
|
EMBEDDING_MODEL_PATH = os.getenv("EMBEDDING_MODEL_PATH", "")
|
||||||
EMBEDDING_MODEL_TYPE = os.getenv("EMBEDDING_MODEL_TYPE", "")
|
EMBEDDING_MODEL_TYPE = os.getenv("EMBEDDING_MODEL_TYPE", "")
|
||||||
EMBEDDING_MODEL_DEVICE = os.getenv("EMBEDDING_MODEL_DEVICE", "")
|
EMBEDDING_MODEL_DEVICE = os.getenv("EMBEDDING_MODEL_DEVICE", "")
|
||||||
|
|
||||||
RERANK_MODEL_PATH = os.getenv("RERANK_MODEL_PATH", "")
|
RERANK_MODEL_PATH = os.getenv("RERANK_MODEL_PATH", "")
|
||||||
RERANK_MODEL_TYPE = os.getenv("RERANK_MODEL_TYPE", "")
|
RERANK_MODEL_TYPE = os.getenv("RERANK_MODEL_TYPE", "")
|
||||||
RERANK_MODEL_DEVICE = os.getenv("RERANK_MODEL_DEVICE", "")
|
RERANK_MODEL_DEVICE = os.getenv("RERANK_MODEL_DEVICE", "")
|
||||||
|
|
|
@ -5,7 +5,7 @@ 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
|
from transformers import AutoTokenizer, AutoModelForCausalLM
|
||||||
import os
|
import os
|
||||||
import time
|
import time
|
||||||
import torch
|
import torch
|
||||||
|
@ -52,7 +52,7 @@ class AsyncRAGService:
|
||||||
if self.tokenizer.pad_token is None:
|
if self.tokenizer.pad_token is None:
|
||||||
self.tokenizer.pad_token = self.tokenizer.eos_token
|
self.tokenizer.pad_token = self.tokenizer.eos_token
|
||||||
self.tokenizer.pad_token_id = self.tokenizer.eos_token_id
|
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
|
self.rerank_model_path
|
||||||
).eval()
|
).eval()
|
||||||
# 确保模型配置与 tokenizer 一致
|
# 确保模型配置与 tokenizer 一致
|
||||||
|
|
|
@ -31,7 +31,7 @@ class AsyncVectorStore:
|
||||||
# 尝试初始化向量编码器,如果网络失败则使用本地方案
|
# 尝试初始化向量编码器,如果网络失败则使用本地方案
|
||||||
try:
|
try:
|
||||||
self.logger.info("正在加载向量编码模型...")
|
self.logger.info("正在加载向量编码模型...")
|
||||||
self.encoder = SentenceTransformer("all-MiniLM-L6-v2")
|
self.encoder = SentenceTransformer(os.getenv("EMBEDDING_MODEL_PATH"))
|
||||||
self.logger.info("✓ 向量编码模型加载成功")
|
self.logger.info("✓ 向量编码模型加载成功")
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
self.logger.error(f"⚠️ 向量编码模型加载失败: {e}")
|
self.logger.error(f"⚠️ 向量编码模型加载失败: {e}")
|
||||||
|
|
Loading…
Reference in New Issue