This commit is contained in:
Tom Aarsen 2025-11-18 12:19:55 +01:00
parent 835edda6fc
commit 23e7ffbe1c
2 changed files with 10 additions and 4 deletions

View file

@ -13,7 +13,9 @@ from sentence_transformers import SentenceTransformer
async def sentence_transformers_embed(
texts: list[str], model: SentenceTransformer
) -> np.ndarray:
async def inner_encode(texts: list[str], model: SentenceTransformer, embedding_dim: int = 1024):
async def inner_encode(
texts: list[str], model: SentenceTransformer, embedding_dim: int = 1024
):
return model.encode(
texts,
truncate_dim=embedding_dim,
@ -22,5 +24,9 @@ async def sentence_transformers_embed(
show_progress_bar=False,
)
embedding_func = EmbeddingFunc(embedding_dim=model.get_sentence_embedding_dimension(), func=inner_encode, max_token_size=model.get_max_seq_length())
embedding_func = EmbeddingFunc(
embedding_dim=model.get_sentence_embedding_dimension(),
func=inner_encode,
max_token_size=model.get_max_seq_length(),
)
return await embedding_func(texts, model=model)