[Improvement] add vector_dimension configuration in embedder config (#1192)

Co-authored-by: Deven Patel <deven298@yahoo.com>
This commit is contained in:
Deven Patel
2024-01-19 10:31:41 +05:30
committed by GitHub
parent e572b5a3dc
commit 59600e2a5b
9 changed files with 20 additions and 8 deletions

View File

@@ -6,7 +6,11 @@ from embedchain.helpers.json_serializable import register_deserializable
@register_deserializable
class BaseEmbedderConfig:
def __init__(
self, model: Optional[str] = None, deployment_name: Optional[str] = None, api_key: Optional[str] = None
self,
model: Optional[str] = None,
deployment_name: Optional[str] = None,
vector_dimension: Optional[int] = None,
api_key: Optional[str] = None,
):
"""
Initialize a new instance of an embedder config class.
@@ -18,4 +22,5 @@ class BaseEmbedderConfig:
"""
self.model = model
self.deployment_name = deployment_name
self.vector_dimension = vector_dimension
self.api_key = api_key

View File

@@ -27,5 +27,5 @@ class GoogleAIEmbedder(BaseEmbedder):
embedding_fn = GoogleAIEmbeddingFunction(config=config)
self.set_embedding_fn(embedding_fn=embedding_fn)
vector_dimension = VectorDimensions.GOOGLE_AI.value
vector_dimension = self.config.vector_dimension or VectorDimensions.GOOGLE_AI.value
self.set_vector_dimension(vector_dimension=vector_dimension)

View File

@@ -16,5 +16,5 @@ class GPT4AllEmbedder(BaseEmbedder):
embedding_fn = BaseEmbedder._langchain_default_concept(embeddings)
self.set_embedding_fn(embedding_fn=embedding_fn)
vector_dimension = VectorDimensions.GPT4ALL.value
vector_dimension = self.config.vector_dimension or VectorDimensions.GPT4ALL.value
self.set_vector_dimension(vector_dimension=vector_dimension)

View File

@@ -15,5 +15,5 @@ class HuggingFaceEmbedder(BaseEmbedder):
embedding_fn = BaseEmbedder._langchain_default_concept(embeddings)
self.set_embedding_fn(embedding_fn=embedding_fn)
vector_dimension = VectorDimensions.HUGGING_FACE.value
vector_dimension = self.config.vector_dimension or VectorDimensions.HUGGING_FACE.value
self.set_vector_dimension(vector_dimension=vector_dimension)

View File

@@ -32,4 +32,5 @@ class OpenAIEmbedder(BaseEmbedder):
model_name=self.config.model,
)
self.set_embedding_fn(embedding_fn=embedding_fn)
self.set_vector_dimension(vector_dimension=VectorDimensions.OPENAI.value)
vector_dimension = self.config.vector_dimension or VectorDimensions.OPENAI.value
self.set_vector_dimension(vector_dimension=vector_dimension)

View File

@@ -15,5 +15,5 @@ class VertexAIEmbedder(BaseEmbedder):
embedding_fn = BaseEmbedder._langchain_default_concept(embeddings)
self.set_embedding_fn(embedding_fn=embedding_fn)
vector_dimension = VectorDimensions.VERTEX_AI.value
vector_dimension = self.config.vector_dimension or VectorDimensions.VERTEX_AI.value
self.set_vector_dimension(vector_dimension=vector_dimension)

View File

@@ -438,6 +438,7 @@ def validate_config(config_data):
Optional("api_key"): str,
Optional("title"): str,
Optional("task_type"): str,
Optional("vector_dimension"): int,
},
},
Optional("embedding_model"): {
@@ -448,6 +449,7 @@ def validate_config(config_data):
Optional("api_key"): str,
Optional("title"): str,
Optional("task_type"): str,
Optional("vector_dimension"): int,
},
},
Optional("chunker"): {