feat: enable dynamic embedding models #1060 (#1063)

* feat: dynamic embed init

* feat: apply embeddingbase class to vectordb

* feat: add tests and optimize codes

* fix: update import path of embedding_models from autorag to autorag.embedding.base

* fix: import err

* fix: autorag.embedding_models -> autorag.embedding.base.embedding_models

* simplify and delete _check_one_item function and use assertion

* return LazyInit at load_from_str and use initialized embedding model at vectordb class

* add docs for newer embedding model configuration

---------

Co-authored-by: Um Changyong <[email protected]>
Co-authored-by: Jeffrey (Dongkyu) Kim <[email protected]>

autorag/__init__.py

@@ -41,13 +41,6 @@ def __getattr__(self, name):
 		return getattr(self._instance, name)
 
 
-class MockEmbeddingRandom(MockEmbedding):
-	"""Mock embedding with random vectors."""
-
-	def _get_vector(self) -> List[float]:
-		return [random() for _ in range(self.embed_dim)]
-
-
 rich_format = "[%(filename)s:%(lineno)s] >> %(message)s"
 logging.basicConfig(
 	level="INFO", format=rich_format, handlers=[RichHandler(rich_tracebacks=True)]
@@ -62,49 +55,6 @@ def handle_exception(exc_type, exc_value, exc_traceback):
 
 sys.excepthook = handle_exception
 
-embedding_models = {
-	# llama index
-	"openai": LazyInit(
-		OpenAIEmbedding
-	),  # default model is OpenAIEmbeddingModelType.TEXT_EMBED_ADA_002
-	"openai_embed_3_large": LazyInit(
-		OpenAIEmbedding, model_name=OpenAIEmbeddingModelType.TEXT_EMBED_3_LARGE
-	),
-	"openai_embed_3_small": LazyInit(
-		OpenAIEmbedding, model_name=OpenAIEmbeddingModelType.TEXT_EMBED_3_SMALL
-	),
-	"mock": LazyInit(MockEmbeddingRandom, embed_dim=768),
-	# langchain
-	"openai_langchain": LazyInit(OpenAIEmbeddings),
-}
-
-try:
-	# you can use your own model in this way.
-	from llama_index.embeddings.huggingface import HuggingFaceEmbedding
-
-	embedding_models["huggingface_baai_bge_small"] = LazyInit(
-		HuggingFaceEmbedding, model_name="BAAI/bge-small-en-v1.5"
-	)
-	embedding_models["huggingface_cointegrated_rubert_tiny2"] = LazyInit(
-		HuggingFaceEmbedding, model_name="cointegrated/rubert-tiny2"
-	)
-	embedding_models["huggingface_all_mpnet_base_v2"] = LazyInit(
-		HuggingFaceEmbedding,
-		model_name="sentence-transformers/all-mpnet-base-v2",
-		max_length=512,
-	)
-	embedding_models["huggingface_bge_m3"] = LazyInit(
-		HuggingFaceEmbedding, model_name="BAAI/bge-m3"
-	)
-	embedding_models["huggingface_multilingual_e5_large"] = LazyInit(
-		HuggingFaceEmbedding, model_name="intfloat/multilingual-e5-large-instruct"
-	)
-except ImportError:
-	logger.info(
-		"You are using API version of AutoRAG."
-		"To use local version, run pip install 'AutoRAG[gpu]'"
-	)
-
 
 class AutoRAGBedrock(Bedrock):
 	async def acomplete(

autorag/data/chunk/base.py

@@ -4,7 +4,7 @@
 
 import pandas as pd
 
-from autorag import embedding_models
+from autorag.embedding.base import embedding_models
 from autorag.data import chunk_modules, sentence_splitter_modules
 from autorag.utils import result_to_dataframe
 

autorag/embedding/base.py

@@ -0,0 +1,122 @@
+import logging
+import sys
+
+from random import random
+from typing import List, Union, Dict
+
+from llama_index.core.embeddings.mock_embed_model import MockEmbedding
+from llama_index.embeddings.openai import OpenAIEmbedding
+from llama_index.embeddings.openai import OpenAIEmbeddingModelType
+from langchain_openai.embeddings import OpenAIEmbeddings
+
+from autorag import LazyInit
+
+logger = logging.getLogger("AutoRAG")
+
+
+class MockEmbeddingRandom(MockEmbedding):
+	"""Mock embedding with random vectors."""
+
+	def _get_vector(self) -> List[float]:
+		return [random() for _ in range(self.embed_dim)]
+
+
+embedding_models = {
+	# llama index
+	"openai": LazyInit(
+		OpenAIEmbedding
+	),  # default model is OpenAIEmbeddingModelType.TEXT_EMBED_ADA_002
+	"openai_embed_3_large": LazyInit(
+		OpenAIEmbedding, model_name=OpenAIEmbeddingModelType.TEXT_EMBED_3_LARGE
+	),
+	"openai_embed_3_small": LazyInit(
+		OpenAIEmbedding, model_name=OpenAIEmbeddingModelType.TEXT_EMBED_3_SMALL
+	),
+	"mock": LazyInit(MockEmbeddingRandom, embed_dim=768),
+	# langchain
+	"openai_langchain": LazyInit(OpenAIEmbeddings),
+}
+
+try:
+	# you can use your own model in this way.
+	from llama_index.embeddings.huggingface import HuggingFaceEmbedding
+
+	embedding_models["huggingface_baai_bge_small"] = LazyInit(
+		HuggingFaceEmbedding, model_name="BAAI/bge-small-en-v1.5"
+	)
+	embedding_models["huggingface_cointegrated_rubert_tiny2"] = LazyInit(
+		HuggingFaceEmbedding, model_name="cointegrated/rubert-tiny2"
+	)
+	embedding_models["huggingface_all_mpnet_base_v2"] = LazyInit(
+		HuggingFaceEmbedding,
+		model_name="sentence-transformers/all-mpnet-base-v2",
+		max_length=512,
+	)
+	embedding_models["huggingface_bge_m3"] = LazyInit(
+		HuggingFaceEmbedding, model_name="BAAI/bge-m3"
+	)
+	embedding_models["huggingface_multilingual_e5_large"] = LazyInit(
+		HuggingFaceEmbedding, model_name="intfloat/multilingual-e5-large-instruct"
+	)
+except ImportError:
+	logger.info(
+		"You are using API version of AutoRAG."
+		"To use local version, run pip install 'AutoRAG[gpu]'"
+	)
+
+
+class EmbeddingModel:
+	@staticmethod
+	def load(config: Union[str, List[Dict]]):
+		if isinstance(config, str):
+			return EmbeddingModel.load_from_str(config)
+		elif isinstance(config, list):
+			return EmbeddingModel.load_from_dict(config)
+		else:
+			raise ValueError("Invalid type of config")
+
+	@staticmethod
+	def load_from_str(name: str):
+		try:
+			return embedding_models[name]
+		except KeyError:
+			raise ValueError(f"Embedding model '{name}' is not supported")
+
+	@staticmethod
+	def load_from_dict(option: List[dict]):
+		def _check_keys(target: dict):
+			if "type" not in target or "model_name" not in target:
+				raise ValueError("Both 'type' and 'model_name' must be provided")
+			if target["type"] not in ["openai", "huggingface", "mock"]:
+				raise ValueError(
+					f"Embedding model type '{target['type']}' is not supported"
+				)
+
+		def _get_huggingface_class():
+			module = sys.modules.get("llama_index.embeddings.huggingface")
+			if not module:
+				logger.info(
+					"You are using API version of AutoRAG. "
+					"To use local version, run `pip install 'AutoRAG[gpu]'`."
+				)
+				return None
+			return getattr(module, "HuggingFaceEmbedding", None)
+
+		if len(option) != 1:
+			raise ValueError("Only one embedding model is supported")
+		_check_keys(option[0])
+
+		model_options = option[0]
+		model_type = model_options.pop("type")
+
+		embedding_map = {
+			"openai": OpenAIEmbedding,
+			"mock": MockEmbeddingRandom,
+			"huggingface": _get_huggingface_class(),
+		}
+
+		embedding_class = embedding_map.get(model_type)
+		if not embedding_class:
+			raise ValueError(f"Embedding model type '{model_type}' is not supported")
+
+		return LazyInit(embedding_class, **model_options)

autorag/evaluation/metric/generation.py

@@ -14,7 +14,7 @@
 from rouge_score.rouge_scorer import RougeScorer
 from sacrebleu.metrics.bleu import BLEU
 
-from autorag import embedding_models
+from autorag.embedding.base import embedding_models
 from autorag.evaluation.metric.deepeval_prompt import FaithfulnessTemplate
 from autorag.evaluation.metric.util import (
 	autorag_metric_loop,

autorag/evaluation/util.py

@@ -1,7 +1,7 @@
 from copy import deepcopy
 from typing import Union, List, Dict, Tuple, Any
 
-from autorag import embedding_models
+from autorag.embedding.base import embedding_models
 
 
 def cast_metrics(

autorag/nodes/passageaugmenter/prev_next_augmenter.py

@@ -3,7 +3,7 @@
 import numpy as np
 import pandas as pd
 
-from autorag import embedding_models
+from autorag.embedding.base import embedding_models
 from autorag.evaluation.metric.util import calculate_cosine_similarity
 from autorag.nodes.passageaugmenter.base import BasePassageAugmenter
 from autorag.utils.util import (

autorag/nodes/passagefilter/similarity_percentile_cutoff.py

@@ -4,7 +4,7 @@
 import numpy as np
 import pandas as pd
 
-from autorag import embedding_models
+from autorag.embedding.base import embedding_models
 from autorag.evaluation.metric.util import calculate_cosine_similarity
 from autorag.nodes.passagefilter.base import BasePassageFilter
 from autorag.nodes.passagefilter.similarity_threshold_cutoff import (

autorag/nodes/passagefilter/similarity_threshold_cutoff.py

@@ -3,7 +3,7 @@
 import numpy as np
 import pandas as pd
 
-from autorag import embedding_models
+from autorag.embedding.base import embedding_models
 from autorag.evaluation.metric.util import calculate_cosine_similarity
 from autorag.nodes.passagefilter.base import BasePassageFilter
 from autorag.utils.util import (

autorag/vectordb/base.py

@@ -1,22 +1,22 @@
 from abc import abstractmethod
-from typing import List, Tuple
+from typing import List, Tuple, Union
 
 from llama_index.embeddings.openai import OpenAIEmbedding
 
-from autorag import embedding_models
 from autorag.utils.util import openai_truncate_by_token
+from autorag.embedding.base import EmbeddingModel
 
 
 class BaseVectorStore:
 	support_similarity_metrics = ["l2", "ip", "cosine"]
 
 	def __init__(
 		self,
-		embedding_model: str,
+		embedding_model: Union[str, List[dict]],
 		similarity_metric: str = "cosine",
 		embedding_batch: int = 100,
 	):
-		self.embedding = embedding_models[embedding_model]()
+		self.embedding = EmbeddingModel.load(embedding_model)()
 		self.embedding_batch = embedding_batch
 		self.embedding.embed_batch_size = embedding_batch
 		assert (

docs/source/local_model.md

@@ -150,7 +150,50 @@ Modules that using an embedding model can take `embedding_model` parameter to sp
 
 - [vectordb](nodes/retrieval/vectordb.md)
 
-### Supporting Embedding models
+### Configure the model name in the YAML file
+
+For easier use, since v0.3.13, we support configuring model name in the YAML file.
+
+We support this new configuration option, plus the legacy configuration option that is described below.
+So if you used before v0.3.13 version, it is okay to use the legacy configuration option.
+
+We support the following embedding model types:
+
+- openai
+- huggingface
+- mock
+
+You can configure the embedding model option directly in the YAML file `vectordb` section.
+If you want to know how to configure vectordb in AutoRAG,
+please go [here](https://docs.auto-rag.com/integration/vectordb/vectordb.html).
+
+For example,
+
+```yaml
+vectordb:
+- name: autorag_test
+  db_type: milvus
+  embedding_model:
+  - type: huggingface
+    model_name: intfloat/multilingual-e5-large-instruct
+  ...
+```
+
+or
+
+```yaml
+vectordb:
+- name: autorag_test
+  db_type: milvus
+  embedding_model:
+  - type: openai
+    model_name: text-embedding-3-small
+  ...
+```
+
+If you want to use your own embedding model, simply change the model_name at huggingface type embedding model configuration.
+
+### Supporting Embedding models (Legacy)
 
 As default, we support OpenAI embedding models and some of the local models.
 To change the embedding model, you can change the `embedding_model` parameter to the following values:
@@ -184,7 +227,7 @@ nodes:
             vectordb: chroma_openai
 ```
 
-### Add your embedding models
+### Add your embedding models (Legacy)
 
 You can add more embedding models for AutoRAG.
 You can add it by simply calling `autorag.embedding_models` and add new key and value.

tests/autorag/embedding/test_base.py

@@ -0,0 +1,54 @@
+import pytest
+from llama_index.embeddings.openai import OpenAIEmbedding
+
+from autorag.embedding.base import EmbeddingModel, MockEmbeddingRandom
+
+
+def test_load_embedding_model():
+	embedding = EmbeddingModel.load("mock")
+	assert embedding is not None
+	assert isinstance(embedding(), MockEmbeddingRandom)
+
+	embedding = EmbeddingModel.load(
+		[{"type": "openai", "model_name": "text-embedding-ada-002"}]
+	)
+	assert embedding is not None
+	assert isinstance(embedding(), OpenAIEmbedding)
+
+
+def test_load_from_str_embedding_model():
+	# Test loading a supported embedding model
+	embedding = EmbeddingModel.load_from_str("mock")
+	assert embedding is not None
+	assert isinstance(embedding(), MockEmbeddingRandom)
+
+	# Test loading an unsupported embedding model
+	with pytest.raises(
+		ValueError, match="Embedding model 'unsupported_model' is not supported"
+	):
+		EmbeddingModel.load_from_str("unsupported_model")
+
+
+def test_load_embedding_model_from_dict():
+	# Test loading with missing keys
+	with pytest.raises(
+		ValueError, match="Both 'type' and 'model_name' must be provided"
+	):
+		EmbeddingModel.load_from_dict([{"type": "openai"}])
+
+	# Test loading with an unsupported type
+	with pytest.raises(
+		ValueError, match="Embedding model type 'unsupported_type' is not supported"
+	):
+		EmbeddingModel.load_from_dict(
+			[{"type": "unsupported_type", "model_name": "some-model"}]
+		)
+
+	# Test loading with multiple items
+	with pytest.raises(ValueError, match="Only one embedding model is supported"):
+		EmbeddingModel.load_from_dict(
+			[
+				{"type": "openai", "model_name": "text-embedding-ada-002"},
+				{"type": "huggingface", "model_name": "BAAI/bge-small-en-v1.5"},
+			]
+		)

tests/autorag/evaluate/test_evaluate_util.py

@@ -1,4 +1,4 @@
-from autorag import embedding_models
+from autorag.embedding.base import embedding_models
 from autorag.evaluation.util import cast_metrics
 
 

tests/autorag/nodes/queryexpansion/test_query_expansion_run.py

@@ -10,7 +10,7 @@
 from llama_index.core.base.llms.types import CompletionResponse
 from llama_index.llms.openai import OpenAI
 
-from autorag import embedding_models, MockEmbeddingRandom, LazyInit
+from autorag.embedding.base import embedding_models, MockEmbeddingRandom, LazyInit
 from autorag.nodes.queryexpansion import QueryDecompose, HyDE
 from autorag.nodes.queryexpansion.run import evaluate_one_query_expansion_node
 from autorag.nodes.queryexpansion.run import run_query_expansion_node