91e6c47fc0
统一 voc_llm 密钥解析与默认模型；向量化改为本地 mlx 模型；更新 README、gitignore 与流水线文档。 Co-authored-by: Cursor <cursoragent@cursor.com>
144 lines
4.3 KiB
Python
144 lines
4.3 KiB
Python
"""
|
||
本地 MLX Qwen3 Embedding（Apple Silicon）。
|
||
|
||
默认模型目录：项目根 ``Qwen3-Embedding-4B-mxfp8``，可用 ``VOC_EMBED_MODEL_PATH`` 覆盖。
|
||
|
||
在 M4 / 16GB、约 8GB 可用内存下实测（mxfp8，500 字符/条）：
|
||
- 加载峰值约 1.5GB
|
||
- batch 1–16 稳定；默认 ``VOC_EMBED_BATCH_SIZE=16``，串行推理（勿多线程并行加载同一 MLX 模型）
|
||
|
||
需 Python 3.10+ 与 ``mlx-embeddings>=0.1.0``（推荐项目 ``310py`` 虚拟环境）。
|
||
"""
|
||
from __future__ import annotations
|
||
|
||
import logging
|
||
import os
|
||
import sys
|
||
import types
|
||
from pathlib import Path
|
||
from typing import List, Sequence
|
||
|
||
logger = logging.getLogger("local_embedding")
|
||
|
||
PROJECT_ROOT = Path(__file__).resolve().parent
|
||
DEFAULT_MODEL_PATH = PROJECT_ROOT / "Qwen3-Embedding-4B-mxfp8"
|
||
# 实测短句 batch=32 仍 <0.5GB增量；长句 500 字符 batch=16 约 1.5GB 峰值
|
||
DEFAULT_BATCH_SIZE = 16
|
||
DEFAULT_MAX_TEXT_CHARS = 512
|
||
|
||
|
||
def _apply_hf_hub_shim() -> None:
|
||
if "huggingface_hub.utils._errors" in sys.modules:
|
||
return
|
||
try:
|
||
from huggingface_hub.errors import RepositoryNotFoundError
|
||
except ImportError:
|
||
try:
|
||
from huggingface_hub.utils._errors import RepositoryNotFoundError # type: ignore
|
||
except ImportError:
|
||
RepositoryNotFoundError = Exception # type: ignore[misc, assignment]
|
||
mod = types.ModuleType("huggingface_hub.utils._errors")
|
||
mod.RepositoryNotFoundError = RepositoryNotFoundError
|
||
sys.modules["huggingface_hub.utils._errors"] = mod
|
||
|
||
|
||
def _resolve_model_path() -> Path:
|
||
raw = os.environ.get("VOC_EMBED_MODEL_PATH", "").strip()
|
||
p = Path(raw).expanduser() if raw else DEFAULT_MODEL_PATH
|
||
if not p.is_dir():
|
||
raise FileNotFoundError(f"本地 embedding 模型目录不存在: {p}")
|
||
return p.resolve()
|
||
|
||
|
||
def _resolve_batch_size(explicit: int | None = None) -> int:
|
||
if explicit is not None and explicit > 0:
|
||
return explicit
|
||
env = os.environ.get("VOC_EMBED_BATCH_SIZE", "").strip()
|
||
if env.isdigit() and int(env) > 0:
|
||
return int(env)
|
||
return DEFAULT_BATCH_SIZE
|
||
|
||
|
||
def _resolve_max_chars() -> int:
|
||
env = os.environ.get("VOC_EMBED_MAX_TEXT_CHARS", "").strip()
|
||
if env.isdigit() and int(env) > 0:
|
||
return int(env)
|
||
return DEFAULT_MAX_TEXT_CHARS
|
||
|
||
|
||
def _truncate(text: str, max_chars: int) -> str:
|
||
t = (text or "").strip()
|
||
if len(t) <= max_chars:
|
||
return t
|
||
return t[: max_chars - 3] + "..."
|
||
|
||
|
||
_MODEL = None
|
||
_PROCESSOR = None
|
||
_DIMENSIONS: int | None = None
|
||
|
||
|
||
def _load():
|
||
global _MODEL, _PROCESSOR
|
||
if _MODEL is not None:
|
||
return _MODEL, _PROCESSOR
|
||
_apply_hf_hub_shim()
|
||
from mlx_embeddings import load
|
||
|
||
path = _resolve_model_path()
|
||
logger.info("加载本地 embedding: %s", path)
|
||
_MODEL, _PROCESSOR = load(str(path))
|
||
return _MODEL, _PROCESSOR
|
||
|
||
|
||
def embedding_dimensions() -> int:
|
||
global _DIMENSIONS
|
||
if _DIMENSIONS is not None:
|
||
return _DIMENSIONS
|
||
vecs = embed_texts(["dimension probe"])
|
||
_DIMENSIONS = len(vecs[0])
|
||
return _DIMENSIONS
|
||
|
||
|
||
def embed_texts(
|
||
texts: Sequence[str],
|
||
*,
|
||
batch_size: int | None = None,
|
||
max_chars: int | None = None,
|
||
) -> List[List[float]]:
|
||
"""返回与输入等长的浮点向量列表（已 L2 归一化）。"""
|
||
if not texts:
|
||
return []
|
||
from mlx_embeddings import generate
|
||
|
||
model, processor = _load()
|
||
bs = _resolve_batch_size(batch_size)
|
||
cap = max_chars if max_chars is not None else _resolve_max_chars()
|
||
cleaned = [_truncate(t, cap) for t in texts]
|
||
|
||
out_all: List[List[float]] = []
|
||
n = len(cleaned)
|
||
n_chunks = (n + bs - 1) // bs
|
||
log_every = max(1, n_chunks // 20)
|
||
|
||
for chunk_idx, i in enumerate(range(0, n, bs)):
|
||
chunk = cleaned[i : i + bs]
|
||
output = generate(model, processor, texts=chunk)
|
||
emb = output.text_embeds
|
||
for row in emb:
|
||
out_all.append([float(x) for x in row.tolist()])
|
||
|
||
done = chunk_idx + 1
|
||
if done == 1 or done == n_chunks or done % log_every == 0:
|
||
rows_done = min(done * bs, n)
|
||
logger.info(
|
||
"向量化进度 %s/%s 批（%s/%s 条）",
|
||
done,
|
||
n_chunks,
|
||
rows_done,
|
||
n,
|
||
)
|
||
global _DIMENSIONS
|
||
if _DIMENSIONS is None and out_all:
|
||
_DIMENSIONS = len(out_all[0])
|
||
return out_all
|