fix: HF_ENDPOINT 环境变量操作添加线程锁防竞态
This commit is contained in:
+10
-5
@@ -18,6 +18,7 @@
|
|||||||
vectors = batch_embed(embedder, long_text_list)
|
vectors = batch_embed(embedder, long_text_list)
|
||||||
"""
|
"""
|
||||||
import os
|
import os
|
||||||
|
import threading
|
||||||
import logging
|
import logging
|
||||||
from typing import Protocol
|
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_MIRROR = os.environ.get("HF_MIRROR", "https://hf-mirror.com")
|
||||||
|
|
||||||
|
_HF_ENV_LOCK = threading.Lock()
|
||||||
|
|
||||||
# -- Provider 默认配置 --
|
# -- Provider 默认配置 --
|
||||||
_PROVIDER_DEFAULTS: dict[str, dict[str, str | int]] = {
|
_PROVIDER_DEFAULTS: dict[str, dict[str, str | int]] = {
|
||||||
"openai": {
|
"openai": {
|
||||||
@@ -78,16 +81,18 @@ class LocalEmbedder:
|
|||||||
# 通过 HF_ENDPOINT 环境变量设置镜像(sentence-transformers 依赖 huggingface_hub)
|
# 通过 HF_ENDPOINT 环境变量设置镜像(sentence-transformers 依赖 huggingface_hub)
|
||||||
# 临时设置仅用于模型下载,下载完成后还原
|
# 临时设置仅用于模型下载,下载完成后还原
|
||||||
old_endpoint = os.environ.get("HF_ENDPOINT")
|
old_endpoint = os.environ.get("HF_ENDPOINT")
|
||||||
os.environ["HF_ENDPOINT"] = _HF_MIRROR
|
with _HF_ENV_LOCK:
|
||||||
|
os.environ["HF_ENDPOINT"] = _HF_MIRROR
|
||||||
try:
|
try:
|
||||||
self._model = SentenceTransformer(
|
self._model = SentenceTransformer(
|
||||||
config.local_model, device=device
|
config.local_model, device=device
|
||||||
)
|
)
|
||||||
finally:
|
finally:
|
||||||
if old_endpoint is not None:
|
with _HF_ENV_LOCK:
|
||||||
os.environ["HF_ENDPOINT"] = old_endpoint
|
if old_endpoint is not None:
|
||||||
else:
|
os.environ["HF_ENDPOINT"] = old_endpoint
|
||||||
os.environ.pop("HF_ENDPOINT", None)
|
else:
|
||||||
|
os.environ.pop("HF_ENDPOINT", None)
|
||||||
|
|
||||||
@property
|
@property
|
||||||
def dimension(self) -> int:
|
def dimension(self) -> int:
|
||||||
|
|||||||
Reference in New Issue
Block a user