add llama index's ollama embedding and fix type hint of embedding_model (#1071)

rjwharry · web-flow · commit cb406f25e050 · 2025-01-21T17:40:58.000+09:00
* add llama index's ollama embedding and fix type hint of embedding_model

* unify the way load embedding models

* add ollama to supporting model types

* fix test code

* add test_load_embedding_model_from_dict

* remove unused import
diff --git a/autorag/data/chunk/base.py b/autorag/data/chunk/base.py
@@ -4,7 +4,7 @@
 
 import pandas as pd
 
-from autorag.embedding.base import embedding_models
+from autorag.embedding.base import EmbeddingModel
 from autorag.data import chunk_modules, sentence_splitter_modules
 from autorag.utils import result_to_dataframe
 
@@ -80,7 +80,7 @@ def get_embedding_model(_embed_model_str: str, _module_type: str):
 		if _embed_model_str == "openai":
 			if _module_type == "langchain_chunk":
 				_embed_model_str = "openai_langchain"
-		return embedding_models[_embed_model_str]()
+		return EmbeddingModel.load(_embed_model_str)()
 
 	# Add embed_model to kwargs
 	embedding_available_methods = ["semantic_llama_index", "semantic_langchain"]
diff --git a/autorag/embedding/base.py b/autorag/embedding/base.py
@@ -7,6 +7,7 @@
 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 llama_index.embeddings.ollama import OllamaEmbedding
 from langchain_openai.embeddings import OpenAIEmbeddings
 
 from autorag import LazyInit
@@ -35,6 +36,7 @@ def _get_vector(self) -> List[float]:
 	"mock": LazyInit(MockEmbeddingRandom, embed_dim=768),
 	# langchain
 	"openai_langchain": LazyInit(OpenAIEmbeddings),
+	"ollama": LazyInit(OllamaEmbedding),
 }
 
 try:
@@ -67,11 +69,13 @@ def _get_vector(self) -> List[float]:
 
 class EmbeddingModel:
 	@staticmethod
-	def load(config: Union[str, List[Dict]]):
+	def load(config: Union[str, Dict, List[Dict]]):
 		if isinstance(config, str):
 			return EmbeddingModel.load_from_str(config)
-		elif isinstance(config, list):
+		elif isinstance(config, dict):
 			return EmbeddingModel.load_from_dict(config)
+		elif isinstance(config, list):
+			return EmbeddingModel.load_from_list(config)
 		else:
 			raise ValueError("Invalid type of config")
 
@@ -83,11 +87,17 @@ def load_from_str(name: str):
 			raise ValueError(f"Embedding model '{name}' is not supported")
 
 	@staticmethod
