kris/ForcePilot
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
Wenjie Zhang 2025-03-07 01:05:50 +08:00
parent e50631c196
commit 15f52bfc8f
4 changed files with 10 additions and 11 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
src/config/__init__.py
View File

@ -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("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("model_local_paths", default={}, des="本地模型路径")
self.add_item("use_rewrite_query", default="off", des="重写查询", choices=["off", "on", "hyde"]) 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() self.load()

10
src/core/graphbase.py
View File

@ -1,14 +1,10 @@
import os import os
import json 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 import warnings
from src.plugins import pdf2txt import torch
from src.plugins.oneke import OneKE from neo4j import GraphDatabase as GD
from src.utils import logger from src.utils import logger
warnings.filterwarnings("ignore", category=UserWarning) warnings.filterwarnings("ignore", category=UserWarning)

8
src/models/embedding.py
View File

@ -19,7 +19,7 @@ class BaseEmbeddingModel:
return self.predict(queries) return self.predict(queries)
def batch_encode(self, messages, batch_size=20): 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 = [] data = []
if len(messages) > batch_size: 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 = config.model_local_paths.get(info["name"], info.get("local_path"))
self.model = self.model or info["name"] 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, super().__init__(self.model,
query_instruction_for_retrieval=info.get("query_instruction", None), 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") logger.info(f"Embedding model {info['name']} loaded")

2
src/models/rerank_model.py
View File

@ -15,7 +15,7 @@ class LocalReranker(FlagReranker):
model_name_or_path = model_name_or_path or model_info["name"] 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}") 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") logger.info(f"Reranker model {config.reranker} loaded")