diff --git a/src/config/__init__.py b/src/config/__init__.py index 9ebdb38d..ab785c52 100644 --- a/src/config/__init__.py +++ b/src/config/__init__.py @@ -64,6 +64,7 @@ class Config(SimpleConfig): self.add_item("reranker", default="siliconflow/BAAI/bge-reranker-v2-m3", des="Re-Ranker 模型", choices=list(RERANKER_LIST.keys())) self.add_item("model_local_paths", default={}, des="本地模型路径") self.add_item("use_rewrite_query", default="off", des="重写查询", choices=["off", "on", "hyde"]) + self.add_item("device", default="cuda", des="运行本地模型的设备", choices=["cpu", "cuda"]) ### <<< 默认配置结束 self.load() diff --git a/src/core/graphbase.py b/src/core/graphbase.py index 6294c683..23de3ee3 100644 --- a/src/core/graphbase.py +++ b/src/core/graphbase.py @@ -1,14 +1,10 @@ import os import json -import torch -from neo4j import GraphDatabase as GD -# from src.plugins import pdf2txt, OneKE -from transformers import AutoTokenizer, AutoModel -from FlagEmbedding import FlagModel, FlagReranker import warnings -from src.plugins import pdf2txt -from src.plugins.oneke import OneKE +import torch +from neo4j import GraphDatabase as GD + from src.utils import logger warnings.filterwarnings("ignore", category=UserWarning) diff --git a/src/models/embedding.py b/src/models/embedding.py index d38129c1..58fc57f5 100644 --- a/src/models/embedding.py +++ b/src/models/embedding.py @@ -19,7 +19,7 @@ class BaseEmbeddingModel: return self.predict(queries) def batch_encode(self, messages, batch_size=20): - logger.info(f"[Embedding: {self.model}] Batch encoding {len(messages)} messages") + logger.info(f"Batch encoding {len(messages)} messages") data = [] if len(messages) > batch_size: @@ -49,11 +49,13 @@ class LocalEmbeddingModel(FlagModel, BaseEmbeddingModel): self.model = config.model_local_paths.get(info["name"], info.get("local_path")) self.model = self.model or info["name"] - logger.info(f"Loading embedding model {info['name']} from {self.model}") + logger.info(f"Loading local model `{info['name']}` from `{self.model}` with device `{config.device}`") super().__init__(self.model, query_instruction_for_retrieval=info.get("query_instruction", None), - use_fp16=False, **kwargs) + use_fp16=False, + device=config.device, + **kwargs) logger.info(f"Embedding model {info['name']} loaded") diff --git a/src/models/rerank_model.py b/src/models/rerank_model.py index d1cf1136..2bd1c2fc 100644 --- a/src/models/rerank_model.py +++ b/src/models/rerank_model.py @@ -15,7 +15,7 @@ class LocalReranker(FlagReranker): model_name_or_path = model_name_or_path or model_info["name"] logger.info(f"Loading Reranker model {config.reranker} from {model_name_or_path}") - super().__init__(model_name_or_path, use_fp16=True, **kwargs) + super().__init__(model_name_or_path, use_fp16=True, device=config.device, **kwargs) logger.info(f"Reranker model {config.reranker} loaded")