Update sentence transformer text embeddings (#628)

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

camel/embeddings/base.py

@@ -11,8 +11,10 @@
 # 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
+
 from abc import ABC, abstractmethod
-from typing import Any, Generic, List, TypeVar
+from typing import Any, Generic, TypeVar
 
 T = TypeVar('T')
 
@@ -23,35 +25,34 @@ class BaseEmbedding(ABC, Generic[T]):
     @abstractmethod
     def embed_list(
         self,
-        objs: List[T],
+        objs: list[T],
         **kwargs: Any,
-    ) -> List[List[float]]:
+    ) -> list[list[float]]:
         r"""Generates embeddings for the given texts.
 
         Args:
-            objs (List[T]): The objects 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:
-            List[List[float]]: A list that represents the
-            generated embedding as a list of floating-point numbers or a
-            numpy matrix with embeddings.
+            list[list[float]]: A list that represents the
+                generated embedding as a list of floating-point numbers.
         """
         pass
 
     def embed(
         self,
         obj: T,
         **kwargs: Any,
-    ) -> List[float]:
+    ) -> list[float]:
         r"""Generates an embedding for the given text.
 
         Args:
             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
+            list[float]: A list of floating-point numbers representing the
                 generated embedding.
         """
         return self.embed_list([obj], **kwargs)[0]

camel/embeddings/sentence_transformers_embeddings.py

@@ -11,56 +11,70 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # =========== Copyright 2023 @ CAMEL-AI.org. All Rights Reserved. ===========
-from typing import Any, List
+from __future__ import annotations
+
+from typing import Any
+
+from numpy import ndarray
 
 from camel.embeddings.base import BaseEmbedding
 
 
 class SentenceTransformerEncoder(BaseEmbedding[str]):
-    r"""This class provides functionalities to generate embeddings
-    using a specified model from `Sentence Transformers`.
+    r"""This class provides functionalities to generate text
+    embeddings using `Sentence Transformers`.
 
     References:
         https://www.sbert.net/
     """
 
-    def __init__(self, model_name: str = 'intfloat/e5-large-v2'):
+    def __init__(
+        self,
+        model_name: str = "intfloat/e5-large-v2",
+        **kwargs,
+    ):
         r"""Initializes the: obj: `SentenceTransformerEmbedding` class
         with the specified transformer model.
 
         Args:
             model_name (str, optional): The name of the model to use.
-                                        Defaults to `intfloat/e5-large-v2`.
+                (default: :obj:`intfloat/e5-large-v2`)
+            **kwargs (optional): Additional arguments of
+                :class:`SentenceTransformer`, such as :obj:`prompts` etc.
         """
         from sentence_transformers import SentenceTransformer
 
-        self.model = SentenceTransformer(model_name)
+        self.model = SentenceTransformer(model_name, **kwargs)
 
     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 using the model.
 
         Args:
-            objs (List[str]): The texts for which to generate the
-            embeddings.
+            objs (list[str]): The texts for which to generate the
+                embeddings.
 
         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.
         """
         if not objs:
             raise ValueError("Input text list is empty")
-        return self.model.encode(
+        embeddings = self.model.encode(
             objs, normalize_embeddings=True, **kwargs
-        ).tolist()
+        )
+        assert isinstance(embeddings, ndarray)
+        return embeddings.tolist()
 
     def get_output_dim(self) -> int:
         r"""Returns the output dimension of the embeddings.
 
         Returns:
             int: The dimensionality of the embeddings.
         """
-        return self.model.get_sentence_embedding_dimension()
+        output_dim = self.model.get_sentence_embedding_dimension()
+        assert isinstance(output_dim, int)
+        return output_dim