From cc4d9973052b39852923631271ee7bc7066c203e Mon Sep 17 00:00:00 2001
From: kimbwook <bwook00@naver.com>
Date: Sat, 3 Feb 2024 14:35:50 +0000
Subject: [PATCH 1/3] add embedding_models and generator_models

---
 autorag/__init__.py | 24 ++++++++++++++++++++++--
 1 file changed, 22 insertions(+), 2 deletions(-)

diff --git a/autorag/__init__.py b/autorag/__init__.py
index 3c1f9b7e7..822d8a688 100644
--- a/autorag/__init__.py
+++ b/autorag/__init__.py
@@ -5,8 +5,9 @@
 
 from rich.logging import RichHandler
 
-from llama_index import OpenAIEmbedding
-from llama_index.llms import OpenAI
+from llama_index.embeddings import OpenAIEmbedding, HuggingFaceEmbedding, InstructorEmbedding, OptimumEmbedding
+from llama_index.llms import OpenAI, Anthropic, AzureOpenAI, HuggingFaceLLM, LangChainLLM, GradientBaseModelLLM, \
+    GradientModelAdapterLLM, LiteLLM, LlamaCPP, OpenAILike, OpenLLM, PaLM, PredibaseLLM, Replicate, Xinference
 
 root_path = os.path.dirname(os.path.dirname(os.path.realpath(__file__)))
 version_path = os.path.join(root_path, 'VERSION')
@@ -16,10 +17,29 @@
 
 embedding_models = {
     'openai': OpenAIEmbedding(),
+    # You can write your own model in this way.
+    'huggingface_baai_bge_small': HuggingFaceEmbedding(model_name="BAAI/bge-small-en-v1.5"),
+    'huggingface_cointegrated_rubert_tiny2': HuggingFaceEmbedding(model_name="cointegrated/rubert-tiny2"),
+    'instructor_hkunlp': InstructorEmbedding(model_name="hkunlp/instructor-base"),
+    'optimum_baai_bge_small': OptimumEmbedding.create_and_save_optimum_model("BAAI/bge-small-en-v1.5", "./bge_onnx"),
 }
 
 generator_models = {
     'openai': OpenAI,
+    'anthropic': Anthropic,
+    'azureopenai': AzureOpenAI,
+    'huggingfacellm': HuggingFaceLLM,
+    'langchainllm': LangChainLLM,
+    'gradientbasemodelllm': GradientBaseModelLLM,
+    'gradientmodeladapterllm': GradientModelAdapterLLM,
+    'litellm': LiteLLM,
+    'llamacpp': LlamaCPP,
+    'openailike': OpenAILike,
+    'openllm': OpenLLM,
+    'palm': PaLM,
+    'predibasellm': PredibaseLLM,
+    'replicate': Replicate,
+    'xinference': Xinference,
 }
 
 rich_format = "[%(filename)s:%(lineno)s] >> %(message)s"

From 3e3593faad7d890dff6cf1b06f5d5247dc4c58f5 Mon Sep 17 00:00:00 2001
From: kimbwook <bwook00@naver.com>
Date: Sat, 3 Feb 2024 15:17:58 +0000
Subject: [PATCH 2/3] edit embedding models

---
 autorag/__init__.py | 13 ++++++++-----
 1 file changed, 8 insertions(+), 5 deletions(-)

diff --git a/autorag/__init__.py b/autorag/__init__.py
index 822d8a688..f1d49835a 100644
--- a/autorag/__init__.py
+++ b/autorag/__init__.py
@@ -5,7 +5,8 @@
 
 from rich.logging import RichHandler
 
-from llama_index.embeddings import OpenAIEmbedding, HuggingFaceEmbedding, InstructorEmbedding, OptimumEmbedding
+from llama_index.embeddings import OpenAIEmbedding, HuggingFaceEmbedding
+from llama_index.embeddings.openai import OpenAIEmbeddingModelType
 from llama_index.llms import OpenAI, Anthropic, AzureOpenAI, HuggingFaceLLM, LangChainLLM, GradientBaseModelLLM, \
     GradientModelAdapterLLM, LiteLLM, LlamaCPP, OpenAILike, OpenLLM, PaLM, PredibaseLLM, Replicate, Xinference
 
@@ -16,12 +17,14 @@
     __version__ = f.read().strip()
 
 embedding_models = {
-    'openai': OpenAIEmbedding(),
-    # You can write your own model in this way.
+    'openai': OpenAIEmbedding(),  # default model is OpenAIEmbeddingModelType.TEXT_EMBED_ADA_002
+    'openai_babbage': OpenAIEmbedding(model=OpenAIEmbeddingModelType.BABBAGE),
+    'openai_ada': OpenAIEmbedding(model=OpenAIEmbeddingModelType.ADA),
+    'openai_davinci': OpenAIEmbedding(model=OpenAIEmbeddingModelType.DAVINCI),
+    'openai_curie': OpenAIEmbedding(model=OpenAIEmbeddingModelType.CURIE),
+    # you can change your own model in this way.
     'huggingface_baai_bge_small': HuggingFaceEmbedding(model_name="BAAI/bge-small-en-v1.5"),
     'huggingface_cointegrated_rubert_tiny2': HuggingFaceEmbedding(model_name="cointegrated/rubert-tiny2"),
-    'instructor_hkunlp': InstructorEmbedding(model_name="hkunlp/instructor-base"),
-    'optimum_baai_bge_small': OptimumEmbedding.create_and_save_optimum_model("BAAI/bge-small-en-v1.5", "./bge_onnx"),
 }
 
 generator_models = {

From 595f9872c7bd13ed4a52a732b056fa10e02ae4c5 Mon Sep 17 00:00:00 2001
From: kimbwook <bwook00@naver.com>
Date: Sat, 3 Feb 2024 15:20:24 +0000
Subject: [PATCH 3/3] edit annotation

---
 autorag/__init__.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/autorag/__init__.py b/autorag/__init__.py
index f1d49835a..df88f4acf 100644
--- a/autorag/__init__.py
+++ b/autorag/__init__.py
@@ -22,7 +22,7 @@
     'openai_ada': OpenAIEmbedding(model=OpenAIEmbeddingModelType.ADA),
     'openai_davinci': OpenAIEmbedding(model=OpenAIEmbeddingModelType.DAVINCI),
     'openai_curie': OpenAIEmbedding(model=OpenAIEmbeddingModelType.CURIE),
-    # you can change your own model in this way.
+    # you can use your own model in this way.
     'huggingface_baai_bge_small': HuggingFaceEmbedding(model_name="BAAI/bge-small-en-v1.5"),
     'huggingface_cointegrated_rubert_tiny2': HuggingFaceEmbedding(model_name="cointegrated/rubert-tiny2"),
 }