Add batch size parameter to CLI for Sentence Transformers #24
Conversation
- Add --batch-size option to CLI with default value of 32 - Pass batch_size parameter through to compute_text_projection() - Update _projection_for_texts() to use SentenceTransformer's built-in batch_size parameter - Include batch_size in cache key for proper caching - Add documentation for the new parameter This allows users to control memory usage and performance when processing embeddings. Larger batch sizes use more memory but may be faster, while smaller batch sizes use less memory at the cost of potentially slower processing.
| @click.option( | ||
| "--batch-size", | ||
| type=int, | ||
| default=32, |
There was a problem hiding this comment.
Could the default from ST ever change? Then it might make sense to set None as the default here and pass that through.
There was a problem hiding this comment.
yeah, maybe None is safer here. cc @tomaarsen do you plan to change ST default batch size ever?
There was a problem hiding this comment.
No plans, Sentence Transformers is too mature to allow such backwards incompatible changes. I imagine it'll most likely never happen.
- Tom Aarsen
There was a problem hiding this comment.
@domoritz are you okay to leave the bs default as 32 for now in this case?
There was a problem hiding this comment.
Sure. That's okay if ST is stable.
There was a problem hiding this comment.
I think we should also add the batch_size parameter to compute_image_projection (which currently has a default of 16), so maybe it's still a good idea to default to None and use different defaults depending on the modality (and in the future we might be able to even auto adjust it by available VRAM)
There was a problem hiding this comment.
Agree that makes sense. Maybe we could have a default of None in the CLI, then pass 32 for text and 16 for images when None? It might also be worth adding a short log message when batch_size is None to let users know they can adjust it if they run out of memory/VRAM or have a powerful GPU and want to speed things up?
There was a problem hiding this comment.
Yes, both sounds great!
There was a problem hiding this comment.
Have implemented
- bs for images
- defaults bs for images/text (when None is passed to CLI, i.e. default behaviour)
- added more info to logging for users
- Add batch_size parameter to image processing functions - Use None as CLI default with modality-specific defaults (32 for text, 16 for images) - Add educational logging when using defaults - Include batch_size in image cache keys
Very nice package! For large datasets / when running on powerful GPUs for the embeddings, it would be nice to be able to either raise or lower the batch size for the embedding step. Currently, it's using the ST default of 32. I've had embedding-atlas crash when I run larger embedding models on my local machine, and for these, it would be nice to run with a batch size of 8 or so.
This PR adds:
This change shouldn't have an impact on other parts of the frontend, etc, but I mostly looked at the Python part of the code base, so I might have missed something.