afx-team / afx-team/hebb-mind

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

Abierto
#42 0 comentarios 0 reacciones 0 asignados Ver en GitHub
enhancement
Lenguaje dominante
Python
Estrellas
52
Forks
18
Métricas de merge de PR
Sin PR fusionados en 30 d

Descripción

## 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".

Guía de contribución

Abrir la guía de contribución

Evaluación

Este issue todavía no se ha evaluado.

Recibe los nuevos issues en tu correo

Un resumen breve de issues de GitHub para principiantes.