fix: HF_ENDPOINT 环境变量操作添加线程锁防竞态
This commit is contained in:
+10
-5
@@ -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:
|
||||
|
||||
Reference in New Issue
Block a user