diff --git a/.gitignore b/.gitignore index a4cf52e..1b07b36 100644 --- a/.gitignore +++ b/.gitignore @@ -1 +1,3 @@ +models/ +hf_cache/ server/server/ diff --git a/docs/models_downloading.py b/docs/models_downloading.py index 8821831..f9c1867 100644 --- a/docs/models_downloading.py +++ b/docs/models_downloading.py @@ -1,69 +1,154 @@ #!/usr/bin/env python3 """ -离线模型下载器(Hugging Face Hub + SOCKS5 自动支持) +models_downloading.py +- 优先级:CN 镜像 (hf-mirror.com) -> 代理(PROXY) -> 官方直连 +- 统一缓存:HF_HOME=./hf_cache(可被环境变量覆盖) +- 进度可见:启用 huggingface_hub 的 tqdm 进度 +- 幂等安全:本地目录已有关键文件则跳过下载 -默认参数: -- MODEL_ID="BAAI/bge-m3" -- MODEL_DIR="./models/bge-m3" -- PROXY="socks5://127.0.0.1:1080" - -可用环境变量覆盖: -- export MODEL_ID="你的模型ID" -- export MODEL_DIR="/保存路径" -- export PROXY="socks5h://ip:port" # 为空表示直连 +可选环境变量: +- MODEL_ID 默认 "BAAI/bge-m3" +- MODEL_DIR 默认 "models/bge-m3" +- HF_HOME 默认 "./hf_cache" +- PROXY 默认 "socks5h://127.0.0.1:1081"(留空表示不走代理) +- HF_ENDPOINT 手动指定镜像时可设置(脚本也会自动探测 cn mirror) """ import os import sys from pathlib import Path -# ==== 自动安装 SOCKS 依赖 ==== -try: - import socks # PySocks -except ImportError: - print("📦 Installing SOCKS proxy support (requests[socks])...") - os.system(f"{sys.executable} -m pip install -U 'requests[socks]'") - import socks +# ---------- 配置 ---------- +MODEL_ID = os.getenv("MODEL_ID", "BAAI/bge-m3") +MODEL_DIR = Path(os.getenv("MODEL_DIR", "models/bge-m3")) +HF_HOME = Path(os.getenv("HF_HOME", Path.cwd() / "hf_cache")) -# ==== 自动安装 huggingface_hub ==== +CN_MIRROR = "https://hf-mirror.com" +PROXY = os.getenv("PROXY", "socks5h://127.0.0.1:1081") + +# ---------- 提前设置缓存目录(在 import 前) ---------- +os.environ["HF_HOME"] = str(HF_HOME) + +# ---------- 依赖安装 ---------- +def _install(pkgs: str): + os.system(f"{sys.executable} -m pip install -U {pkgs}") + +try: + import requests +except ImportError: + _install("requests") + import requests + +# 若走 socks 代理需要 PySocks +if PROXY and "socks" in PROXY: + try: + import socks # noqa: F401 + except ImportError: + _install("'requests[socks]'") + +# ---------- 选择网络模式(镜像 → 代理 → 官方) ---------- +def set_network_mode(): + # 若外部已设置 HF_ENDPOINT,尊重外部配置 + if os.getenv("HF_ENDPOINT"): + print(f"🌏 Using custom HF endpoint: {os.getenv('HF_ENDPOINT')}") + return + + # 1) 尝试 CN 镜像 + try: + r = requests.get(CN_MIRROR, timeout=2) + if r.status_code == 200: + os.environ["HF_ENDPOINT"] = CN_MIRROR + print(f"🌏 Using Hugging Face CN mirror: {CN_MIRROR}") + return + except Exception: + pass + + # 2) 走代理 + if PROXY: + os.environ["HTTP_PROXY"] = PROXY + os.environ["HTTPS_PROXY"] = PROXY + print(f"🌐 Using proxy: {PROXY}") + return + + # 3) 官方直连 + print("⚠️ No mirror or proxy, using official huggingface.co") + +set_network_mode() + +# 现在再导入 huggingface_hub,确保拿到正确的 endpoint/proxy 设置 try: from huggingface_hub import snapshot_download except ImportError: - print("📦 Installing huggingface_hub...") - os.system(f"{sys.executable} -m pip install -U huggingface_hub") + _install("'huggingface_hub[tqdm]'") from huggingface_hub import snapshot_download -# ==== 默认配置 ==== -DEFAULT_MODEL_ID = "BAAI/bge-m3" -DEFAULT_MODEL_DIR = "models/bge-m3" -DEFAULT_PROXY = "socks5://127.0.0.1:1080" - -# ==== 从环境变量读取 ==== -MODEL_ID = os.environ.get("MODEL_ID", DEFAULT_MODEL_ID) -MODEL_DIR = Path(os.environ.get("MODEL_DIR", DEFAULT_MODEL_DIR)) -PROXY = os.environ.get("PROXY", DEFAULT_PROXY) - -# ==== 设置代理 ==== -if PROXY: - os.environ["HTTP_PROXY"] = PROXY - os.environ["HTTPS_PROXY"] = PROXY - print(f"🌐 Using proxy: {PROXY}") -else: - print("🚫 No proxy configured, direct connection.") - -# ==== 创建保存目录 ==== -MODEL_DIR.parent.mkdir(parents=True, exist_ok=True) - -# ==== 下载模型 ==== -print(f"⬇️ Downloading model from Hugging Face...") -print(f" Model ID: {MODEL_ID}") -print(f" Save dir: {MODEL_DIR}") - -snapshot_download( - repo_id=MODEL_ID, - local_dir=str(MODEL_DIR), - local_dir_use_symlinks=False +# ---------- 工具函数 ---------- +KEY_FILES = ( + "tokenizer.json", + "config.json", + "sentencepiece.bpe.model", + "onnx/model.onnx", + "pytorch_model.bin", + "model.safetensors", ) -print(f"✅ Model cached to {MODEL_DIR}") +def has_local_model(root: Path) -> bool: + if not root.exists(): + return False + for k in KEY_FILES: + if any(root.rglob(k)): + return True + # 兜底:只要非空也算有内容(对应部分仓库布局) + return any(root.iterdir()) +# ---------- 主流程 ---------- +def main(): + print("⬇️ Downloading model from Hugging Face…") + print(f" Model ID : {MODEL_ID}") + print(f" Save dir : {MODEL_DIR}") + print(f" HF_HOME : {HF_HOME}") + if os.getenv("HF_ENDPOINT"): + print(f" Endpoint : {os.getenv('HF_ENDPOINT')}") + elif os.getenv("HTTP_PROXY"): + print(f" Proxy : {os.getenv('HTTP_PROXY')}") + else: + print(" Endpoint : official (huggingface.co)") + + MODEL_DIR.parent.mkdir(parents=True, exist_ok=True) + HF_HOME.mkdir(parents=True, exist_ok=True) + + # 已有可用文件 → 跳过下载 + if has_local_model(MODEL_DIR): + print(f"📂 Local model exists, skip download: {MODEL_DIR}") + print("💡 To force re-download, remove the folder and rerun.") + return + + # 下载(显示进度) + try: + snapshot_download( + repo_id=MODEL_ID, + local_dir=str(MODEL_DIR), + local_dir_use_symlinks=False, + tqdm_class=None, # 使用默认 tqdm 进度条 + ) + except Exception as e: + # 失败时检查是否已经有部分或全部文件 + if has_local_model(MODEL_DIR): + print(f"⚠️ Online fetch failed but local files exist: {MODEL_DIR}") + print(f" Error: {e}") + else: + print("❌ Download failed and no local files found.") + print(f" Error: {e}") + print("🔁 Try: 1) 切换镜像/代理 2) 检查网络 3) 稍后重试") + sys.exit(1) + + # 最终确认 + if has_local_model(MODEL_DIR): + print(f"✅ Model cached to {MODEL_DIR}") + print("💡 To run offline later, set: export HF_HUB_OFFLINE=1") + else: + print("❌ No model files found after download attempt.") + sys.exit(1) + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/docs/offline_embed_server.py b/docs/offline_embed_server.py index 711c5c3..be4b00b 100644 --- a/docs/offline_embed_server.py +++ b/docs/offline_embed_server.py @@ -2,44 +2,134 @@ import os, sys, numpy as np from pathlib import Path -# 自动装依赖 +# ==== 自动安装依赖 ==== +def install(pkg): + os.system(f"{sys.executable} -m pip install -U {pkg}") + try: from flask import Flask, request, jsonify from fastembed import TextEmbedding from huggingface_hub import snapshot_download except ImportError: - os.system(f"{sys.executable} -m pip install -U flask fastembed numpy huggingface_hub") + install("flask fastembed numpy huggingface_hub") from flask import Flask, request, jsonify from fastembed import TextEmbedding from huggingface_hub import snapshot_download -# 模型路径 -MODEL_DIR = Path(os.getenv("BGE_M3_DIR", "models/bge-m3")) +# ==== 配置 ==== +MODEL_ID = os.getenv("MODEL_ID", "BAAI/bge-m3") +MODEL_DIR = Path(os.getenv("MODEL_DIR", "models/bge-m3")) +HF_HOME = Path(os.getenv("HF_HOME", Path.cwd() / "hf_cache")) -# 如果本地无模型,先下载 -if not MODEL_DIR.exists(): - print(f"⬇️ Downloading BGE-M3 to {MODEL_DIR} ...") - snapshot_download("BAAI/bge-m3", local_dir=str(MODEL_DIR), local_dir_use_symlinks=False) +# 设置 HF 缓存目录 +os.environ["HF_HOME"] = str(HF_HOME) -# 离线模式 -os.environ["HF_HOME"] = str(Path.cwd() / "hf_cache") +# ==== 确保模型已存在(不在这里判断支持性) ==== +if not MODEL_DIR.exists() or not any(MODEL_DIR.iterdir()): + print(f"⬇️ Downloading model {MODEL_ID} to {MODEL_DIR} ...") + snapshot_download(repo_id=MODEL_ID, local_dir=str(MODEL_DIR), local_dir_use_symlinks=False) + +# ==== 离线模式 ==== os.environ["HF_HUB_OFFLINE"] = "1" -# 启动服务 +# ==== 启动 Flask 服务 ==== app = Flask(__name__) -model = TextEmbedding(str(MODEL_DIR)) + +# 这里用模型 ID 初始化,而不是路径 +model = TextEmbedding(MODEL_ID) +DIM = 1024 # bge-m3 的维度 @app.post("/v1/embeddings") def embeddings(): data = request.get_json(force=True) or {} texts = [data["input"]] if isinstance(data.get("input"), str) else data.get("input", []) vecs = [np.asarray(v, np.float32) / (np.linalg.norm(v) + 1e-12) for v in model.embed(texts)] - return jsonify({"object": "list", "data": [ - {"object": "embedding", "index": i, "embedding": v.tolist()} for i, v in enumerate(vecs) - ], "model": data.get("model", "BAAI/bge-m3")}) + return jsonify({ + "object": "list", + "data": [ + {"object": "embedding", "index": i, "embedding": v.tolist()} + for i, v in enumerate(vecs) + ], + "model": data.get("model", MODEL_ID) + }) @app.get("/healthz") -def healthz(): return "ok", 200 +def healthz(): + return "ok", 200 if __name__ == "__main__": - app.run(host=os.getenv("EMBED_HOST", "0.0.0.0"), port=int(os.getenv("EMBED_PORT", 9000))) + host = os.getenv("EMBED_HOST", "0.0.0.0") + port = int(os.getenv("EMBED_PORT", 9000)) + print(f"🚀 Starting embedding server on http://{host}:{port}") + print(f" Model: {MODEL_ID}") + print(f" Cache dir: {HF_HOME}") + app.run(host=host, port=port) +shenlan@MacBook-Pro-3 XControl % clear +shenlan@MacBook-Pro-3 XControl % +shenlan@MacBook-Pro-3 XControl % cat docs/offline_embed_server.py +#!/usr/bin/env python3 +import os, sys, numpy as np +from pathlib import Path + +# ==== 自动安装依赖 ==== +def install(pkg): + os.system(f"{sys.executable} -m pip install -U {pkg}") + +try: + from flask import Flask, request, jsonify + from fastembed import TextEmbedding + from huggingface_hub import snapshot_download +except ImportError: + install("flask fastembed numpy huggingface_hub") + from flask import Flask, request, jsonify + from fastembed import TextEmbedding + from huggingface_hub import snapshot_download + +# ==== 配置 ==== +MODEL_ID = os.getenv("MODEL_ID", "BAAI/bge-m3") +MODEL_DIR = Path(os.getenv("MODEL_DIR", "models/bge-m3")) +HF_HOME = Path(os.getenv("HF_HOME", Path.cwd() / "hf_cache")) + +# 设置 HF 缓存目录 +os.environ["HF_HOME"] = str(HF_HOME) + +# ==== 确保模型已存在(不在这里判断支持性) ==== +if not MODEL_DIR.exists() or not any(MODEL_DIR.iterdir()): + print(f"⬇️ Downloading model {MODEL_ID} to {MODEL_DIR} ...") + snapshot_download(repo_id=MODEL_ID, local_dir=str(MODEL_DIR), local_dir_use_symlinks=False) + +# ==== 离线模式 ==== +os.environ["HF_HUB_OFFLINE"] = "1" + +# ==== 启动 Flask 服务 ==== +app = Flask(__name__) + +# 这里用模型 ID 初始化,而不是路径 +model = TextEmbedding(MODEL_ID) +DIM = 1024 # bge-m3 的维度 + +@app.post("/v1/embeddings") +def embeddings(): + data = request.get_json(force=True) or {} + texts = [data["input"]] if isinstance(data.get("input"), str) else data.get("input", []) + vecs = [np.asarray(v, np.float32) / (np.linalg.norm(v) + 1e-12) for v in model.embed(texts)] + return jsonify({ + "object": "list", + "data": [ + {"object": "embedding", "index": i, "embedding": v.tolist()} + for i, v in enumerate(vecs) + ], + "model": data.get("model", MODEL_ID) + }) + +@app.get("/healthz") +def healthz(): + return "ok", 200 + +if __name__ == "__main__": + host = os.getenv("EMBED_HOST", "0.0.0.0") + port = int(os.getenv("EMBED_PORT", 9000)) + print(f"🚀 Starting embedding server on http://{host}:{port}") + print(f" Model: {MODEL_ID}") + print(f" Cache dir: {HF_HOME}") + app.run(host=host, port=port) \ No newline at end of file