fix: 修复 MarkdownSplitter 死循环 + 添加 GPU 自动检测
- ingest.py: _split_single_paragraph 中分隔符距 start 小于 overlap 时 start 会回退为负数,Python rfind 负索引导致无限循环 修复: start = max(start + 1, next_start) 确保始终前进 - embedder.py: LocalEmbedder 自动检测 CUDA,优先使用 GPU - ingest_obsidian.py: MAX_SIZE 调整为 200KB 适配 GPU 嵌入 - 清理临时日志文件
This commit is contained in:
+15
-3
@@ -52,21 +52,33 @@ class Embedder(Protocol):
|
||||
|
||||
# -- 本地模型 --
|
||||
class LocalEmbedder:
|
||||
"""sentence-transformers 本地模型嵌入器."""
|
||||
"""sentence-transformers 本地模型嵌入器. 自动检测 GPU."""
|
||||
|
||||
def __init__(self, config: EmbedConfig):
|
||||
from sentence_transformers import SentenceTransformer
|
||||
|
||||
# 自动检测 GPU
|
||||
try:
|
||||
import torch
|
||||
device = "cuda" if torch.cuda.is_available() else "cpu"
|
||||
if device == "cuda":
|
||||
logger.info("检测到 GPU: %s", torch.cuda.get_device_name(0))
|
||||
except ImportError:
|
||||
device = "cpu"
|
||||
|
||||
self._config = config
|
||||
try:
|
||||
self._model = SentenceTransformer(
|
||||
config.local_model, local_files_only=True
|
||||
config.local_model, local_files_only=True, device=device
|
||||
)
|
||||
except Exception:
|
||||
logger.info("模型未缓存, 通过镜像下载 %s", config.local_model)
|
||||
old_endpoint = os.environ.get("HF_ENDPOINT")
|
||||
os.environ["HF_ENDPOINT"] = _HF_MIRROR
|
||||
try:
|
||||
self._model = SentenceTransformer(config.local_model)
|
||||
self._model = SentenceTransformer(
|
||||
config.local_model, device=device
|
||||
)
|
||||
finally:
|
||||
if old_endpoint is not None:
|
||||
os.environ["HF_ENDPOINT"] = old_endpoint
|
||||
|
||||
Reference in New Issue
Block a user