Complete SentenceTransformer integration

Introduced a way to mix and match embedding.

Only marginal improvement with Sentence Embedding.

todo:
clean up the config mess. new var EMBEDDING introduced

api/db/vector_store.py

@@ -1,10 +1,10 @@
 from abc import abstractmethod
 import os
 from qdrant_client import QdrantClient
-from langchain.embeddings import OpenAIEmbeddings, ElasticsearchEmbeddings
 from langchain.embeddings.sentence_transformer import SentenceTransformerEmbeddings
 from langchain.vectorstores import Qdrant, ElasticVectorSearch, VectorStore
 from qdrant_client.models import VectorParams, Distance
+from db.embedding import Embedding, EMBEDDINGS
 
 
 class ToyVectorStore:
@@ -12,13 +12,10 @@ class ToyVectorStore:
     @staticmethod
     def get_embedding():
         embedding = os.getenv("EMBEDDING")
-        if "SENTENCE" == embedding:
-            return SentenceTransformerEmbeddings()
-        elif "ELASTIC" == embedding:
-            return ElasticsearchEmbeddings()
-        else:
-            return OpenAIEmbeddings()
-
+        if not embedding:
+            return EMBEDDINGS["OPEN_AI"]
+        return EMBEDDINGS[embedding]
+    
     @staticmethod
     def get_instance():
         vector_store = os.getenv("STORE")
@@ -31,8 +28,8 @@ class ToyVectorStore:
             raise ValueError(f"Invalid vector store {vector_store}")
     
 
-    def __init__(self, embeddings):
-        self.embeddings = embeddings
+    def __init__(self, embedding: Embedding):
+        self.embedding = embedding
 
     @abstractmethod
     def get_collection(self, collection: str="test") -> VectorStore:
@@ -41,7 +38,7 @@ class ToyVectorStore:
         of collection
         """
         pass
-    
+
     @abstractmethod
     def create_collection(self, collection: str) -> None:
         """
@@ -51,11 +48,13 @@ class ToyVectorStore:
         pass
 
 class ElasticVectorStore(ToyVectorStore):
+    def __init__(self, embeddings):
+        super().__init__(embeddings)
 
     def get_collection(self, collection:str) -> ElasticVectorSearch:
         return ElasticVectorSearch(elasticsearch_url= os.getenv("ES_URL"),
-                               index_name= collection, embedding=self.embeddings)
-
+                               index_name= collection, embedding=self.embedding.embedding)
+    
     def create_collection(self, collection: str) -> None:
         store = self.get_collection(collection)
         store.create_index(store.client,collection, dict())
@@ -63,15 +62,17 @@ class ElasticVectorStore(ToyVectorStore):
 
 class QdrantVectorStore(ToyVectorStore):
 
-    def __init__(self):
+    def __init__(self, embeddings):
+        super().__init__(embeddings)
         self.client = QdrantClient(url=os.getenv("QDRANT_URL"),
                                         api_key=os.getenv("QDRANT_API_KEY"))
 
     def get_collection(self, collection: str) -> Qdrant:  
         return Qdrant(client=self.client,collection_name=collection,
-                      embeddings=self.embeddings)
+                      embeddings=self.embedding.embedding)
 
     def create_collection(self, collection: str) -> None:
         self.client.create_collection(collection_name=collection, 
-                        vectors_config=VectorParams(size=1536, distance=Distance.COSINE))
+                        vectors_config=VectorParams(size=self.embedding.dimension, 
+                                                    distance=Distance.COSINE))
     

api/routes/search.py

@@ -11,6 +11,7 @@ from langchain.vectorstores import Qdrant
 from langchain.schema import Document
 from langchain.chains.question_answering import load_qa_chain
 from langchain.llms import OpenAI
+from langchain.text_splitter import SentenceTransformersTokenTextSplitter
 from db.vector_store import ToyVectorStore
 
 router = APIRouter()
@@ -50,20 +51,26 @@ async def answer(name: str, query: str):
 
 async def generate_documents(file: UploadFile, file_name: str):
     num=0
-    async for txt in convert_documents(file):
+    async for txts in convert_documents(file):
         num += 1
-        document = Document(page_content=txt,metadata={"file": file_name, "page": num})
-        yield document
+        for txt in txts:
+            document = Document(page_content=txt,metadata={"file": file_name, "page": num})
+            yield document
  
 async def convert_documents(file: UploadFile):
+    splitter = SentenceTransformersTokenTextSplitter(chunk_overlap=0)
+
     #parse pdf document
     if file.content_type == 'application/pdf':
         content = await file.read()
         pdf_reader = PdfReader(io.BytesIO(content))
         try:
             for page in pdf_reader.pages:
-                yield page.extract_text()
+                yield splitter.split_text(page.extract_text())
         except Exception as e:
             print(f"Exception {e}")
+    elif "text" in file.content_type:
+        content = await file.read()
+        yield splitter.split_text(content.decode("utf-8"))
     else:
         return

requirements.txt

@@ -9,3 +9,5 @@ tiktoken
 faiss-cpu
 qdrant-client
 elasticsearch
+sentence_transformers
+