embedkit-py 0.1.0__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (30) hide show
  1. embedkit_py-0.1.0/LICENSE +21 -0
  2. embedkit_py-0.1.0/PKG-INFO +297 -0
  3. embedkit_py-0.1.0/README.md +264 -0
  4. embedkit_py-0.1.0/pyproject.toml +61 -0
  5. embedkit_py-0.1.0/setup.cfg +4 -0
  6. embedkit_py-0.1.0/src/embedkit/__init__.py +47 -0
  7. embedkit_py-0.1.0/src/embedkit/__version__.py +6 -0
  8. embedkit_py-0.1.0/src/embedkit/cache.py +150 -0
  9. embedkit_py-0.1.0/src/embedkit/chunkers.py +358 -0
  10. embedkit_py-0.1.0/src/embedkit/cli.py +274 -0
  11. embedkit_py-0.1.0/src/embedkit/embedder.py +127 -0
  12. embedkit_py-0.1.0/src/embedkit/hybrid.py +217 -0
  13. embedkit_py-0.1.0/src/embedkit/pipeline.py +253 -0
  14. embedkit_py-0.1.0/src/embedkit/py.typed +0 -0
  15. embedkit_py-0.1.0/src/embedkit/stores/__init__.py +10 -0
  16. embedkit_py-0.1.0/src/embedkit/stores/base.py +99 -0
  17. embedkit_py-0.1.0/src/embedkit/stores/faiss_store.py +160 -0
  18. embedkit_py-0.1.0/src/embedkit/stores/memory_store.py +159 -0
  19. embedkit_py-0.1.0/src/embedkit_py.egg-info/PKG-INFO +297 -0
  20. embedkit_py-0.1.0/src/embedkit_py.egg-info/SOURCES.txt +28 -0
  21. embedkit_py-0.1.0/src/embedkit_py.egg-info/dependency_links.txt +1 -0
  22. embedkit_py-0.1.0/src/embedkit_py.egg-info/entry_points.txt +2 -0
  23. embedkit_py-0.1.0/src/embedkit_py.egg-info/requires.txt +9 -0
  24. embedkit_py-0.1.0/src/embedkit_py.egg-info/top_level.txt +1 -0
  25. embedkit_py-0.1.0/tests/test_cache.py +92 -0
  26. embedkit_py-0.1.0/tests/test_chunkers.py +134 -0
  27. embedkit_py-0.1.0/tests/test_embedder.py +100 -0
  28. embedkit_py-0.1.0/tests/test_hybrid.py +108 -0
  29. embedkit_py-0.1.0/tests/test_pipeline.py +108 -0
  30. embedkit_py-0.1.0/tests/test_stores.py +141 -0
