afx-team / afx-team/hebb-mind

Make the heavy ML stack (torch/transformers/sentence-transformers) an opt-in install

未關閉
#42 0 則留言 0 個 reaction 已指派 0 人 在 GitHub 檢視
enhancement
主要語言
Python
星號
52
分支
18
PR 合併指標
30 天內沒有已合併 PR

描述

## Problem

A default `pip install hebb-mind` eagerly pulls the full ML stack — `sentence-transformers` and its transitive deps `torch` / `transformers` / `huggingface_hub` / `tokenizers` / `safetensors`. On Linux the default `torch` wheel bundles CUDA and the install can reach **~2–2.5 GB**, even for users who never run a local model (e.g. they intend to use an API embedding provider).

The root cause is a single line: `sentence-transformers>=3.0.0` is declared as a **core** dependency in `pyproject.toml` (`[project.dependencies]`).

### Two different "downloads" — only one is the problem

| Type | What | When | Status |
|---|---|---|---|
| **A. pip packages** (torch/transformers/sentence-transformers) | ~2–2.5 GB (CUDA) | at `pip install` | ❌ eager — **this is the pain** |
| **B. model weights** (all-MiniLM-L6-v2, bge-reranker-base) | 90 MB – 2 GB | first use / `hebb setup` | ✅ already on-demand |

So "download only when needed" already works for model weights (B). The gap is the install-time pip packages (A).

## Why this is cheap to fix

The codebase is already structured for it:

- **All heavy imports are already lazy** (function/method-level) — `import hebb` is sub-second and never imports torch. Sites: `embedding/local.py:135` (`SentenceTransformer`), `retrieval/rerank/local.py:58-59` (`CrossEncoder`, `torch.nn.Sigmoid`).
- **Factory layers already degrade gracefully** — `embedding/factory.py:57-65` falls back to `NoopEmbedder`; `retrieval/rerank/factory.py:32-44` returns `None`. A lean install won't crash; it just disables vector/rerank.
- **The exact pattern to copy already exists** — the `[pg]` extra: `storage/factory.py:70-73` does `try: import asyncpg / except ImportError: raise ImportError("... pip install hebb-mind[pg]")`.
- **An API-only path already exists** that needs neither torch nor sentence-transformers — `embedding_provider="api"` (litellm / custom HTTP) + `rerank_enabled=false`.

## Proposed solution

1. **Packaging** — move `sentence-transformers` out of core deps into a new `local` extra:
```toml
[project.optional-dependencies]
local = ["sentence-transformers>=3.0.0"]
```
`pip install hebb-mind` → lean. `pip install hebb-mind[local]` → full local stack. Document/use the CPU index for torch (`--extra-index-url https://download.pytorch.org/whl/cpu`) to cut ~2.5 GB → ~250 MB on Linux.

2. **Import guards** — wrap the lazy imports (`embedding/local.py`, `retrieval/rerank/local.py`) in `try/except ModuleNotFoundError` with an actionable message; and make the factory `except` blocks special-case the missing stack so degradation is **loud and actionable** instead of a generic warning.

3. **`hebb setup` owns the install** (per the project's User Path Ownership principle) — when the local provider is selected and the stack is missing, `hebb setup` runs `pip install` for the user (subprocess + progress, CPU index), right alongside where it already downloads model weights (`cli/commands/setup.py`). User path becomes: `pip install hebb-mind` (fast) → `hebb setup` (installs exactly what the chosen config needs).

### Explicitly NOT doing
- Flipping the default provider to API — would change the published benchmark provenance (LoCoMo / MemBench / LongMemEval are all measured on the local `all-MiniLM-384 + bge-reranker-base` default) and the zero-config local-first experience. Keep local default; only move *when* the stack installs.
- Silent runtime `pip install` at arbitrary call sites — fragile (CPU/CUDA/index-url matrix) and surprising. Confine auto-install to `hebb setup`.

## Acceptance criteria
- `pip install hebb-mind` no longer pulls torch/transformers/sentence-transformers; install size drops to the ~100 MB range.
- Lean install: `import hebb` works; API-provider path works end-to-end; local-provider gives a clear, actionable message instead of silently disabling vector search.
- `hebb setup` (local provider) auto-installs the CPU ML stack + model weights with no manual pip/torch command.
- CI (incl. `slow` tests + mypy strict) green under `[dev,local]`; default eval config and published numbers unchanged.

## Notes for implementers
- Add `local` to the `dev` extra (or run CI as `.[dev,local]`) so `slow` tests + mypy still have the stack.
- Add `sentence_transformers.*` / `torch.*` / `transformers.*` to `[tool.mypy.overrides] ignore_missing_imports`.
- Backward-compat: existing users upgrading will lose the local stack — call this out in the release notes; consider a transitional `full`/`all` extra = `[local,pg]`.
- `hebb doctor` should diagnose "local provider but ML stack missing".

貢獻指南

開啟貢獻指南

評估

這個 Issue 還沒有評估資料。

把新 issue 寄到你的電子郵件信箱

精選適合新手參與的 GitHub issue 摘要。