[ENH] Added batch_size as parameter to SentenceTransformerEmbeddingFu… #2759
Conversation
…nction class ## Description of changes Added the hyperparameter batch_size to the SentenceTransformerEmbeddingFunction class in order to have more control over memory requirements when it comes to deploying large sentence transformer. ## Test plan - [x] Tests pass locally with `pytest` for python. Local installation and usage also worked. ## Documentation Changes Added docstring for the new hyperparameter
Reviewer ChecklistPlease leverage this checklist to ensure your code review is thorough before approving Testing, Bugs, Errors, Logs, Documentation
System Compatibility
Quality
|
There was a problem hiding this comment.
LGTM with a minor nit.
@cowolff Thanks for this. Reading on the effects of batch_size (huggingface/transformers#2401), I feel this is a somewhat contentious topic, but having explicit is better DX (also Zen of Python)
| @@ -17,6 +17,7 @@ def __init__( | |||
| model_name: str = "all-MiniLM-L6-v2", | |||
| device: str = "cpu", | |||
| normalize_embeddings: bool = False, | |||
| batch_size: int = 32, | |||
There was a problem hiding this comment.
Can we make this Optional[int] = 32
There was a problem hiding this comment.
Also, technically, this is already supported by the kwargs (not that we pass them to the encode method), but perhaps making it explicit is arguably a better DX.
|
Our underlying impl has changed and so this PR is not landable as is. That being said - we'd still like to add this functionality and that is now tracked in this issue. |
…nction class
Description of changes
Added the hyperparameter batch_size to the SentenceTransformerEmbeddingFunction class in order to have more control over memory requirements when it comes to deploying large sentence transformers.
Test plan
pytestfor python. Local installation and usage also worked.Documentation Changes
Added docstring for the new hyperparameter