From 15f52bfc8ffa91ce223e397ea10c1b4a9542c29a Mon Sep 17 00:00:00 2001 From: Wenjie Zhang Date: Fri, 7 Mar 2025 01:05:50 +0800 Subject: [PATCH] =?UTF-8?q?=E4=BF=AE=E5=A4=8D=E3=80=80embedding=E3=80=80?= =?UTF-8?q?=E5=92=8C=20reranker=20=E6=97=A0=E6=B3=95=E6=9C=AC=E5=9C=B0?= =?UTF-8?q?=E8=B0=83=E7=94=A8=20gpu=20=E7=9A=84=E9=97=AE=E9=A2=98=EF=BC=8C?= =?UTF-8?q?=E9=9C=80=E8=A6=81=E5=AE=89=E8=A3=85=20nvidia-container-toolkit?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/config/__init__.py | 1 + src/core/graphbase.py | 10 +++------- src/models/embedding.py | 8 +++++--- src/models/rerank_model.py | 2 +- 4 files changed, 10 insertions(+), 11 deletions(-) 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")