You signed in with another tab or window. Reload to refresh your session.You signed out in another tab or window. Reload to refresh your session.You switched accounts on another tab or window. Reload to refresh your session.Dismiss alert
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 types
0 commit comments