-	def load_from_dict(option: List[dict]):
+	def load_from_list(option: List[dict]):
+		if len(option) != 1:
+			raise ValueError("Only one embedding model is supported")
+		return EmbeddingModel.load_from_dict(option[0])
+
+	@staticmethod
+	def load_from_dict(option: 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"]:
+			if target["type"] not in ["openai", "huggingface", "mock", "ollama"]:
 				raise ValueError(
 					f"Embedding model type '{target['type']}' is not supported"
 				)
@@ -102,17 +112,16 @@ def _get_huggingface_class():
 				return None
 			return getattr(module, "HuggingFaceEmbedding", None)
 
-		if len(option) != 1:
-			raise ValueError("Only one embedding model is supported")
-		_check_keys(option[0])
+		_check_keys(option)
 
-		model_options = option[0]
+		model_options = option
 		model_type = model_options.pop("type")
 
 		embedding_map = {
 			"openai": OpenAIEmbedding,
 			"mock": MockEmbeddingRandom,
 			"huggingface": _get_huggingface_class(),
+			"ollama": OllamaEmbedding,
 		}
 
 		embedding_class = embedding_map.get(model_type)
diff --git a/autorag/evaluation/util.py b/autorag/evaluation/util.py
@@ -1,7 +1,7 @@
 from copy import deepcopy
 from typing import Union, List, Dict, Tuple, Any
 
-from autorag.embedding.base import embedding_models
+from autorag.embedding.base import EmbeddingModel
 
 
 def cast_metrics(
@@ -38,6 +38,6 @@ def cast_metrics(
 
 def cast_embedding_model(key, value):
 	if key == "embedding_model":
-		return key, embedding_models[value]()
+		return key, EmbeddingModel.load(value)()
 	else:
 		return key, value
diff --git a/autorag/nodes/passageaugmenter/prev_next_augmenter.py b/autorag/nodes/passageaugmenter/prev_next_augmenter.py
@@ -1,9 +1,9 @@
-from typing import List
+from typing import List, Union
 
 import numpy as np
 import pandas as pd
 
-from autorag.embedding.base import embedding_models
+from autorag.embedding.base import EmbeddingModel
 from autorag.evaluation.metric.util import calculate_cosine_similarity
 from autorag.nodes.passageaugmenter.base import BasePassageAugmenter
 from autorag.utils.util import (
@@ -17,7 +17,7 @@
 
 class PrevNextPassageAugmenter(BasePassageAugmenter):
 	def __init__(
-		self, project_dir: str, embedding_model: str = "openai", *args, **kwargs
+		self, project_dir: str, embedding_model: Union[str, dict] = "openai", *args, **kwargs
 	):
 		"""
 		Initialize the PrevNextPassageAugmenter module.
@@ -35,7 +35,7 @@ def __init__(
 		self.slim_corpus_df = slim_corpus_df
 
 		# init embedding model
-		self.embedding_model = embedding_models[embedding_model]()
+		self.embedding_model = EmbeddingModel.load(embedding_model)()
 
 	def __del__(self):
 		del self.embedding_model
diff --git a/autorag/nodes/passagefilter/similarity_percentile_cutoff.py b/autorag/nodes/passagefilter/similarity_percentile_cutoff.py
@@ -4,7 +4,7 @@
 import numpy as np
 import pandas as pd
 
-from autorag.embedding.base import embedding_models
+from autorag.embedding.base import EmbeddingModel
 from autorag.evaluation.metric.util import calculate_cosine_similarity
 from autorag.nodes.passagefilter.base import BasePassageFilter
 from autorag.nodes.passagefilter.similarity_threshold_cutoff import (
@@ -24,8 +24,8 @@ def __init__(self, project_dir: Union[str, Path], *args, **kwargs):
 		        Default is "openai" which is OpenAI text-embedding-ada-002 embedding model.
 		"""
 		super().__init__(project_dir, *args, **kwargs)
-		embedding_model_str = kwargs.pop("embedding_model", "openai")
-		self.embedding_model = embedding_models[embedding_model_str]()
+		embedding_model = kwargs.pop("embedding_model", "openai")
+		self.embedding_model = EmbeddingModel.load(embedding_model)()
 
 	def __del__(self):
 		super().__del__()
diff --git a/autorag/nodes/passagefilter/similarity_threshold_cutoff.py b/autorag/nodes/passagefilter/similarity_threshold_cutoff.py
@@ -3,7 +3,7 @@
 import numpy as np
 import pandas as pd
 
-from autorag.embedding.base import embedding_models
+from autorag.embedding.base import EmbeddingModel
 from autorag.evaluation.metric.util import calculate_cosine_similarity
 from autorag.nodes.passagefilter.base import BasePassageFilter
 from autorag.utils.util import (
@@ -24,8 +24,8 @@ def __init__(self, project_dir: str, *args, **kwargs):
 		        Default is "openai" which is OpenAI text-embedding-ada-002 embedding model.
 		"""
 		super().__init__(project_dir, *args, **kwargs)
-		embedding_model_str = kwargs.get("embedding_model", "openai")
-		self.embedding_model = embedding_models[embedding_model_str]()
+		embedding_model= kwargs.get("embedding_model", "openai")
+		self.embedding_model = EmbeddingModel.load(embedding_model)()
 
 	def __del__(self):
 		del self.embedding_model
diff --git a/autorag/schema/module.py b/autorag/schema/module.py
@@ -22,3 +22,4 @@ def from_dict(cls, module_dict: Dict) -> "Module":
 		module_type = _module_dict.pop("module_type")
 		module_params = _module_dict
 		return cls(module_type, module_params)
+
diff --git a/autorag/utils/util.py b/autorag/utils/util.py
@@ -153,6 +153,7 @@ def make_combinations(target_dict: Dict[str, Any]) -> List[Dict[str, Any]]:
 		)
 	)
 
+
 	def delete_duplicate(x):
 		def is_hashable(obj):
 			try:
diff --git a/autorag/vectordb/chroma.py b/autorag/vectordb/chroma.py
@@ -1,4 +1,4 @@
-from typing import List, Optional, Dict, Tuple
+from typing import List, Optional, Dict, Tuple, Union
 
 from chromadb import (
 	EphemeralClient,
@@ -18,7 +18,7 @@
 class Chroma(BaseVectorStore):
 	def __init__(
 		self,
-		embedding_model: str,
+		embedding_model: Union[str, List[dict]],
 		collection_name: str,
 		embedding_batch: int = 100,
 		client_type: str = "persistent",
diff --git a/autorag/vectordb/couchbase.py b/autorag/vectordb/couchbase.py
@@ -6,7 +6,7 @@
 from couchbase.cluster import Cluster
 from couchbase.options import ClusterOptions
 
-from typing import List, Tuple, Optional
+from typing import List, Tuple, Optional, Union
 
 from autorag.utils.util import make_batch
 from autorag.vectordb import BaseVectorStore
@@ -17,7 +17,7 @@
 class Couchbase(BaseVectorStore):
 	def __init__(
 		self,
-		embedding_model: str,
+		embedding_model: Union[str, List[dict]],
 		bucket_name: str,
 		scope_name: str,
 		collection_name: str,
diff --git a/autorag/vectordb/milvus.py b/autorag/vectordb/milvus.py
@@ -1,5 +1,5 @@
 import logging
-from typing import Any, Dict, List, Tuple, Optional
+from typing import Any, Dict, List, Tuple, Optional, Union
 
 from pymilvus import (
 	DataType,
@@ -21,7 +21,7 @@
 class Milvus(BaseVectorStore):
 	def __init__(
 		self,
-		embedding_model: str,
+		embedding_model: Union[str, List[dict]],
 		collection_name: str,
 		embedding_batch: int = 100,
 		similarity_metric: str = "cosine",
diff --git a/autorag/vectordb/pinecone.py b/autorag/vectordb/pinecone.py
@@ -3,7 +3,7 @@
 from pinecone.grpc import PineconeGRPC as Pinecone_client
 from pinecone import ServerlessSpec
 
-from typing import List, Optional, Tuple
+from typing import List, Optional, Tuple, Union
 
 from autorag.utils.util import make_batch, apply_recursive
 from autorag.vectordb import BaseVectorStore
@@ -14,7 +14,7 @@
 class Pinecone(BaseVectorStore):
 	def __init__(
 		self,
-		embedding_model: str,
+		embedding_model: Union[str, List[dict]],
 		index_name: str,
 		embedding_batch: int = 100,
 		dimension: int = 1536,
diff --git a/autorag/vectordb/qdrant.py b/autorag/vectordb/qdrant.py
@@ -11,7 +11,7 @@
 	SearchRequest,
 )
 
-from typing import List, Tuple
+from typing import List, Tuple, Union
 
 from autorag.vectordb import BaseVectorStore
 
@@ -21,7 +21,7 @@
 class Qdrant(BaseVectorStore):
 	def __init__(
 		self,
-		embedding_model: str,
+		embedding_model: Union[str, List[dict]],
 		collection_name: str,
 		embedding_batch: int = 100,
 		similarity_metric: str = "cosine",
diff --git a/autorag/vectordb/weaviate.py b/autorag/vectordb/weaviate.py
@@ -6,7 +6,7 @@
 import weaviate.classes as wvc
 from weaviate.classes.query import MetadataQuery
 
-from typing import List, Optional, Tuple
+from typing import List, Optional, Tuple, Union
 
 from autorag.vectordb import BaseVectorStore
 
@@ -16,7 +16,7 @@
 class Weaviate(BaseVectorStore):
 	def __init__(
 		self,
-		embedding_model: str,
+		embedding_model: Union[str, List[dict]],
 		collection_name: str,
 		embedding_batch: int = 100,
 		similarity_metric: str = "cosine",
diff --git a/docs/source/local_model.md b/docs/source/local_model.md
@@ -162,6 +162,7 @@ We support the following embedding model types:
 - openai
 - huggingface
 - mock
+- ollama
 
 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,
@@ -191,7 +192,7 @@ vectordb:
   ...
 ```
 
-If you want to use your own embedding model, simply change the model_name at huggingface type embedding model configuration.
+If you want to use your own embedding model, simply change the model_name at huggingface(or ollama) type embedding model configuration.
 
 ### Supporting Embedding models (Legacy)
 
diff --git a/requirements.txt b/requirements.txt
@@ -40,6 +40,7 @@ llama-index-core>=0.11.0
 llama-index-readers-file
 # Embeddings
 llama-index-embeddings-openai
+llama-index-embeddings-ollama
 # LLMs
 llama-index-llms-openai>=0.2.7
 llama-index-llms-openai-like
diff --git a/tests/autorag/embedding/test_base.py b/tests/autorag/embedding/test_base.py
@@ -1,6 +1,5 @@
 import pytest
 from llama_index.embeddings.openai import OpenAIEmbedding
-
 from autorag.embedding.base import EmbeddingModel, MockEmbeddingRandom
 
 
@@ -29,26 +28,44 @@ def test_load_from_str_embedding_model():
 		EmbeddingModel.load_from_str("unsupported_model")
 
 
-def test_load_embedding_model_from_dict():
+def test_load_embedding_model_from_list():
 	# Test loading with missing keys
 	with pytest.raises(
 		ValueError, match="Both 'type' and 'model_name' must be provided"
 	):
-		EmbeddingModel.load_from_dict([{"type": "openai"}])
+		EmbeddingModel.load_from_list([{"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(
+		EmbeddingModel.load_from_list(
 			[{"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(
+		EmbeddingModel.load_from_list(
 			[
 				{"type": "openai", "model_name": "text-embedding-ada-002"},
 				{"type": "huggingface", "model_name": "BAAI/bge-small-en-v1.5"},
 			]
 		)
+
+def test_load_embedding_model_from_dict():
+	with pytest.raises(
+		ValueError, match="Both 'type' and 'model_name' must be provided"
+	):
+		embedding = 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"}
+		)
+
+	embedding = EmbeddingModel.load_from_dict({"type": "openai", "model_name": "text-embedding-ada-002"})
+	assert embedding is not None
+	assert isinstance(embedding(), OpenAIEmbedding)

Original file line number	Diff line number	Diff line change
`@@ -153,6 +153,7 @@ def make_combinations(target_dict: Dict[str, Any]) -> List[Dict[str, Any]]:`
`153`	`153`	`)`
`154`	`154`	`)`
`155`	`155`
	`156`	`+`
`156`	`157`	`def delete_duplicate(x):`
`157`	`158`	`def is_hashable(obj):`
`158`	`159`	`try:`