Skip to content

Commit 191909a

Browse files
committed
fix parallel worker init
1 parent 4c2f997 commit 191909a

File tree

1 file changed

+5
-1
lines changed

1 file changed

+5
-1
lines changed

fastembed/late_interaction/token_embeddings.py

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
from typing import Union, Iterable, Optional, List, Dict, Any
1+
from typing import Union, Iterable, Optional, List, Dict, Any, Type
22

33
import numpy as np
44

@@ -31,6 +31,10 @@ def list_supported_models(cls) -> List[Dict[str, Any]]:
3131
"""
3232
return supported_token_embeddings_models
3333

34+
@classmethod
35+
def _get_worker_class(cls) -> Type[TextEmbeddingWorker]:
36+
return TokensEmbeddingWorker
37+
3438
def _post_process_onnx_output(self, output: OnnxOutputContext) -> Iterable[np.ndarray]:
3539
# Size: (batch_size, sequence_length, hidden_size)
3640
embeddings = output.model_output

0 commit comments

Comments
 (0)