[Enhancement] Add sentence-transformers embedding support

datastore/datastore.py

@@ -11,7 +11,7 @@
     QueryWithEmbedding,
 )
 from services.chunks import get_document_chunks
-from services.openai import get_embeddings
+from services.embedding import Embedding
 
 
 class DataStore(ABC):
@@ -56,7 +56,7 @@ async def query(self, queries: List[Query]) -> List[QueryResult]:
         """
         # get a list of of just the queries from the Query list
         query_texts = [query.query for query in queries]
-        query_embeddings = get_embeddings(query_texts)
+        query_embeddings = Embedding.instance().get_embeddings(query_texts)
         # hydrate the queries with embeddings
         queries_with_embeddings = [
             QueryWithEmbedding(**query.dict(), embedding=embedding)

datastore/factory.py

@@ -1,11 +1,14 @@
 from datastore.datastore import DataStore
+from services.embedding import Embedding
 import os
 
 
 async def get_datastore() -> DataStore:
     datastore = os.environ.get("DATASTORE")
     assert datastore is not None
 
+    dimension = Embedding.instance().dimension
+
     match datastore:
         case "chroma":
             from datastore.providers.chroma_datastore import ChromaDataStore
@@ -19,7 +22,7 @@ async def get_datastore() -> DataStore:
         case "pinecone":
             from datastore.providers.pinecone_datastore import PineconeDataStore
 
-            return PineconeDataStore()
+            return PineconeDataStore(dimension=dimension)
         case "weaviate":
             from datastore.providers.weaviate_datastore import WeaviateDataStore
 
@@ -35,11 +38,11 @@ async def get_datastore() -> DataStore:
         case "redis":
             from datastore.providers.redis_datastore import RedisDataStore
 
-            return await RedisDataStore.init()
+            return await RedisDataStore.init(dim=dimension)
         case "qdrant":
             from datastore.providers.qdrant_datastore import QdrantDataStore
 
-            return QdrantDataStore()
+            return QdrantDataStore(vector_size=dimension)
         case "azuresearch":
             from datastore.providers.azuresearch_datastore import AzureSearchDataStore
 

datastore/providers/analyticdb_datastore.py

@@ -30,17 +30,17 @@
     "host": os.environ.get("PG_HOST", "localhost"),
     "port": int(os.environ.get("PG_PORT", "5432")),
 }
-OUTPUT_DIM = 1536
 
 
 class AnalyticDBDataStore(DataStore):
-    def __init__(self, config: Dict[str, str] = PG_CONFIG):
+    def __init__(self, config: Dict[str, str] = PG_CONFIG, dimension=1536):
         self.collection_name = config["collection"]
         self.user = config["user"]
         self.password = config["password"]
         self.database = config["database"]
         self.host = config["host"]
         self.port = config["port"]
+        self.dimension = dimension
 
         self.connection_pool = SimpleConnectionPool(
             minconn=1,
@@ -99,7 +99,7 @@ def _create_embedding_index(self, cur: psycopg2.extensions.cursor):
                 USING ann(embedding)
                 WITH (
                     distancemeasure=L2,
-                    dim=OUTPUT_DIM,
+                    dim={self.dimension},
                     pq_segments=64,
                     hnsw_m=100,
                     pq_centers=2048

datastore/providers/pinecone_datastore.py

@@ -33,7 +33,7 @@
 
 
 class PineconeDataStore(DataStore):
-    def __init__(self):
+    def __init__(self, dimension=1536):
         # Check if the index name is specified and exists in Pinecone
         if PINECONE_INDEX and PINECONE_INDEX not in pinecone.list_indexes():
 
@@ -47,7 +47,7 @@ def __init__(self):
                 )
                 pinecone.create_index(
                     PINECONE_INDEX,
-                    dimension=1536,  # dimensionality of OpenAI ada v2 embeddings
+                    dimension=dimension,
                     metadata_config={"indexed": fields_to_index},
                 )
                 self.index = pinecone.Index(PINECONE_INDEX)

pyproject.toml

@@ -40,6 +40,7 @@ azure-search-documents = {version = "11.4.0a20230509004", source = "azure-sdk-de
 pgvector = "^0.1.7"
 psycopg2cffi = {version = "^2.9.0", optional = true}
 loguru = "^0.7.0"
+sentence-transformers = "^2.2.2"
 
 [tool.poetry.scripts]
 start = "server.main:start"

services/chunks.py

@@ -5,7 +5,7 @@
 
 import tiktoken
 
-from services.openai import get_embeddings
+from services.embedding import Embedding
 
 # Global variables
 tokenizer = tiktoken.get_encoding(
@@ -190,7 +190,7 @@ def get_document_chunks(
         ]
 
         # Get the embeddings for the batch texts
-        batch_embeddings = get_embeddings(batch_texts)
+        batch_embeddings = Embedding.instance().get_embeddings(batch_texts)
 
         # Append the batch embeddings to the embeddings list
         embeddings.extend(batch_embeddings)

services/embedding.py

@@ -0,0 +1,59 @@
+import os
+from typing import List
+from abc import ABC, abstractmethod
+
+import services.openai
+
+EMBEDDING_ENGINE = 'openai'
+_instance = None
+
+
+class Embedding(ABC):
+    @property
+    @abstractmethod
+    def dimension(self) -> int:
+        raise NotImplementedError
+
+    @abstractmethod
+    def get_embeddings(self, texts: List[str]) -> List[List[float]]:
+        raise NotImplementedError
+
+    @staticmethod
+    def instance() -> 'Embedding':
+        global _instance
+
+        if _instance is None:
+            datastore = os.getenv("EMBEDDING_ENGINE", EMBEDDING_ENGINE)
+            if datastore == 'openai':
+                _instance = OpenaiEmbedding()
+            elif datastore == 'sentence':
+                _instance = SentenceTransformerEmbedding()
+            else:
+                raise ValueError(
+                    f"Unsupported embedding engine: {datastore}. "
+                    f"Try one of the following: openai, sentence"
+                )
+
+        return _instance
+
+
+class OpenaiEmbedding(Embedding):
+    @property
+    def dimension(self) -> int:
+        return 1536
+
+    def get_embeddings(self, texts: List[str]) -> List[List[float]]:
+        return services.openai.get_embeddings(texts)
+
+
+class SentenceTransformerEmbedding(Embedding):
+    def __init__(self):
+        from sentence_transformers import SentenceTransformer
+        self.model = SentenceTransformer('sentence-transformers/all-MiniLM-L6-v2')
+
+    @property
+    def dimension(self) -> int:
+        return 384
+
+    def get_embeddings(self, texts: List[str]) -> List[List[float]]:
+        return self.model.encode(texts).tolist()