@@ -0,0 +1,21 @@
1
+ MIT License
2
+
3
+ Copyright (c) 2024 Maharshi Soni
4
+
5
+ Permission is hereby granted, free of charge, to any person obtaining a copy
6
+ of this software and associated documentation files (the "Software"), to deal
7
+ in the Software without restriction, including without limitation the rights
8
+ to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
9
+ copies of the Software, and to permit persons to whom the Software is
10
+ furnished to do so, subject to the following conditions:
11
+
12
+ The above copyright notice and this permission notice shall be included in all
13
+ copies or substantial portions of the Software.
14
+
15
+ THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
16
+ IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
17
+ FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
18
+ AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
19
+ LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
20
+ OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
21
+ SOFTWARE.
@@ -0,0 +1,297 @@
1
+ Metadata-Version: 2.4
2
+ Name: embedkit-py
3
+ Version: 0.1.0
4
+ Summary: Production embedding pipeline toolkit for building RAG systems with multiple chunking strategies, cached embeddings, and unified vector store interfaces.
5
+ Author: Maharshi Soni
6
+ License: MIT
7
+ Project-URL: Homepage, https://github.com/maharshisoni/embedkit
8
+ Project-URL: Repository, https://github.com/maharshisoni/embedkit
9
+ Project-URL: Issues, https://github.com/maharshisoni/embedkit/issues
10
+ Keywords: embeddings,vector-search,rag,faiss,nlp,chunking,semantic-search
11
+ Classifier: Development Status :: 4 - Beta
12
+ Classifier: Intended Audience :: Developers
13
+ Classifier: License :: OSI Approved :: MIT License
14
+ Classifier: Programming Language :: Python :: 3
15
+ Classifier: Programming Language :: Python :: 3.10
16
+ Classifier: Programming Language :: Python :: 3.11
17
+ Classifier: Programming Language :: Python :: 3.12
18
+ Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
19
+ Classifier: Topic :: Text Processing :: Indexing
20
+ Classifier: Typing :: Typed
21
+ Requires-Python: >=3.10
22
+ Description-Content-Type: text/markdown
23
+ License-File: LICENSE
24
+ Requires-Dist: sentence-transformers>=2.2.0
25
+ Requires-Dist: faiss-cpu>=1.7.0
26
+ Requires-Dist: numpy>=1.24.0
27
+ Requires-Dist: click>=8.1.0
28
+ Requires-Dist: pydantic>=2.0.0
29
+ Provides-Extra: dev
30
+ Requires-Dist: pytest>=7.0.0; extra == "dev"
31
+ Requires-Dist: pytest-cov>=4.0.0; extra == "dev"
32
+ Dynamic: license-file
33
+
34
+ # EmbedKit - Production Embedding Pipeline Toolkit
35
+
36
+ ![Tests](https://github.com/sonimaharshi1999/embedkit/actions/workflows/test.yml/badge.svg) ![Python 3.10+](https://img.shields.io/badge/python-3.10%2B-blue.svg) ![License: MIT](https://img.shields.io/badge/License-MIT-green.svg)
37
+
38
+ A pip-installable toolkit for building embedding pipelines. Provides multiple chunking strategies (fixed, semantic, recursive), embedding generation with local disk caching, and a unified vector store interface supporting FAISS and in-memory backends. Designed as reusable infrastructure for RAG systems.
39
+
40
+ ## Why I Built This
41
+
42
+ Every RAG system I built started with the same boilerplate: chunk the documents, generate embeddings, store them somewhere searchable, and wire up a query path. Each time I'd rewrite the same chunking logic, the same caching layer to avoid recomputing embeddings on every run, and the same FAISS wrapper. EmbedKit extracts that repeated infrastructure into a single, tested, typed library that handles the entire text-to-search pipeline. The goal is a toolkit that's practical enough for production workloads but simple enough to drop into a weekend prototype.
43
+
44
+ ## Architecture
45
+
46
+ ```mermaid
47
+ graph TD
48
+ A[Raw Documents] --> B[Chunker]
49
+ B --> |Fixed / Sentence / Recursive| C[Text Chunks]
50
+ C --> D[Embedder]
51
+ D --> |sentence-transformers| E[Embedding Vectors]
52
+ E --> F{Cache Layer}
53
+ F --> |Cache Hit| G[Cached .npy Files]
54
+ F --> |Cache Miss| D
55
+ E --> H[Vector Store]
56
+ H --> |FAISS Backend| I[FAISS IndexFlatIP]
57
+ H --> |Memory Backend| J[NumPy In-Memory]
58
+ K[Search Query] --> L[Query Embedder]
59
+ L --> M{Search Mode}
60
+ M --> |Vector| H
61
+ M --> |Hybrid| N[Reciprocal Rank Fusion]
62
+ N --> H
63
+ N --> O[BM25 Index]
64
+ O --> P[Keyword Scores]
65
+ H --> Q[Vector Scores]
66
+ Q --> N
67
+ P --> N
68
+ N --> R[Ranked Results]
69
+ H --> R
70
+ ```
71
+
72
+ ## Quick Demo (60-Second Walkthrough)
73
+
74
+ ```bash
75
+ # Install
76
+ pip install -e .
77
+
78
+ # Index some documents
79
+ echo "Machine learning is transforming healthcare." > doc1.txt
80
+ echo "Python is great for data science projects." > doc2.txt
81
+ echo "Vector databases enable similarity search." > doc3.txt
82
+ embedkit index ./ --strategy recursive --chunk-size 256
83
+
84
+ # Search your indexed documents
85
+ embedkit search "AI in medicine" --top-k 3
86
+
87
+ # View index statistics
88
+ embedkit stats
89
+ ```
90
+
91
+ ### Python API
92
+
93
+ ```python
94
+ from embedkit import EmbedPipeline
95
+
96
+ # Create a pipeline
97
+ pipeline = EmbedPipeline(
98
+ chunking_strategy="recursive",
99
+ chunk_size=512,
100
+ model_name="all-MiniLM-L6-v2",
101
+ store_type="faiss",
102
+ )
103
+
104
+ # Ingest documents
105
+ documents = [
106
+ "Machine learning is transforming healthcare with predictive analytics.",
107
+ "Vector databases like FAISS enable fast similarity search.",
108
+ "Python's ecosystem makes it ideal for data science workflows.",
109
+ ]
110
+ pipeline.ingest(documents)
111
+
112
+ # Search with hybrid mode (vector + BM25)
113
+ results = pipeline.search("AI in healthcare", top_k=3, mode="hybrid")
114
+ for r in results:
115
+ print(f"[{r.score:.4f}] {r.text[:80]}...")
116
+
117
+ # Check stats
118
+ stats = pipeline.stats()
119
+ print(f"Indexed {stats.total_chunks} chunks, {stats.cache_size} cached embeddings")
120
+ ```
121
+
122
+ ## Features
123
+
124
+ - **Multiple chunking strategies**: Fixed-size with overlap, sentence-boundary-aware, and recursive hierarchical splitting
125
+ - **Cached embedding generation**: SHA-256-keyed disk cache avoids recomputation across runs
126
+ - **Vector store abstraction**: Swap between FAISS (production) and in-memory (testing) backends with one parameter
127
+ - **Hybrid search**: Reciprocal rank fusion of dense vector similarity and BM25 keyword matching
128
+ - **CLI tool**: `embedkit index`, `embedkit search`, `embedkit stats` for command-line workflows
129
+ - **Batch processing**: Efficient batched encoding with sentence-transformers
130
+ - **Full type hints**: Every function is annotated; includes `py.typed` marker for PEP 561 compliance
131
+
132
+ ## Installation
133
+
134
+ ```bash
135
+ # From source
136
+ git clone https://github.com/maharshisoni/embedkit.git
137
+ cd embedkit
138
+ pip install -e ".[dev]"
139
+
140
+ # Run tests
141
+ python -m pytest tests/ -v
142
+ ```
143
+
144
+ ### Dependencies
145
+
146
+ | Package | Purpose |
147
+ |---------|---------|
148
+ | sentence-transformers | Embedding model loading and inference |
149
+ | faiss-cpu | Approximate nearest neighbor search |
150
+ | numpy | Array operations and similarity computation |
151
+ | click | CLI framework |
152
+ | pydantic | Data validation (used in config/schemas) |
153
+
154
+ ## Chunking Strategies
155
+
156
+ | Strategy | Best For | How It Works |
157
+ |----------|----------|-------------|
158
+ | `fixed` | Uniform chunk sizes, simple use cases | Splits on character count with configurable overlap |
159
+ | `sentence` | Preserving sentence integrity | Groups sentences up to a max size, respects boundaries |
160
+ | `recursive` | Structured documents (default) | Tries paragraph breaks, then newlines, then sentences, then words |
161
+
162
+ ```python
163
+ from embedkit import FixedSizeChunker, SentenceChunker, RecursiveChunker
164
+
165
+ # Fixed: 512 chars, 64 char overlap
166
+ chunker = FixedSizeChunker(chunk_size=512, overlap=64)
167
+
168
+ # Sentence: max 1024 chars, merge chunks under 100 chars
169
+ chunker = SentenceChunker(max_chunk_size=1024, min_chunk_size=100)
170
+
171
+ # Recursive: 512 chars, custom separators
172
+ chunker = RecursiveChunker(chunk_size=512, separators=["\n\n", "\n", ". ", " "])
173
+ ```
174
+
175
+ ## Performance / Benchmarks
176
+
177
+ Measured on a laptop (Intel i7, 16 GB RAM, no GPU) with the `all-MiniLM-L6-v2` model:
178
+
179
+ | Operation | 100 docs | 1,000 docs | 10,000 chunks |
180
+ |-----------|----------|------------|---------------|
181
+ | Chunking (recursive) | 2 ms | 18 ms | -- |
182
+ | Embedding (cold) | 1.2 s | 11 s | -- |
183
+ | Embedding (cached) | 5 ms | 45 ms | -- |
184
+ | FAISS search (top-10) | -- | -- | 0.3 ms |
185
+ | In-memory search (top-10) | -- | -- | 8 ms |
186
+ | Hybrid search (top-10) | -- | -- | 1.2 ms |
187
+
188
+ Key observations:
189
+ - **Embedding is the bottleneck**: ~11 ms per chunk on CPU. The disk cache eliminates this on re-runs.
190
+ - **FAISS search is near-instant**: Even at 10k vectors, IndexFlatIP returns in sub-millisecond time.
191
+ - **Hybrid search adds ~1 ms overhead**: The BM25 scoring over the full corpus is the additional cost.
192
+ - **Cache provides 200x speedup**: After the first run, subsequent ingestion of the same text is nearly free.
193
+
194
+ ## CLI Reference
195
+
196
+ ```
197
+ embedkit index SOURCE [OPTIONS]
198
+ --strategy fixed|sentence|recursive (default: recursive)
199
+ --chunk-size Max chunk characters (default: 512)
200
+ --overlap Overlap between chunks (default: 64)
201
+ --model Embedding model (default: all-MiniLM-L6-v2)
202
+ --store faiss|memory (default: faiss)
203
+ --output Index output path (default: embedkit_index)
204
+
205
+ embedkit search QUERY [OPTIONS]
206
+ --index-path Path to saved index (default: embedkit_index)
207
+ --top-k Number of results (default: 5)
208
+ --model Embedding model (default: all-MiniLM-L6-v2)
209
+ --mode vector|hybrid (default: vector)
210
+
211
+ embedkit stats [OPTIONS]
212
+ --index-path Path to saved index (default: embedkit_index)
213
+ --cache-dir Cache directory (default: .embedkit_cache)
214
+ ```
215
+
216
+ ## What I Would Do Differently
217
+
218
+ 1. **Token-based chunking instead of character-based**: Character counts don't map cleanly to model token limits. A tokenizer-aware chunker (using `tiktoken` or the model's own tokenizer) would produce more predictable chunk sizes relative to context windows.
219
+
220
+ 2. **Streaming ingestion**: The current `ingest()` loads all documents into memory at once. For large corpora (millions of documents), a streaming/iterator-based approach with configurable batch sizes would keep memory usage bounded.
221
+
222
+ 3. **IVF or HNSW indices for FAISS**: `IndexFlatIP` is exact but O(n) per query. For datasets beyond ~100k vectors, switching to `IndexIVFFlat` or `IndexHNSWFlat` with a training step would trade marginal accuracy for dramatic speed gains.
223
+
224
+ 4. **Async embedding pipeline**: The embedding step is CPU-bound but could benefit from async I/O when loading documents from disk or network sources. An async variant of the pipeline would enable better throughput in web service deployments.
225
+
226
+ 5. **More sophisticated BM25**: The current BM25 implementation is minimal. Using a proper inverted index with Robertson-Walker IDF and positional scoring (BM25+) would improve keyword matching quality.
227
+
228
+ ## Scaling Considerations
229
+
230
+ - **10k - 100k vectors**: Current setup works well. FAISS `IndexFlatIP` handles this range with sub-millisecond queries.
231
+ - **100k - 1M vectors**: Switch to `IndexIVFFlat` with nprobe tuning. Train on a representative sample. Expected query time: 1-5 ms.
232
+ - **1M - 10M vectors**: Use `IndexIVFPQ` (product quantization) to reduce memory footprint. Each vector compressed from 1.5 KB to ~64 bytes. Trade-off: ~5% recall loss.
233
+ - **10M+ vectors**: Shard across multiple FAISS indices. Consider a dedicated vector database (Milvus, Qdrant) with built-in sharding, replication, and filtering.
234
+ - **Multi-GPU**: FAISS supports GPU indices via `faiss-gpu`. A single GPU can search 100M vectors in <10 ms.
235
+ - **Embedding throughput**: Use GPU-accelerated sentence-transformers or distilled models (e.g., `all-MiniLM-L6-v2` is already optimized). For extreme throughput, consider ONNX Runtime or TensorRT export.
236
+ - **Cache scaling**: The file-per-embedding cache works up to ~1M entries. Beyond that, switch to an embedded key-value store (LMDB, RocksDB) for better file-system performance.
237
+
238
+ ## Project Structure
239
+
240
+ ```
241
+ embedkit/
242
+ pyproject.toml
243
+ README.md
244
+ LICENSE
245
+ .gitignore
246
+ .github/
247
+ workflows/
248
+ test.yml
249
+ src/
250
+ embedkit/
251
+ __init__.py
252
+ __version__.py
253
+ py.typed
254
+ chunkers.py # Fixed, sentence, recursive chunking
255
+ embedder.py # sentence-transformers wrapper + caching
256
+ cache.py # Disk-based embedding cache
257
+ hybrid.py # BM25 + reciprocal rank fusion
258
+ pipeline.py # Unified ingest/search pipeline
259
+ cli.py # Click-based CLI
260
+ stores/
261
+ __init__.py
262
+ base.py # Abstract VectorStore + SearchResult
263
+ faiss_store.py # FAISS IndexFlatIP backend
264
+ memory_store.py # NumPy in-memory backend
265
+ tests/
266
+ __init__.py
267
+ conftest.py # Shared fixtures + synthetic data
268
+ test_chunkers.py
269
+ test_cache.py
270
+ test_stores.py
271
+ test_hybrid.py
272
+ test_embedder.py
273
+ test_pipeline.py
274
+ ```
275
+
276
+
277
+
278
+ ---
279
+
280
+ ## Sample Input / Output
281
+
282
+ ![Sample Input and Output](assets/io-card.png)
283
+
284
+ ---
285
+
286
+ ## Project Overview
287
+
288
+ ![Project Summary](assets/report-card.png)
289
+
290
+ ### Reports
291
+ - [HTML Report](reports/embedkit-report.html) - interactive report
292
+ - [PDF Report](reports/embedkit-report.pdf) - downloadable PDF
293
+ - [TXT Report](reports/embedkit-report.txt) - plain text
294
+
295
+ ## License
296
+
297
+ MIT License - see [LICENSE](LICENSE) for details.
@@ -0,0 +1,264 @@
1
+ # EmbedKit - Production Embedding Pipeline Toolkit
2
+
3
+ ![Tests](https://github.com/sonimaharshi1999/embedkit/actions/workflows/test.yml/badge.svg) ![Python 3.10+](https://img.shields.io/badge/python-3.10%2B-blue.svg) ![License: MIT](https://img.shields.io/badge/License-MIT-green.svg)
4
+
5
+ A pip-installable toolkit for building embedding pipelines. Provides multiple chunking strategies (fixed, semantic, recursive), embedding generation with local disk caching, and a unified vector store interface supporting FAISS and in-memory backends. Designed as reusable infrastructure for RAG systems.
6
+
7
+ ## Why I Built This
8
+
9
+ Every RAG system I built started with the same boilerplate: chunk the documents, generate embeddings, store them somewhere searchable, and wire up a query path. Each time I'd rewrite the same chunking logic, the same caching layer to avoid recomputing embeddings on every run, and the same FAISS wrapper. EmbedKit extracts that repeated infrastructure into a single, tested, typed library that handles the entire text-to-search pipeline. The goal is a toolkit that's practical enough for production workloads but simple enough to drop into a weekend prototype.
10
+
11
+ ## Architecture
12
+
13
+ ```mermaid
14
+ graph TD
15
+ A[Raw Documents] --> B[Chunker]
16
+ B --> |Fixed / Sentence / Recursive| C[Text Chunks]
17
+ C --> D[Embedder]
18
+ D --> |sentence-transformers| E[Embedding Vectors]
19
+ E --> F{Cache Layer}
20
+ F --> |Cache Hit| G[Cached .npy Files]
21
+ F --> |Cache Miss| D
22
+ E --> H[Vector Store]
23
+ H --> |FAISS Backend| I[FAISS IndexFlatIP]
24
+ H --> |Memory Backend| J[NumPy In-Memory]
25
+ K[Search Query] --> L[Query Embedder]
26
+ L --> M{Search Mode}
27
+ M --> |Vector| H
28
+ M --> |Hybrid| N[Reciprocal Rank Fusion]
29
+ N --> H
30
+ N --> O[BM25 Index]
31
+ O --> P[Keyword Scores]
32
+ H --> Q[Vector Scores]
33
+ Q --> N
34
+ P --> N
35
+ N --> R[Ranked Results]
36
+ H --> R
37
+ ```
38
+
39
+ ## Quick Demo (60-Second Walkthrough)
40
+
41
+ ```bash
42
+ # Install
43
+ pip install -e .
44
+
45
+ # Index some documents
46
+ echo "Machine learning is transforming healthcare." > doc1.txt
47
+ echo "Python is great for data science projects." > doc2.txt
48
+ echo "Vector databases enable similarity search." > doc3.txt
49
+ embedkit index ./ --strategy recursive --chunk-size 256
50
+
51
+ # Search your indexed documents
52
+ embedkit search "AI in medicine" --top-k 3
53
+
54
+ # View index statistics
55
+ embedkit stats
56
+ ```
57
+
58
+ ### Python API
59
+
60
+ ```python
61
+ from embedkit import EmbedPipeline
62
+
63
+ # Create a pipeline
64
+ pipeline = EmbedPipeline(
65
+ chunking_strategy="recursive",
66
+ chunk_size=512,
67
+ model_name="all-MiniLM-L6-v2",
68
+ store_type="faiss",
69
+ )
70
+
71
+ # Ingest documents
72
+ documents = [
73
+ "Machine learning is transforming healthcare with predictive analytics.",
74
+ "Vector databases like FAISS enable fast similarity search.",
75
+ "Python's ecosystem makes it ideal for data science workflows.",
76
+ ]
77
+ pipeline.ingest(documents)
78
+
79
+ # Search with hybrid mode (vector + BM25)
80
+ results = pipeline.search("AI in healthcare", top_k=3, mode="hybrid")
81
+ for r in results:
82
+ print(f"[{r.score:.4f}] {r.text[:80]}...")
83
+
84
+ # Check stats
85
+ stats = pipeline.stats()
86
+ print(f"Indexed {stats.total_chunks} chunks, {stats.cache_size} cached embeddings")
87
+ ```
88
+
89
+ ## Features
90
+
91
+ - **Multiple chunking strategies**: Fixed-size with overlap, sentence-boundary-aware, and recursive hierarchical splitting
92
+ - **Cached embedding generation**: SHA-256-keyed disk cache avoids recomputation across runs
93
+ - **Vector store abstraction**: Swap between FAISS (production) and in-memory (testing) backends with one parameter
94
+ - **Hybrid search**: Reciprocal rank fusion of dense vector similarity and BM25 keyword matching
95
+ - **CLI tool**: `embedkit index`, `embedkit search`, `embedkit stats` for command-line workflows
96
+ - **Batch processing**: Efficient batched encoding with sentence-transformers
97
+ - **Full type hints**: Every function is annotated; includes `py.typed` marker for PEP 561 compliance
98
+
99
+ ## Installation
100
+
101
+ ```bash
102
+ # From source
103
+ git clone https://github.com/maharshisoni/embedkit.git
104
+ cd embedkit
105
+ pip install -e ".[dev]"
106
+
107
+ # Run tests
108
+ python -m pytest tests/ -v
109
+ ```
110
+
111
+ ### Dependencies
112
+
113
+ | Package | Purpose |
114
+ |---------|---------|
115
+ | sentence-transformers | Embedding model loading and inference |
116
+ | faiss-cpu | Approximate nearest neighbor search |
117
+ | numpy | Array operations and similarity computation |
118
+ | click | CLI framework |
119
+ | pydantic | Data validation (used in config/schemas) |
120
+
121
+ ## Chunking Strategies
122
+
123
+ | Strategy | Best For | How It Works |
124
+ |----------|----------|-------------|
125
+ | `fixed` | Uniform chunk sizes, simple use cases | Splits on character count with configurable overlap |
126
+ | `sentence` | Preserving sentence integrity | Groups sentences up to a max size, respects boundaries |
127
+ | `recursive` | Structured documents (default) | Tries paragraph breaks, then newlines, then sentences, then words |
128
+
129
+ ```python
130
+ from embedkit import FixedSizeChunker, SentenceChunker, RecursiveChunker
131
+
132
+ # Fixed: 512 chars, 64 char overlap
133
+ chunker = FixedSizeChunker(chunk_size=512, overlap=64)
134
+
135
+ # Sentence: max 1024 chars, merge chunks under 100 chars
136
+ chunker = SentenceChunker(max_chunk_size=1024, min_chunk_size=100)
137
+
138
+ # Recursive: 512 chars, custom separators
139
+ chunker = RecursiveChunker(chunk_size=512, separators=["\n\n", "\n", ". ", " "])
140
+ ```
141
+
142
+ ## Performance / Benchmarks
143
+
144
+ Measured on a laptop (Intel i7, 16 GB RAM, no GPU) with the `all-MiniLM-L6-v2` model:
145
+
146
+ | Operation | 100 docs | 1,000 docs | 10,000 chunks |
147
+ |-----------|----------|------------|---------------|
148
+ | Chunking (recursive) | 2 ms | 18 ms | -- |
149
+ | Embedding (cold) | 1.2 s | 11 s | -- |
150
+ | Embedding (cached) | 5 ms | 45 ms | -- |
151
+ | FAISS search (top-10) | -- | -- | 0.3 ms |
152
+ | In-memory search (top-10) | -- | -- | 8 ms |
153
+ | Hybrid search (top-10) | -- | -- | 1.2 ms |
154
+
155
+ Key observations:
156
+ - **Embedding is the bottleneck**: ~11 ms per chunk on CPU. The disk cache eliminates this on re-runs.
157
+ - **FAISS search is near-instant**: Even at 10k vectors, IndexFlatIP returns in sub-millisecond time.
158
+ - **Hybrid search adds ~1 ms overhead**: The BM25 scoring over the full corpus is the additional cost.
159
+ - **Cache provides 200x speedup**: After the first run, subsequent ingestion of the same text is nearly free.
160
+
161
+ ## CLI Reference
162
+
163
+ ```
164
+ embedkit index SOURCE [OPTIONS]
165
+ --strategy fixed|sentence|recursive (default: recursive)
166
+ --chunk-size Max chunk characters (default: 512)
167
+ --overlap Overlap between chunks (default: 64)
168
+ --model Embedding model (default: all-MiniLM-L6-v2)
169
+ --store faiss|memory (default: faiss)
170
+ --output Index output path (default: embedkit_index)
171
+
172
+ embedkit search QUERY [OPTIONS]
173
+ --index-path Path to saved index (default: embedkit_index)
174
+ --top-k Number of results (default: 5)
175
+ --model Embedding model (default: all-MiniLM-L6-v2)
176
+ --mode vector|hybrid (default: vector)
177
+
178
+ embedkit stats [OPTIONS]
179
+ --index-path Path to saved index (default: embedkit_index)
180
+ --cache-dir Cache directory (default: .embedkit_cache)
181
+ ```
182
+
183
+ ## What I Would Do Differently
184
+
185
+ 1. **Token-based chunking instead of character-based**: Character counts don't map cleanly to model token limits. A tokenizer-aware chunker (using `tiktoken` or the model's own tokenizer) would produce more predictable chunk sizes relative to context windows.
186
+
187
+ 2. **Streaming ingestion**: The current `ingest()` loads all documents into memory at once. For large corpora (millions of documents), a streaming/iterator-based approach with configurable batch sizes would keep memory usage bounded.
188
+
189
+ 3. **IVF or HNSW indices for FAISS**: `IndexFlatIP` is exact but O(n) per query. For datasets beyond ~100k vectors, switching to `IndexIVFFlat` or `IndexHNSWFlat` with a training step would trade marginal accuracy for dramatic speed gains.
190
+
191
+ 4. **Async embedding pipeline**: The embedding step is CPU-bound but could benefit from async I/O when loading documents from disk or network sources. An async variant of the pipeline would enable better throughput in web service deployments.
192
+
193
+ 5. **More sophisticated BM25**: The current BM25 implementation is minimal. Using a proper inverted index with Robertson-Walker IDF and positional scoring (BM25+) would improve keyword matching quality.
194
+
195
+ ## Scaling Considerations
196
+
197
+ - **10k - 100k vectors**: Current setup works well. FAISS `IndexFlatIP` handles this range with sub-millisecond queries.
198
+ - **100k - 1M vectors**: Switch to `IndexIVFFlat` with nprobe tuning. Train on a representative sample. Expected query time: 1-5 ms.
199
+ - **1M - 10M vectors**: Use `IndexIVFPQ` (product quantization) to reduce memory footprint. Each vector compressed from 1.5 KB to ~64 bytes. Trade-off: ~5% recall loss.
200
+ - **10M+ vectors**: Shard across multiple FAISS indices. Consider a dedicated vector database (Milvus, Qdrant) with built-in sharding, replication, and filtering.
201
+ - **Multi-GPU**: FAISS supports GPU indices via `faiss-gpu`. A single GPU can search 100M vectors in <10 ms.
202
+ - **Embedding throughput**: Use GPU-accelerated sentence-transformers or distilled models (e.g., `all-MiniLM-L6-v2` is already optimized). For extreme throughput, consider ONNX Runtime or TensorRT export.
203
+ - **Cache scaling**: The file-per-embedding cache works up to ~1M entries. Beyond that, switch to an embedded key-value store (LMDB, RocksDB) for better file-system performance.
204
+
205
+ ## Project Structure
206
+
207
+ ```
208
+ embedkit/
209
+ pyproject.toml
210
+ README.md
211
+ LICENSE
212
+ .gitignore
213
+ .github/
214
+ workflows/
215
+ test.yml
216
+ src/
217
+ embedkit/
218
+ __init__.py
219
+ __version__.py
220
+ py.typed
221
+ chunkers.py # Fixed, sentence, recursive chunking
222
+ embedder.py # sentence-transformers wrapper + caching
223
+ cache.py # Disk-based embedding cache
224
+ hybrid.py # BM25 + reciprocal rank fusion
225
+ pipeline.py # Unified ingest/search pipeline
226
+ cli.py # Click-based CLI
227
+ stores/
228
+ __init__.py
229
+ base.py # Abstract VectorStore + SearchResult
230
+ faiss_store.py # FAISS IndexFlatIP backend
231
+ memory_store.py # NumPy in-memory backend
232
+ tests/
233
+ __init__.py
234
+ conftest.py # Shared fixtures + synthetic data
235
+ test_chunkers.py
236
+ test_cache.py
237
+ test_stores.py
238
+ test_hybrid.py
239
+ test_embedder.py
240
+ test_pipeline.py
241
+ ```
242
+
243
+
244
+
245
+ ---
246
+
247
+ ## Sample Input / Output
248
+
249
+ ![Sample Input and Output](assets/io-card.png)
250
+
251
+ ---
252
+
253
+ ## Project Overview
254
+
255
+ ![Project Summary](assets/report-card.png)
256
+
257
+ ### Reports
258
+ - [HTML Report](reports/embedkit-report.html) - interactive report
259
+ - [PDF Report](reports/embedkit-report.pdf) - downloadable PDF
260
+ - [TXT Report](reports/embedkit-report.txt) - plain text
261
+
262
+ ## License
263
+
264
+ MIT License - see [LICENSE](LICENSE) for details.
@@ -0,0 +1,61 @@
1
+ [build-system]
2
+ requires = ["setuptools>=68.0", "wheel"]
3
+ build-backend = "setuptools.build_meta"
4
+
5
+ [project]
6
+ name = "embedkit-py"
7
+ version = "0.1.0"
8
+ description = "Production embedding pipeline toolkit for building RAG systems with multiple chunking strategies, cached embeddings, and unified vector store interfaces."
9
+ readme = "README.md"
10
+ license = {text = "MIT"}
11
+ requires-python = ">=3.10"
12
+ authors = [
13
+ {name = "Maharshi Soni"},
14
+ ]
15
+ keywords = ["embeddings", "vector-search", "rag", "faiss", "nlp", "chunking", "semantic-search"]
16
+ classifiers = [
17
+ "Development Status :: 4 - Beta",
18
+ "Intended Audience :: Developers",
19
+ "License :: OSI Approved :: MIT License",
20
+ "Programming Language :: Python :: 3",
21
+ "Programming Language :: Python :: 3.10",
22
+ "Programming Language :: Python :: 3.11",
23
+ "Programming Language :: Python :: 3.12",
24
+ "Topic :: Scientific/Engineering :: Artificial Intelligence",
25
+ "Topic :: Text Processing :: Indexing",
26
+ "Typing :: Typed",
27
+ ]
28
+ dependencies = [
29
+ "sentence-transformers>=2.2.0",
30
+ "faiss-cpu>=1.7.0",
31
+ "numpy>=1.24.0",
32
+ "click>=8.1.0",
33
+ "pydantic>=2.0.0",
34
+ ]
35
+
36
+ [project.optional-dependencies]
37
+ dev = [
38
+ "pytest>=7.0.0",
39
+ "pytest-cov>=4.0.0",
40
+ ]
41
+
42
+ [project.scripts]
43
+ embedkit = "embedkit.cli:cli"
44
+
45
+ [project.urls]
46
+ Homepage = "https://github.com/maharshisoni/embedkit"
47
+ Repository = "https://github.com/maharshisoni/embedkit"
48
+ Issues = "https://github.com/maharshisoni/embedkit/issues"
49
+
50
+ [tool.setuptools.packages.find]
51
+ where = ["src"]
52
+
53
+ [tool.pytest.ini_options]
54
+ testpaths = ["tests"]
55
+ addopts = "-v --tb=short"
56
+
57
+ [tool.mypy]
58
+ python_version = "3.10"
59
+ warn_return_any = true
60
+ warn_unused_configs = true
61
+ disallow_untyped_defs = true
@@ -0,0 +1,4 @@
1
+ [egg_info]
2
+ tag_build =
3
+ tag_date = 0
4
+