Spaces:
Sleeping
Sleeping
docs: design spec for ONNX lightweight migration
Browse filesReplace PyTorch + sentence-transformers (~2GB) with ONNX Runtime direct
loading (~50MB, already a chromadb dep). Switch bge-large to bge-small,
drop cross-encoder reranker. Target: HF Spaces deployable image ~400MB.
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
docs/superpowers/specs/2026-03-30-onnx-lightweight-migration-design.md
ADDED
|
@@ -0,0 +1,91 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# Design: Migration ONNX Lightweight
|
| 2 |
+
|
| 3 |
+
## Problème
|
| 4 |
+
|
| 5 |
+
Le build HuggingFace Spaces échoue — l'image Docker pèse ~2.7 GB, principalement à cause de PyTorch (800 MB) + bge-large (1.34 GB) + sentence-transformers/transformers (200 MB). Le free tier HF Spaces ne supporte pas cette taille.
|
| 6 |
+
|
| 7 |
+
## Décision
|
| 8 |
+
|
| 9 |
+
Remplacer la stack embedding lourde (sentence-transformers → PyTorch) par un chargement ONNX Runtime direct, et passer de bge-large à bge-small.
|
| 10 |
+
|
| 11 |
+
## Changements
|
| 12 |
+
|
| 13 |
+
### Dépendances
|
| 14 |
+
|
| 15 |
+
**Supprimé :**
|
| 16 |
+
- `sentence-transformers` (et ses transitives : `torch`, `transformers`, `scipy`, `scikit-learn`)
|
| 17 |
+
|
| 18 |
+
**Ajouté :**
|
| 19 |
+
- Rien — `onnxruntime` et `tokenizers` sont déjà des dépendances de `chromadb 1.5.5`
|
| 20 |
+
|
| 21 |
+
**Impact taille :** ~2.7 GB → ~400 MB estimé
|
| 22 |
+
|
| 23 |
+
### Modèle d'embedding
|
| 24 |
+
|
| 25 |
+
| | Avant | Après |
|
| 26 |
+
|---|---|---|
|
| 27 |
+
| Modèle | BAAI/bge-large-en-v1.5 | BAAI/bge-small-en-v1.5 |
|
| 28 |
+
| Dimensions | 1024 | 384 |
|
| 29 |
+
| Taille | 1.34 GB | ~133 MB (ou ~35 MB quantifié int8) |
|
| 30 |
+
| MTEB score | 64.23 | 62.17 |
|
| 31 |
+
| Moteur | PyTorch via sentence-transformers | ONNX Runtime direct |
|
| 32 |
+
|
| 33 |
+
### Reranker
|
| 34 |
+
|
| 35 |
+
**Supprimé.** Le cross-encoder `ms-marco-MiniLM-L-12-v2` est retiré.
|
| 36 |
+
|
| 37 |
+
**Justification :** Avec 350 stories et hybrid retrieval (BM25 + dense + RRF), le reranker apporte ~3-5% de qualité (NDCG 0.85 → 0.82 sans). Le consensus expert situe le seuil d'utilité à ~1000 docs. Le gain en taille (-900 MB) et latence (-250-400ms) justifie la suppression.
|
| 38 |
+
|
| 39 |
+
**Filet de sécurité :** Si la qualité baisse, rerank via Gemini Flash sur les top-5 (appel API, pas de dep lourde).
|
| 40 |
+
|
| 41 |
+
## Fichiers impactés
|
| 42 |
+
|
| 43 |
+
| Fichier | Action |
|
| 44 |
+
|---|---|
|
| 45 |
+
| `src/mediastorm/config.py` | Modèle → bge-small, dimension → 384, ajout ONNX_MODEL_PATH |
|
| 46 |
+
| `src/mediastorm/vectorize/embedder.py` | Réécriture : onnxruntime.InferenceSession + tokenizers + mean pooling + L2 norm |
|
| 47 |
+
| `src/mediastorm/rag/reranker.py` | Supprimé |
|
| 48 |
+
| `src/mediastorm/rag/retriever.py` | Retirer l'appel au reranker |
|
| 49 |
+
| `pyproject.toml` | Retirer `sentence-transformers` |
|
| 50 |
+
| `Dockerfile` | Supprimer install PyTorch, ajouter COPY models/ |
|
| 51 |
+
| `tests/test_embedder.py` | Adapter assertions (1024 → 384) |
|
| 52 |
+
| `models/bge-small-en-v1.5/` | Nouveau dossier : model.onnx + tokenizer.json |
|
| 53 |
+
|
| 54 |
+
## Embedder ONNX — interface
|
| 55 |
+
|
| 56 |
+
L'interface publique reste identique :
|
| 57 |
+
|
| 58 |
+
```python
|
| 59 |
+
class Embedder:
|
| 60 |
+
def __init__(self):
|
| 61 |
+
# Charge tokenizer.json via tokenizers.Tokenizer
|
| 62 |
+
# Charge model.onnx via onnxruntime.InferenceSession
|
| 63 |
+
|
| 64 |
+
def embed_texts(self, texts: list[str]) -> list[list[float]]:
|
| 65 |
+
# Tokenize → inference ONNX → mean pooling → L2 normalize
|
| 66 |
+
```
|
| 67 |
+
|
| 68 |
+
Aucun impact sur `store.py`, `retriever.py` (côté embedding), `chunker.py`, `app.py`.
|
| 69 |
+
|
| 70 |
+
## Modèle ONNX — préparation
|
| 71 |
+
|
| 72 |
+
Export one-shot du modèle bge-small au format ONNX :
|
| 73 |
+
- Via `optimum-cli export onnx` ou script Python
|
| 74 |
+
- Quantification int8 optionnelle (133 MB → ~35 MB)
|
| 75 |
+
- Fichiers résultants : `model.onnx` + `tokenizer.json`
|
| 76 |
+
- Stockés dans `models/bge-small-en-v1.5/`
|
| 77 |
+
- Embarqués dans l'image Docker (`COPY models/ models/`)
|
| 78 |
+
|
| 79 |
+
## Re-vectorisation
|
| 80 |
+
|
| 81 |
+
Obligatoire après migration — les embeddings 384d ne sont pas compatibles avec les 1024d actuels.
|
| 82 |
+
|
| 83 |
+
```bash
|
| 84 |
+
python cli.py vectorize # re-embed 350 stories dans ChromaDB via upsert
|
| 85 |
+
```
|
| 86 |
+
|
| 87 |
+
## Risques
|
| 88 |
+
|
| 89 |
+
- **Qualité retrieval :** -2 points MTEB (bge-large → small) + -3-5% NDCG (pas de reranker). Acceptable pour 350 docs. Mesurable via `python cli.py audit`.
|
| 90 |
+
- **Re-vectorisation :** One-shot, ~5 min pour 350 stories. ChromaDB upsert écrase les anciens vecteurs.
|
| 91 |
+
- **ONNX pooling/norm custom :** 10 lignes de code, bien documenté. Testable unitairement.
|