fix(embeddings): Fix embeddings names (camel-ai#384)

Author: dandansamax

camel/embeddings/base.py

@@ -12,22 +12,24 @@
 # limitations under the License.
 # =========== Copyright 2023 @ CAMEL-AI.org. All Rights Reserved. ===========
 from abc import ABC, abstractmethod
-from typing import Any, List
+from typing import Any, Generic, List, TypeVar
 
+T = TypeVar('T')
 
-class BaseEmbedding(ABC):
+
+class BaseEmbedding(ABC, Generic[T]):
     r"""Abstract base class for text embedding functionalities."""
 
     @abstractmethod
-    def embed_texts(
+    def embed_list(
         self,
-        texts: List[str],
+        objs: List[T],
         **kwargs: Any,
     ) -> List[List[float]]:
         r"""Generates embeddings for the given texts.
 
         Args:
-            texts (List[str]): The texts for which to generate the embeddings.
+            objs (List[T]): The objects for which to generate the embeddings.
             **kwargs (Any): Extra kwargs passed to the embedding API.
 
         Returns:
@@ -36,22 +38,22 @@ def embed_texts(
         """
         pass
 
-    def embed_text(
+    def embed(
         self,
-        text: str,
+        obj: T,
         **kwargs: Any,
     ) -> List[float]:
         r"""Generates an embedding for the given text.
 
         Args:
-            text (str): The text for which to generate the embedding.
+            obj (T): The object for which to generate the embedding.
             **kwargs (Any): Extra kwargs passed to the embedding API.
 
         Returns:
             List[float]: A list of floating-point numbers representing the
                 generated embedding.
         """
-        return self.embed_texts([text], **kwargs)[0]
+        return self.embed_list([obj], **kwargs)[0]
 
     @abstractmethod
     def get_output_dim(self) -> int:

camel/embeddings/openai_embedding.py

@@ -20,20 +20,20 @@
 from camel.utils import openai_api_key_required
 
 
-class OpenAIEmbedding(BaseEmbedding):
+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.ADA2`)
+            generating embeddings. (default: :obj:`ModelType.ADA_2`)
 
     Raises:
         RuntimeError: If an unsupported model type is specified.
     """
 
     def __init__(
         self,
-        model_type: EmbeddingModelType = EmbeddingModelType.ADA2,
+        model_type: EmbeddingModelType = EmbeddingModelType.ADA_2,
     ) -> None:
         if not model_type.is_openai:
             raise ValueError("Invalid OpenAI embedding model type.")
@@ -42,15 +42,15 @@ def __init__(
         self.client = OpenAI()
 
     @openai_api_key_required
-    def embed_texts(
+    def embed_list(
         self,
-        texts: List[str],
+        objs: List[str],
         **kwargs: Any,
     ) -> List[List[float]]:
         r"""Generates embeddings for the given texts.
 
         Args:
-            texts (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:
@@ -59,7 +59,7 @@ def embed_texts(
         """
         # TODO: count tokens
         response = self.client.embeddings.create(
-            input=texts,
+            input=objs,
             model=self.model_type.value,
             **kwargs,
         )

camel/types/enums.py

@@ -114,34 +114,34 @@ def validate_model_name(self, model_name: str) -> bool:
 
 
 class EmbeddingModelType(Enum):
-    ADA2 = "text-embedding-ada-002"
-    ADA1 = "text-embedding-ada-001"
-    BABBAGE1 = "text-embedding-babbage-001"
-    CURIE1 = "text-embedding-curie-001"
-    DAVINCI1 = "text-embedding-davinci-001"
+    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"
 
     @property
     def is_openai(self) -> bool:
         r"""Returns whether this type of models is an OpenAI-released model."""
         return self in {
-            EmbeddingModelType.ADA1,
-            EmbeddingModelType.ADA2,
-            EmbeddingModelType.BABBAGE1,
-            EmbeddingModelType.CURIE1,
-            EmbeddingModelType.DAVINCI1,
+            EmbeddingModelType.ADA_2,
+            EmbeddingModelType.ADA_1,
+            EmbeddingModelType.BABBAGE_1,
+            EmbeddingModelType.CURIE_1,
+            EmbeddingModelType.DAVINCI_1,
         }
 
     @property
     def output_dim(self) -> int:
-        if self is EmbeddingModelType.ADA2:
+        if self is EmbeddingModelType.ADA_2:
             return 1536
-        elif self is EmbeddingModelType.ADA1:
+        elif self is EmbeddingModelType.ADA_1:
             return 1024
-        elif self is EmbeddingModelType.BABBAGE1:
+        elif self is EmbeddingModelType.BABBAGE_1:
             return 2048
-        elif self is EmbeddingModelType.CURIE1:
+        elif self is EmbeddingModelType.CURIE_1:
             return 4096
-        elif self is EmbeddingModelType.DAVINCI1:
+        elif self is EmbeddingModelType.DAVINCI_1:
             return 12288
         else:
             raise ValueError(f"Unknown model type {self}.")

test/embeddings/test_openai_embedding.py

@@ -19,5 +19,5 @@
 @pytest.mark.parametrize("embedding_model", [OpenAIEmbedding()])
 def test_embedding(embedding_model: BaseEmbedding):
     text = "test embedding text."
-    vector = embedding_model.embed_text(text)
+    vector = embedding_model.embed(text)
     assert len(vector) == embedding_model.get_output_dim()