47 lines
1.4 KiB
Python
47 lines
1.4 KiB
Python
|
"""集中管理项目本地缓存目录（HuggingFace / tiktoken）。
|
|||
|
|
|
|||
|
|
所有 Python 库通过环境变量引用 ``<repo_root>/.cache/{huggingface,tiktoken}``，
|
|||
|
|
避免写入用户家目录，也避免不同部署路径（本地 / Docker）下缓存位置不一致。
|
|||
|
|
"""
|
|||
|
|
|
|||
|
|
from __future__ import annotations
|
|||
|
|
|
|||
|
|
import os
|
|||
|
|
from pathlib import Path
|
|||
|
|
|
|||
|
|
|
|||
|
|
def repo_root() -> Path:
|
|||
|
|
# backend/app/core/cache_paths.py -> backend -> 仓库根目录
|
|||
|
|
return Path(__file__).resolve().parents[3]
|
|||
|
|
|
|||
|
|
|
|||
|
|
def cache_root() -> Path:
|
|||
|
|
return repo_root() / ".cache"
|
|||
|
|
|
|||
|
|
|
|||
|
|
def huggingface_cache_dir() -> Path:
|
|||
|
|
return cache_root() / "huggingface"
|
|||
|
|
|
|||
|
|
|
|||
|
|
def tiktoken_cache_dir() -> Path:
|
|||
|
|
return cache_root() / "tiktoken"
|
|||
|
|
|
|||
|
|
|
|||
|
|
def setup_local_caches() -> None:
|
|||
|
|
"""进程启动时统一设置 HF / tiktoken 缓存环境变量，并确保目录存在。
|
|||
|
|
|
|||
|
|
必须在 import ``docling`` / ``tiktoken`` 等依赖之前调用，否则首次使用会
|
|||
|
|
仍然走到默认 ``~/.cache`` 路径。
|
|||
|
|
"""
|
|||
|
|
|
|||
|
|
hf_dir = huggingface_cache_dir()
|
|||
|
|
tiktoken_dir = tiktoken_cache_dir()
|
|||
|
|
hf_dir.mkdir(parents=True, exist_ok=True)
|
|||
|
|
(hf_dir / "hub").mkdir(parents=True, exist_ok=True)
|
|||
|
|
tiktoken_dir.mkdir(parents=True, exist_ok=True)
|
|||
|
|
|
|||
|
|
os.environ["HF_HOME"] = str(hf_dir)
|
|||
|
|
os.environ["HUGGINGFACE_HUB_CACHE"] = str(hf_dir / "hub")
|
|||
|
|
os.environ["HF_HUB_CACHE"] = str(hf_dir / "hub")
|
|||
|
|
os.environ["TIKTOKEN_CACHE_DIR"] = str(tiktoken_dir)
|