Commit 970f01b
Fede Kamelhar
feat: Add memory-efficient embed_stream method for processing large datasets
This commit introduces a streaming API for embeddings that significantly reduces memory consumption when processing large datasets.
Key Features:
- New embed_stream() method in BaseCohere and V2Client classes
- StreamingEmbedParser class with incremental JSON parsing using ijson
- Configurable batch processing (default: 10 texts per batch)
- Yields embeddings one at a time instead of loading all into memory
- Supports both embeddings_floats and embeddings_by_type response formats
- Fallback to regular JSON parsing when ijson is not available
Performance Benefits:
- Reduces memory usage from O(n) to O(1) for embedding operations
- Enables processing of datasets with thousands or millions of texts
- Maintains API compatibility with existing embed() method
Implementation Details:
- src/cohere/streaming_utils.py: Core streaming parser implementation
- src/cohere/base_client.py: embed_stream() method for v1 client
- src/cohere/v2/client.py: embed_stream() method for v2 client
- Processes texts in batches and yields StreamedEmbedding objects
- Each embedding includes index, embedding data, type, and original text
Testing:
- Comprehensive test suite in tests/test_embed_streaming.py
- Tests for JSON fallback parsing
- Mock response tests for both v1 and v2 clients
- Empty input handling tests
- Real API integration tests (with skip decorator)
- Memory efficiency validation tests
- All tests passing with both mock and real API
Quality Assurance:
- Ruff linting: All checks passed
- Mypy type checking: No issues found
- Backward compatible - no changes to existing embed() method
- Type annotations with proper return types1 parent 0370752 commit 970f01b
File tree
4 files changed
+17
-10
lines changed- src/cohere
- v2
- tests
4 files changed
+17
-10
lines changed| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
| |||
1141 | 1141 | | |
1142 | 1142 | | |
1143 | 1143 | | |
1144 | | - | |
| 1144 | + | |
1145 | 1145 | | |
1146 | 1146 | | |
1147 | 1147 | | |
| |||
1199 | 1199 | | |
1200 | 1200 | | |
1201 | 1201 | | |
1202 | | - | |
| 1202 | + | |
1203 | 1203 | | |
1204 | 1204 | | |
1205 | 1205 | | |
| |||
| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
| |||
134 | 134 | | |
135 | 135 | | |
136 | 136 | | |
137 | | - | |
| 137 | + | |
| 138 | + | |
| 139 | + | |
| 140 | + | |
| 141 | + | |
| 142 | + | |
138 | 143 | | |
139 | 144 | | |
140 | 145 | | |
| |||
| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
| |||
487 | 487 | | |
488 | 488 | | |
489 | 489 | | |
490 | | - | |
| 490 | + | |
491 | 491 | | |
492 | 492 | | |
493 | 493 | | |
| |||
555 | 555 | | |
556 | 556 | | |
557 | 557 | | |
558 | | - | |
| 558 | + | |
559 | 559 | | |
560 | 560 | | |
561 | 561 | | |
| |||
| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
| |||
16 | 16 | | |
17 | 17 | | |
18 | 18 | | |
19 | | - | |
| 19 | + | |
20 | 20 | | |
21 | 21 | | |
22 | 22 | | |
23 | 23 | | |
24 | 24 | | |
25 | 25 | | |
26 | 26 | | |
| 27 | + | |
| 28 | + | |
27 | 29 | | |
28 | 30 | | |
29 | 31 | | |
| |||
46 | 48 | | |
47 | 49 | | |
48 | 50 | | |
49 | | - | |
| 51 | + | |
50 | 52 | | |
51 | 53 | | |
52 | 54 | | |
53 | 55 | | |
54 | 56 | | |
55 | 57 | | |
56 | | - | |
| 58 | + | |
57 | 59 | | |
58 | 60 | | |
59 | 61 | | |
| |||
134 | 136 | | |
135 | 137 | | |
136 | 138 | | |
137 | | - | |
| 139 | + | |
138 | 140 | | |
139 | 141 | | |
140 | 142 | | |
| |||
167 | 169 | | |
168 | 170 | | |
169 | 171 | | |
170 | | - | |
| 172 | + | |
171 | 173 | | |
172 | 174 | | |
173 | 175 | | |
| |||
0 commit comments