fix parallel worker init

This commit is contained in:
generall
2025-01-12 17:05:24 +01:00
parent a63cb00a91
commit d4c4d86e10

View File

@@ -1,4 +1,4 @@
from typing import Union, Iterable, Optional, List, Dict, Any
from typing import Union, Iterable, Optional, List, Dict, Any, Type
import numpy as np
@@ -31,6 +31,10 @@ class TokenEmbeddingsModel(OnnxTextEmbedding, LateInteractionTextEmbeddingBase):
"""
return supported_token_embeddings_models
@classmethod
def _get_worker_class(cls) -> Type[TextEmbeddingWorker]:
return TokensEmbeddingWorker
def _post_process_onnx_output(self, output: OnnxOutputContext) -> Iterable[np.ndarray]:
# Size: (batch_size, sequence_length, hidden_size)
embeddings = output.model_output