diff --git a/src/core/embedder.py b/src/core/embedder.py index a11fada..ff35556 100644 --- a/src/core/embedder.py +++ b/src/core/embedder.py @@ -18,6 +18,7 @@ vectors = batch_embed(embedder, long_text_list) """ import os +import threading import logging from typing import Protocol @@ -29,6 +30,8 @@ logger = logging.getLogger("md-vector-db") _HF_MIRROR = os.environ.get("HF_MIRROR", "https://hf-mirror.com") +_HF_ENV_LOCK = threading.Lock() + # -- Provider 默认配置 -- _PROVIDER_DEFAULTS: dict[str, dict[str, str | int]] = { "openai": { @@ -78,16 +81,18 @@ class LocalEmbedder: # 通过 HF_ENDPOINT 环境变量设置镜像(sentence-transformers 依赖 huggingface_hub) # 临时设置仅用于模型下载,下载完成后还原 old_endpoint = os.environ.get("HF_ENDPOINT") - os.environ["HF_ENDPOINT"] = _HF_MIRROR + with _HF_ENV_LOCK: + os.environ["HF_ENDPOINT"] = _HF_MIRROR try: self._model = SentenceTransformer( config.local_model, device=device ) finally: - if old_endpoint is not None: - os.environ["HF_ENDPOINT"] = old_endpoint - else: - os.environ.pop("HF_ENDPOINT", None) + with _HF_ENV_LOCK: + if old_endpoint is not None: + os.environ["HF_ENDPOINT"] = old_endpoint + else: + os.environ.pop("HF_ENDPOINT", None) @property def dimension(self) -> int: