Update OpenAI text embeddings (#627)

Co-authored-by: Wendong-Fan <[email protected]>
Co-authored-by: Wendong <[email protected]>

Author: zechengz

camel/embeddings/openai_embedding.py

@@ -11,10 +11,12 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # =========== Copyright 2023 @ CAMEL-AI.org. All Rights Reserved. ===========
+from __future__ import annotations
+
 import os
-from typing import Any, List, Optional
+from typing import Any
 
-from openai import OpenAI
+from openai import NOT_GIVEN, NotGiven, OpenAI
 
 from camel.embeddings.base import BaseEmbedding
 from camel.types import EmbeddingModelType
@@ -25,47 +27,58 @@ class OpenAIEmbedding(BaseEmbedding[str]):
     r"""Provides text embedding functionalities using OpenAI's models.
 
     Args:
-        model (OpenAiEmbeddingModel, optional): The model type to be used for
-            generating embeddings. (default: :obj:`ModelType.ADA_2`)
-        api_key (Optional[str]): The API key for authenticating with the
+        model_type (EmbeddingModelType, optional): The model type to be
+            used for text embeddings.
+            (default: :obj:`TEXT_EMBEDDING_3_SMALL`)
+        api_key (str, optional): The API key for authenticating with the
             OpenAI service. (default: :obj:`None`)
+        dimensions (int, optional): The text embedding output dimensions.
+            (default: :obj:`NOT_GIVEN`)
 
     Raises:
         RuntimeError: If an unsupported model type is specified.
     """
 
     def __init__(
         self,
-        model_type: EmbeddingModelType = EmbeddingModelType.ADA_2,
-        api_key: Optional[str] = None,
+        model_type: EmbeddingModelType = (
+            EmbeddingModelType.TEXT_EMBEDDING_3_SMALL
+        ),
+        api_key: str | None = None,
+        dimensions: int | NotGiven = NOT_GIVEN,
     ) -> None:
         if not model_type.is_openai:
             raise ValueError("Invalid OpenAI embedding model type.")
         self.model_type = model_type
-        self.output_dim = model_type.output_dim
+        if dimensions == NOT_GIVEN:
+            self.output_dim = model_type.output_dim
+        else:
+            assert isinstance(dimensions, int)
+            self.output_dim = dimensions
         self._api_key = api_key or os.environ.get("OPENAI_API_KEY")
         self.client = OpenAI(timeout=60, max_retries=3, api_key=self._api_key)
 
     @model_api_key_required
     def embed_list(
         self,
-        objs: List[str],
+        objs: list[str],
         **kwargs: Any,
-    ) -> List[List[float]]:
+    ) -> list[list[float]]:
         r"""Generates embeddings for the given texts.
 
         Args:
-            objs (List[str]): The texts for which to generate the embeddings.
+            objs (list[str]): The texts for which to generate the embeddings.
             **kwargs (Any): Extra kwargs passed to the embedding API.
 
         Returns:
-            List[List[float]]: A list that represents the generated embedding
+            list[list[float]]: A list that represents the generated embedding
                 as a list of floating-point numbers.
         """
         # TODO: count tokens
         response = self.client.embeddings.create(
             input=objs,
             model=self.model_type.value,
+            dimensions=self.output_dim,
             **kwargs,
         )
         return [data.embedding for data in response.data]

camel/retrievers/auto_retriever.py

