Skip to content

Commit 8fe754b

Browse files
authored
Add cache option to vectorizer from dict (#325)
1 parent 184f521 commit 8fe754b

File tree

1 file changed

+22
-8
lines changed

1 file changed

+22
-8
lines changed

redisvl/utils/vectorize/__init__.py

+22-8
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,7 @@
1+
import os
2+
from typing import Optional
3+
4+
from redisvl.extensions.cache.embeddings import EmbeddingsCache
15
from redisvl.utils.vectorize.base import BaseVectorizer, Vectorizers
26
from redisvl.utils.vectorize.text.azureopenai import AzureOpenAITextVectorizer
37
from redisvl.utils.vectorize.text.bedrock import BedrockTextVectorizer
@@ -23,22 +27,32 @@
2327
]
2428

2529

26-
def vectorizer_from_dict(vectorizer: dict) -> BaseVectorizer:
30+
def vectorizer_from_dict(
31+
vectorizer: dict,
32+
cache: dict = {},
33+
cache_folder=os.getenv("SENTENCE_TRANSFORMERS_HOME"),
34+
) -> BaseVectorizer:
2735
vectorizer_type = Vectorizers(vectorizer["type"])
2836
model = vectorizer["model"]
37+
38+
args = {"model": model}
39+
if cache:
40+
emb_cache = EmbeddingsCache(**cache)
41+
args["cache"] = emb_cache
42+
2943
if vectorizer_type == Vectorizers.cohere:
30-
return CohereTextVectorizer(model=model)
44+
return CohereTextVectorizer(**args)
3145
elif vectorizer_type == Vectorizers.openai:
32-
return OpenAITextVectorizer(model=model)
46+
return OpenAITextVectorizer(**args)
3347
elif vectorizer_type == Vectorizers.azure_openai:
34-
return AzureOpenAITextVectorizer(model=model)
48+
return AzureOpenAITextVectorizer(**args)
3549
elif vectorizer_type == Vectorizers.hf:
36-
return HFTextVectorizer(model=model)
50+
return HFTextVectorizer(**args)
3751
elif vectorizer_type == Vectorizers.mistral:
38-
return MistralAITextVectorizer(model=model)
52+
return MistralAITextVectorizer(**args)
3953
elif vectorizer_type == Vectorizers.vertexai:
40-
return VertexAITextVectorizer(model=model)
54+
return VertexAITextVectorizer(**args)
4155
elif vectorizer_type == Vectorizers.voyageai:
42-
return VoyageAITextVectorizer(model=model)
56+
return VoyageAITextVectorizer(**args)
4357
else:
4458
raise ValueError(f"Unsupported vectorizer type: {vectorizer_type}")

0 commit comments

Comments
 (0)