diff --git a/semantic_router/encoders/huggingface.py b/semantic_router/encoders/huggingface.py index 7ca7580d..1a7a1e2f 100644 --- a/semantic_router/encoders/huggingface.py +++ b/semantic_router/encoders/huggingface.py @@ -211,18 +211,21 @@ def __call__(self, docs: List[str]) -> List[List[float]]: Raises: ValueError: If no embeddings are returned for a document. """ + + batch_size=32 embeddings = [] - for d in docs: + for i in range(0, len(docs), batch_size): + batch = docs[i:i + batch_size] try: - output = self.query({"inputs": d, "parameters": {}}) - if not output or len(output) == 0: + outputs = self.query({"inputs": batch, "parameters": {}}) + if not outputs or len(outputs) == 0: raise ValueError("No embeddings returned from the query.") - embeddings.append(output) - + embeddings.extend(outputs) except Exception as e: raise ValueError( - f"No embeddings returned for document. Error: {e}" + f"No embeddings returned for batch. Error: {e}" ) from e + return embeddings def query(self, payload, max_retries=3, retry_interval=5): @@ -261,6 +264,7 @@ def query(self, payload, max_retries=3, retry_interval=5): continue else: response.raise_for_status() + break except requests.exceptions.RequestException: if attempt < max_retries - 1: