修复 embedding 和 reranker 无法本地调用 gpu 的问题，需要安装 nvidia-container-toolkit

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()

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)

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

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