remdms Claude Opus 4.6 commited on
Commit
bdb27e7
·
1 Parent(s): 17a0c29

docs: design spec for ONNX lightweight migration

Browse files

Replace 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.