Make the heavy ML stack (torch/transformers/sentence-transformers) an opt-in install
- Dominant language
- Python
- Stars
- 52
- Forks
- 18
- PR merge metrics
- No merged PRs in 30d
Description
## 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".
Contributor guide
Assessment
This issue has not been assessed yet.