@@ -159,11 +159,11 @@ def _get_file_modified_date_from_storage(
     ) -> str:
         r"""Retrieves the last modified date and time of a given file. This
         function takes vector storage instance as input and returns the last
-        modified date from the meta data.
+        modified date from the metadata.
 
         Args:
             vector_storage_instance (BaseVectorStorage): The vector storage
-                where modified date is to be retrieved from meta data.
+                where modified date is to be retrieved from metadata.
 
         Returns:
             str: The last modified date from vector storage.

camel/types/enums.py

@@ -186,35 +186,27 @@ def validate_model_name(self, model_name: str) -> bool:
 
 
 class EmbeddingModelType(Enum):
-    ADA_2 = "text-embedding-ada-002"
-    ADA_1 = "text-embedding-ada-001"
-    BABBAGE_1 = "text-embedding-babbage-001"
-    CURIE_1 = "text-embedding-curie-001"
-    DAVINCI_1 = "text-embedding-davinci-001"
+    TEXT_EMBEDDING_ADA_2 = "text-embedding-ada-002"
+    TEXT_EMBEDDING_3_SMALL = "text-embedding-3-small"
+    TEXT_EMBEDDING_3_LARGE = "text-embedding-3-large"
 
     @property
     def is_openai(self) -> bool:
         r"""Returns whether this type of models is an OpenAI-released model."""
         return self in {
-            EmbeddingModelType.ADA_2,
-            EmbeddingModelType.ADA_1,
-            EmbeddingModelType.BABBAGE_1,
-            EmbeddingModelType.CURIE_1,
-            EmbeddingModelType.DAVINCI_1,
+            EmbeddingModelType.TEXT_EMBEDDING_ADA_2,
+            EmbeddingModelType.TEXT_EMBEDDING_3_SMALL,
+            EmbeddingModelType.TEXT_EMBEDDING_3_LARGE,
         }
 
     @property
     def output_dim(self) -> int:
-        if self is EmbeddingModelType.ADA_2:
+        if self is EmbeddingModelType.TEXT_EMBEDDING_ADA_2:
             return 1536
-        elif self is EmbeddingModelType.ADA_1:
-            return 1024
-        elif self is EmbeddingModelType.BABBAGE_1:
-            return 2048
-        elif self is EmbeddingModelType.CURIE_1:
-            return 4096
-        elif self is EmbeddingModelType.DAVINCI_1:
-            return 12288
+        elif self is EmbeddingModelType.TEXT_EMBEDDING_3_SMALL:
+            return 1536
+        elif self is EmbeddingModelType.TEXT_EMBEDDING_3_LARGE:
+            return 3072
         else:
             raise ValueError(f"Unknown model type {self}.")
 

docs/key_modules/embeddings.md

@@ -12,7 +12,7 @@ Consider two sentences:
 
 - "A child is kicking a football on a playground."
 
-Text embedding models would transform these sentences into two high-dimensional vector (*e.g.*, 1536 dimension if using `text-embedding-ada-002`). Despite different wordings, the vectors will be similar, capturing the shared concept of a child playing a ball game outdoors. This transformation into vectors allows machines to understand and compare the semantic similarities between the context.
+Text embedding models would transform these sentences into two high-dimensional vector (*e.g.*, 1536 dimension if using `text-embedding-3-small`). Despite different wordings, the vectors will be similar, capturing the shared concept of a child playing a ball game outdoors. This transformation into vectors allows machines to understand and compare the semantic similarities between the context.
 
 ### 1.2. Image Embeddings (WIP)
 Image embeddings convert images into numerical vectors, capturing essential features like shapes, colors, textures, and spatial hierarchies. This transformation is typically performed by Convolutional Neural Networks (CNNs) or other advanced neural network architectures designed for image processing. The resulting embeddings can be used for tasks like image classification, similarity comparison, and retrieval.
@@ -39,7 +39,7 @@ from camel.embeddings import OpenAIEmbedding
 from camel.types import EmbeddingModelType
 
 # Initialize the OpenAI embedding with a specific model
-openai_embedding = OpenAIEmbedding(model_type=EmbeddingModelType.ADA_2)
+openai_embedding = OpenAIEmbedding(model_type=EmbeddingModelType.TEXT_EMBEDDING_3_SMALL)
 
 # Generate embeddings for a list of texts
 embeddings = openai_embedding.embed_list(["Hello, world!", "Another example"])

test/embeddings/test_openai_embedding.py

@@ -15,7 +15,12 @@
 
 
 def test_openai_embedding():
-    text = "test embedding text."
     embedding_model = OpenAIEmbedding()
+    text = "test 1."
     vector = embedding_model.embed(text)
     assert len(vector) == embedding_model.get_output_dim()
+
+    embedding_model = OpenAIEmbedding(dimensions=256)
+    text = "test 2"
+    vector = embedding_model.embed(text)
+    assert len(vector) == embedding_model.get_output_dim() == 256

test/retrievers/test_auto_retriever.py

@@ -68,11 +68,11 @@ def test_get_file_modified_date_from_file(auto_retriever):
 
 def test_run_vector_retriever(auto_retriever):
     # Define mock data for testing
-    query_related = "what is camel"
+    query_related = "what is camel ai"
     query_unrealted = "unrelated query"
     content_input_paths = "https://www.camel-ai.org/"
     top_k = 1
-    similarity_threshold = 0.75
+    similarity_threshold = 0.5
 
     # Test with query related to the content in mock data
     result_related = auto_retriever.run_vector_retriever(