From 75c5f0aa80b7d6f9bd40ec5155c2b84e306cbd24 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E6=9C=B1=E6=BD=AE?= Date: Wed, 22 Oct 2025 10:44:22 +0800 Subject: [PATCH] modify env --- mcp/semantic_search_server.py | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/mcp/semantic_search_server.py b/mcp/semantic_search_server.py index c9d09dd..c989942 100644 --- a/mcp/semantic_search_server.py +++ b/mcp/semantic_search_server.py @@ -31,13 +31,19 @@ def get_model(model_name_or_path='sentence-transformers/paraphrase-multilingual- # 优先使用本地模型路径 local_model_path = "./models/paraphrase-multilingual-MiniLM-L12-v2" + # 从环境变量获取设备配置，默认为 CPU + device = os.environ.get('SENTENCE_TRANSFORMER_DEVICE', 'cpu') + if device not in ['cpu', 'cuda', 'mps']: + print(f"警告: 不支持的设备类型 '{device}'，使用默认 CPU") + device = 'cpu' + # 检查本地模型是否存在 if os.path.exists(local_model_path): print(f"使用本地模型: {local_model_path}") - embedder = SentenceTransformer(local_model_path, device='cpu') + embedder = SentenceTransformer(local_model_path, device=device) else: print(f"本地模型不存在，使用HuggingFace模型: {model_name_or_path}") - embedder = SentenceTransformer(model_name_or_path, device='cpu') + embedder = SentenceTransformer(model_name_or_path, device=device) return embedder