WIP: expose dimensions in the Embedders

Author: stellasia

examples/customize/embeddings/cohere_embeddings.py

@@ -4,8 +4,14 @@
 api_key = None
 
 embeder = CohereEmbeddings(
-    model="embed-english-v3.0",
+    model="embed-v4.0",
     api_key=api_key,
 )
-res = embeder.embed_query("my question")
+res = embeder.embed_query(
+    "my question",
+    # optionally, set output dimensions if it's supported by the model
+    dimensions=256,
+    input_type="search_query",
+)
+print("Embedding dimensions", len(res))
 print(res[:10])

examples/customize/embeddings/custom_embeddings.py

@@ -8,7 +8,7 @@ class CustomEmbeddings(Embedder):
     def __init__(self, dimension: int = 10, **kwargs: Any):
         self.dimension = dimension
 
-    def embed_query(self, input: str) -> list[float]:
+    def embed_query(self, input: str, **kwargs) -> list[float]:
         return [random.random() for _ in range(self.dimension)]
 
 

examples/customize/embeddings/mistalai_embeddings.py

@@ -8,5 +8,10 @@
 api_key = None
 
 embeder = MistralAIEmbeddings(model="mistral-embed", api_key=api_key)
-res = embeder.embed_query("my question")
+res = embeder.embed_query(
+    "my question",
+    # optionally, set output dimensions
+    dimensions=256,
+)
+print("Embedding dimensions", len(res))
 print(res[:10])

examples/customize/embeddings/openai_embeddings.py

@@ -7,6 +7,12 @@
 # set api key here on in the OPENAI_API_KEY env var
 api_key = None
 
-embeder = OpenAIEmbeddings(model="text-embedding-ada-002", api_key=api_key)
-res = embeder.embed_query("my question")
+embeder = OpenAIEmbeddings(model="text-embedding-3-small", api_key=api_key)
+res = embeder.embed_query(
+    "my question",
+    # optionally, set output dimensions
+    # dimensions=256,
+)
+
+print("Embedding dimensions", len(res))
 print(res[:10])

examples/customize/embeddings/vertexai_embeddings.py

@@ -5,5 +5,9 @@
 from neo4j_graphrag.embeddings import VertexAIEmbeddings
 
 embeder = VertexAIEmbeddings(model="text-embedding-005")
-res = embeder.embed_query("my question")
+res = embeder.embed_query(
+    "my question",
+    dimensions=256,
+)
+print("Embedding dimensions", len(res))
 print(res[:10])

src/neo4j_graphrag/embeddings/base.py

@@ -24,11 +24,14 @@ class Embedder(ABC):
     """
 
     @abstractmethod
-    def embed_query(self, text: str) -> list[float]:
+    def embed_query(
+        self, text: str, dimensions: int | None = None, **kwargs
+    ) -> list[float]:
         """Embed query text.
 
         Args:
             text (str): Text to convert to vector embedding
+            dimensions (Optional[int]): The number of dimensions the resulting output embeddings should have. Only for models supporting it.
 
         Returns:
             list[float]: A vector embedding.

src/neo4j_graphrag/embeddings/cohere.py

@@ -32,12 +32,23 @@ def __init__(self, model: str = "", **kwargs: Any) -> None:
                 Please install it with `pip install "neo4j-graphrag[cohere]"`."""
             )
         self.model = model
-        self.client = cohere.Client(**kwargs)
+        self.client = cohere.ClientV2(**kwargs)
 
-    def embed_query(self, text: str, **kwargs: Any) -> list[float]:
+    def embed_query(
+        self, text: str, dimensions: int | None = None, **kwargs: Any
+    ) -> list[float]:
+        """
+        Generate embeddings for a given query using a Cohere text embedding model.
+
+        Args:
+            text (str): The text to generate an embedding for.
+            dimensions (Optional[int]): The number of dimensions the resulting output embeddings should have. Only for models supporting it.
+            **kwargs (Any): Additional keyword arguments to pass to the Cohere ClientV2.embed method.
+        """
         response = self.client.embed(
             texts=[text],
             model=self.model,
+            output_dimension=dimensions,
             **kwargs,
         )
-        return response.embeddings[0]  # type: ignore
+        return response.embeddings.float[0]  # type: ignore

src/neo4j_graphrag/embeddings/mistral.py

@@ -48,16 +48,19 @@ def __init__(self, model: str = "mistral-embed", **kwargs: Any) -> None:
         self.model = model
         self.mistral_client = Mistral(api_key=api_key, **kwargs)
 
-    def embed_query(self, text: str, **kwargs: Any) -> list[float]:
+    def embed_query(
+        self, text: str, dimensions: int | None = None, **kwargs: Any
+    ) -> list[float]:
         """
         Generate embeddings for a given query using a Mistral AI text embedding model.
 
         Args:
             text (str): The text to generate an embedding for.
-            **kwargs (Any): Additional keyword arguments to pass to the Mistral AI client.
+            dimensions (Optional[int]): The number of dimensions the resulting output embeddings should have. Only for models supporting it.
+            **kwargs (Any): Additional keyword arguments to pass to the embeddings.create method.
         """
         embeddings_batch_response = self.mistral_client.embeddings.create(
-            model=self.model, inputs=[text], **kwargs
+            model=self.model, inputs=[text], output_dimension=dimensions, **kwargs
         )
         if embeddings_batch_response is None or not embeddings_batch_response.data:
             raise EmbeddingsGenerationError("Failed to retrieve embeddings.")

src/neo4j_graphrag/embeddings/openai.py

@@ -51,15 +51,21 @@ def _initialize_client(self, **kwargs: Any) -> Any:
         """
         pass
 
-    def embed_query(self, text: str, **kwargs: Any) -> list[float]:
+    def embed_query(
+        self, text: str, dimensions: int | None = None, **kwargs: Any
+    ) -> list[float]:
         """
         Generate embeddings for a given query using an OpenAI text embedding model.
 
         Args:
             text (str): The text to generate an embedding for.
+            dimensions (Optional[int]): The number of dimensions the resulting output embeddings should have. Only for models supporting it.
+
             **kwargs (Any): Additional arguments to pass to the OpenAI embedding generation function.
         """
-        response = self.client.embeddings.create(input=text, model=self.model, **kwargs)
+        response = self.client.embeddings.create(
+            input=text, model=self.model, dimensions=dimensions, **kwargs
+        )
         embedding: list[float] = response.data[0].embedding
         return embedding
 

src/neo4j_graphrag/embeddings/sentence_transformers.py

@@ -35,7 +35,7 @@ def __init__(
         self.np = np
         self.model = sentence_transformers.SentenceTransformer(model, *args, **kwargs)
 
-    def embed_query(self, text: str) -> Any:
+    def embed_query(self, text: str, **kwargs) -> Any:
         result = self.model.encode([text])
         if isinstance(result, self.torch.Tensor) or isinstance(result, self.np.ndarray):
             return result.flatten().tolist()