mirror of
https://github.com/catlog22/Claude-Code-Workflow.git
synced 2026-03-18 18:48:48 +08:00
feat: add model download manager with HF mirror support and fix defaults
- Add lightweight model_manager.py: cache detection (with fastembed name remapping), HF mirror download via huggingface_hub, auto model.onnx fallback from quantized variants - Config defaults: embed_model -> bge-small-en-v1.5 (384d), reranker -> Xenova/ms-marco-MiniLM-L-6-v2 (fastembed 0.7.4 compatible) - Add model_cache_dir and hf_mirror config options - embed/local.py and rerank/local.py use model_manager for cache-aware loading - Fix FastEmbedReranker to handle both float list and RerankResult formats - E2E test uses real FastEmbedReranker instead of mock KeywordReranker Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
@@ -8,10 +8,14 @@ log = logging.getLogger(__name__)
|
||||
@dataclass
|
||||
class Config:
|
||||
# Embedding
|
||||
embed_model: str = "jinaai/jina-embeddings-v2-base-code"
|
||||
embed_dim: int = 768
|
||||
embed_model: str = "BAAI/bge-small-en-v1.5"
|
||||
embed_dim: int = 384
|
||||
embed_batch_size: int = 64
|
||||
|
||||
# Model download / cache
|
||||
model_cache_dir: str = "" # empty = fastembed default cache
|
||||
hf_mirror: str = "" # HuggingFace mirror URL, e.g. "https://hf-mirror.com"
|
||||
|
||||
# GPU / execution providers
|
||||
device: str = "auto" # 'auto', 'cuda', 'cpu'
|
||||
embed_providers: list[str] | None = None # explicit ONNX providers override
|
||||
@@ -35,7 +39,7 @@ class Config:
|
||||
ann_top_k: int = 50
|
||||
|
||||
# Reranker
|
||||
reranker_model: str = "BAAI/bge-reranker-v2-m3"
|
||||
reranker_model: str = "Xenova/ms-marco-MiniLM-L-6-v2"
|
||||
reranker_top_k: int = 20
|
||||
reranker_batch_size: int = 32
|
||||
|
||||
|
||||
@@ -24,16 +24,23 @@ class FastEmbedEmbedder(BaseEmbedder):
|
||||
"""Lazy-load the fastembed TextEmbedding model on first use."""
|
||||
if self._model is not None:
|
||||
return
|
||||
from .. import model_manager
|
||||
model_manager.ensure_model(self._config.embed_model, self._config)
|
||||
|
||||
from fastembed import TextEmbedding
|
||||
providers = self._config.resolve_embed_providers()
|
||||
cache_kwargs = model_manager.get_cache_kwargs(self._config)
|
||||
try:
|
||||
self._model = TextEmbedding(
|
||||
model_name=self._config.embed_model,
|
||||
providers=providers,
|
||||
**cache_kwargs,
|
||||
)
|
||||
except TypeError:
|
||||
# Older fastembed versions may not accept providers kwarg
|
||||
self._model = TextEmbedding(model_name=self._config.embed_model)
|
||||
self._model = TextEmbedding(
|
||||
model_name=self._config.embed_model,
|
||||
**cache_kwargs,
|
||||
)
|
||||
|
||||
def embed_single(self, text: str) -> np.ndarray:
|
||||
"""Embed a single text, returns float32 ndarray of shape (dim,)."""
|
||||
|
||||
145
codex-lens-v2/src/codexlens_search/model_manager.py
Normal file
145
codex-lens-v2/src/codexlens_search/model_manager.py
Normal file
@@ -0,0 +1,145 @@
|
||||
"""Lightweight model download manager for fastembed models.
|
||||
|
||||
Handles HuggingFace mirror configuration and cache pre-population so that
|
||||
fastembed can load models from local cache without network access.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import os
|
||||
from pathlib import Path
|
||||
|
||||
from .config import Config
|
||||
|
||||
log = logging.getLogger(__name__)
|
||||
|
||||
# Models that fastembed maps internally (HF repo may differ from model_name)
|
||||
_EMBED_MODEL_FILES = ["*.onnx", "*.json"]
|
||||
_RERANK_MODEL_FILES = ["*.onnx", "*.json"]
|
||||
|
||||
|
||||
def _resolve_cache_dir(config: Config) -> str | None:
|
||||
"""Return cache_dir for fastembed, or None for default."""
|
||||
return config.model_cache_dir or None
|
||||
|
||||
|
||||
def _apply_mirror(config: Config) -> None:
|
||||
"""Set HF_ENDPOINT env var if mirror is configured."""
|
||||
if config.hf_mirror:
|
||||
os.environ["HF_ENDPOINT"] = config.hf_mirror
|
||||
|
||||
|
||||
def _model_is_cached(model_name: str, cache_dir: str | None) -> bool:
|
||||
"""Check if a model already exists in the fastembed/HF hub cache.
|
||||
|
||||
Note: fastembed may remap model names internally (e.g. BAAI/bge-small-en-v1.5
|
||||
-> qdrant/bge-small-en-v1.5-onnx-q), so we also search by partial name match.
|
||||
"""
|
||||
base = cache_dir or _default_fastembed_cache()
|
||||
base_path = Path(base)
|
||||
if not base_path.exists():
|
||||
return False
|
||||
|
||||
# Exact match first
|
||||
safe_name = model_name.replace("/", "--")
|
||||
model_dir = base_path / f"models--{safe_name}"
|
||||
if _dir_has_onnx(model_dir):
|
||||
return True
|
||||
|
||||
# Partial match: fastembed remaps some model names
|
||||
short_name = model_name.split("/")[-1].lower()
|
||||
for d in base_path.iterdir():
|
||||
if short_name in d.name.lower() and _dir_has_onnx(d):
|
||||
return True
|
||||
|
||||
return False
|
||||
|
||||
|
||||
def _dir_has_onnx(model_dir: Path) -> bool:
|
||||
"""Check if a model directory has at least one ONNX file in snapshots."""
|
||||
snapshots = model_dir / "snapshots"
|
||||
if not snapshots.exists():
|
||||
return False
|
||||
for snap in snapshots.iterdir():
|
||||
if list(snap.rglob("*.onnx")):
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
def _default_fastembed_cache() -> str:
|
||||
"""Return fastembed's default cache directory."""
|
||||
return os.path.join(os.environ.get("TMPDIR", os.path.join(
|
||||
os.environ.get("LOCALAPPDATA", os.path.expanduser("~")),
|
||||
)), "fastembed_cache")
|
||||
|
||||
|
||||
def ensure_model(model_name: str, config: Config) -> None:
|
||||
"""Ensure a model is available in the local cache.
|
||||
|
||||
If the model is already cached, this is a no-op.
|
||||
If not cached, attempts to download via huggingface_hub with mirror support.
|
||||
"""
|
||||
cache_dir = _resolve_cache_dir(config)
|
||||
|
||||
if _model_is_cached(model_name, cache_dir):
|
||||
log.debug("Model %s found in cache", model_name)
|
||||
return
|
||||
|
||||
log.info("Model %s not in cache, downloading...", model_name)
|
||||
_apply_mirror(config)
|
||||
|
||||
try:
|
||||
from huggingface_hub import snapshot_download
|
||||
|
||||
kwargs: dict = {
|
||||
"repo_id": model_name,
|
||||
"allow_patterns": ["*.onnx", "*.json"],
|
||||
}
|
||||
if cache_dir:
|
||||
kwargs["cache_dir"] = cache_dir
|
||||
if config.hf_mirror:
|
||||
kwargs["endpoint"] = config.hf_mirror
|
||||
|
||||
path = snapshot_download(**kwargs)
|
||||
log.info("Model %s downloaded to %s", model_name, path)
|
||||
|
||||
# fastembed for some reranker models expects model.onnx but repo may
|
||||
# only have quantized variants. Create a symlink/copy if needed.
|
||||
_ensure_model_onnx(Path(path))
|
||||
|
||||
except ImportError:
|
||||
log.warning(
|
||||
"huggingface_hub not installed. Cannot download models. "
|
||||
"Install with: pip install huggingface-hub"
|
||||
)
|
||||
except Exception as e:
|
||||
log.warning("Failed to download model %s: %s", model_name, e)
|
||||
|
||||
|
||||
def _ensure_model_onnx(model_dir: Path) -> None:
|
||||
"""If model.onnx is missing but a quantized variant exists, copy it."""
|
||||
onnx_dir = model_dir / "onnx"
|
||||
if not onnx_dir.exists():
|
||||
onnx_dir = model_dir # some models put onnx at root
|
||||
|
||||
target = onnx_dir / "model.onnx"
|
||||
if target.exists():
|
||||
return
|
||||
|
||||
# Look for quantized alternatives
|
||||
for name in ["model_quantized.onnx", "model_optimized.onnx",
|
||||
"model_int8.onnx", "model_uint8.onnx"]:
|
||||
candidate = onnx_dir / name
|
||||
if candidate.exists():
|
||||
import shutil
|
||||
shutil.copy2(candidate, target)
|
||||
log.info("Copied %s -> model.onnx", name)
|
||||
return
|
||||
|
||||
|
||||
def get_cache_kwargs(config: Config) -> dict:
|
||||
"""Return kwargs to pass to fastembed constructors for cache_dir."""
|
||||
cache_dir = _resolve_cache_dir(config)
|
||||
if cache_dir:
|
||||
return {"cache_dir": cache_dir}
|
||||
return {}
|
||||
@@ -13,12 +13,26 @@ class FastEmbedReranker(BaseReranker):
|
||||
|
||||
def _load(self) -> None:
|
||||
if self._model is None:
|
||||
from .. import model_manager
|
||||
model_manager.ensure_model(self._config.reranker_model, self._config)
|
||||
|
||||
from fastembed.rerank.cross_encoder import TextCrossEncoder
|
||||
self._model = TextCrossEncoder(model_name=self._config.reranker_model)
|
||||
cache_kwargs = model_manager.get_cache_kwargs(self._config)
|
||||
self._model = TextCrossEncoder(
|
||||
model_name=self._config.reranker_model,
|
||||
**cache_kwargs,
|
||||
)
|
||||
|
||||
def score_pairs(self, query: str, documents: list[str]) -> list[float]:
|
||||
self._load()
|
||||
results = list(self._model.rerank(query, documents))
|
||||
if not results:
|
||||
return [0.0] * len(documents)
|
||||
# fastembed may return list[float] or list[RerankResult] depending on version
|
||||
first = results[0]
|
||||
if isinstance(first, (int, float)):
|
||||
return [float(s) for s in results]
|
||||
# Older format: objects with .index and .score
|
||||
scores = [0.0] * len(documents)
|
||||
for r in results:
|
||||
scores[r.index] = float(r.score)
|
||||
|
||||
Reference in New Issue
Block a user