Add dimension support

pinecone_text/dense/base_dense_ecoder.py

@@ -1,8 +1,18 @@
-from typing import List, Union
+from typing import List, Union, Any, Optional
 from abc import ABC, abstractmethod
+from functools import cached_property
 
 
 class BaseDenseEncoder(ABC):
+    def __init__(self, *, dimension: Optional[int] = None, **kwargs: Any):
+        self._dimension = dimension
+
+    @cached_property
+    def dimension(self) -> int:
+        if self._dimension is None:
+            return len(self.encode_documents("hello"))
+        return self._dimension
+
     @abstractmethod
     def encode_documents(
         self, texts: Union[str, List[str]]

pinecone_text/dense/cohere_encoder.py

@@ -66,6 +66,7 @@ def __init__(
                 + "\n"
                 + "\n".join([f"- {x}" for x in CohereEncoderName.list_models()])
             )
+        super().__init__()
         self._model_name = model_name
         self._client = cohere.Client(api_key=api_key, **kwargs)
 

pinecone_text/dense/jina_encoder.py

@@ -33,6 +33,7 @@ def __init__(
             raise ValueError(
                 "JinaEncoder requires an API key to work. Please provide `api_key` argument or set `JINA_API_KEY` environment variable"
             )
+        super().__init__()
         self._model_name = model_name
         self._session = requests.Session()
         self._session.headers.update(

pinecone_text/dense/openai_encoder.py

@@ -1,5 +1,5 @@
 import os
-from typing import Union, List, Any, Optional
+from typing import Union, List, Any, Optional, Dict
 from pinecone_text.dense.base_dense_ecoder import BaseDenseEncoder
 
 try:
@@ -38,6 +38,8 @@ class OpenAIEncoder(BaseDenseEncoder):
     def __init__(
         self,
         model_name: str = "text-embedding-ada-002",
+        *,
+        dimension: Optional[int] = None,
         **kwargs: Any,
     ):
         if not _openai_installed:
@@ -46,7 +48,9 @@ def __init__(
                 "dependencies by running: "
                 "`pip install pinecone-text[openai]"
             )
+        super().__init__(dimension=dimension)
         self._model_name = model_name
+        self._dimension = dimension
         self._client = self._create_client(**kwargs)
 
     @staticmethod
@@ -76,9 +80,13 @@ def _encode(
             )
 
         try:
-            response = self._client.embeddings.create(
-                input=texts_input, model=self._model_name
+            params: Dict[str, Any] = dict(
+                input=texts_input,
+                model=self._model_name,
             )
+            if self._dimension:
+                params["dimensions"] = self._dimension
+            response = self._client.embeddings.create(**params)
         except OpenAIError as e:
             # TODO: consider wrapping external provider errors
             raise e

pinecone_text/dense/sentence_transformer_encoder.py

@@ -48,7 +48,7 @@ def __init__(
                 "Failed to import sentence transformers. Make sure you install dense "
                 "extra dependencies by running: `pip install pinecone-text[dense]`"
             )
-
+        super().__init__()
         device = device or ("cuda" if torch.cuda.is_available() else "cpu")
         self.document_encoder = SentenceTransformer(
             document_encoder_name, device=device

pyproject.toml

@@ -22,6 +22,7 @@ numpy = [
 ]
 requests = "^2.25.0"
 types-requests = "^2.25.0"
+python-dotenv = "^1.0.1"
 
 
 [tool.poetry.extras]
@@ -39,6 +40,7 @@ mypy = "^1.0.1"
 black = "^23.1.0"
 pytest-cov = "^4.0.0"
 pdoc = "^13.0.0"
+pytest-dotenv = "^0.5.2"
 
 [build-system]
 requires = ["poetry-core"]

tests/system/test_openai_encoder.py

@@ -3,7 +3,6 @@
 from pinecone_text.dense import OpenAIEncoder, AzureOpenAIEncoder
 from openai import BadRequestError, AuthenticationError
 
-
 DEFAULT_DIMENSION = 1536
 
 
@@ -87,3 +86,15 @@ def test_encode_invalid_input(openai_encoder, encoding_function):
 def test_encode_error_handling(openai_encoder, encoding_function):
     with pytest.raises(BadRequestError):
         encode_by_type(openai_encoder, encoding_function, "text is too long" * 10000)
+
+
+def test_dimension(openai_encoder):
+    assert openai_encoder.dimension == DEFAULT_DIMENSION
+
+
+def test_openai_encoder_with_dimension():
+    dimension = 10
+    encoder = OpenAIEncoder(model_name="text-embedding-3-small", dimension=dimension)
+    assert encoder.dimension == dimension
+    result = encoder.encode_documents("test text")
+    assert len(result) == dimension