added Faiss support

.env.example

@@ -57,7 +57,8 @@ LLAMACPP_MODEL_PATH="./models/wizardLM-7B.ggmlv3.q4_1.bin"
 LLAMACPP_DOWNLOAD_LINK=https://huggingface.co/TheBloke/wizardLM-7B-GGML/resolve/main/wizardLM-7B.ggmlv3.q4_1.bin
 
 # Index for AI Books PDF files - chunk_size=1024 chunk_overlap=512
-CHROMADB_INDEX_PATH="./data/chromadb_1024_512/"
+# CHROMADB_INDEX_PATH="./data/chromadb_1024_512/"
+FAISS_INDEX_PATH="./data/faiss_1024_512/"
 
 QUESTIONS_FILE_PATH="./data/questions.txt"
 

.gitattributes

@@ -33,3 +33,5 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+data/faiss_1024_512/index.faiss filter=lfs diff=lfs merge=lfs -text
+data/faiss_1024_512/index.pkl filter=lfs diff=lfs merge=lfs -text

data/faiss_1024_512/index.faiss

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:dcb86f8f32c953c7d5c99662a27e43d6261da7b7b4342bac638e6d19bf7ee530
+size 78975021

data/faiss_1024_512/index.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:313a047fb82ef5c43661a12b2424aeae88688d7631e4bdaf7de283a0b0763dc9
+size 26672894

ingest.py

@@ -6,7 +6,9 @@ from typing import List
 from langchain.document_loaders import PyPDFDirectoryLoader
 from langchain.embeddings import HuggingFaceInstructEmbeddings
 from langchain.text_splitter import RecursiveCharacterTextSplitter
+from langchain.vectorstores.base import VectorStore
 from langchain.vectorstores.chroma import Chroma
+from langchain.vectorstores.faiss import FAISS
 
 from app_modules.utils import *
 
@@ -24,13 +26,23 @@ def split_chunks(documents: List, chunk_size, chunk_overlap) -> List:
     return text_splitter.split_documents(documents)
 
 
-def generate_index(chunks: List, embeddings: HuggingFaceInstructEmbeddings) -> Chroma:
-    chromadb_instructor_embeddings = Chroma.from_documents(
-        documents=chunks, embedding=embeddings, persist_directory=index_path
-    )
+def generate_index(
+    chunks: List, embeddings: HuggingFaceInstructEmbeddings
+) -> VectorStore:
+    if using_faiss:
+        faiss_instructor_embeddings = FAISS.from_documents(
+            documents=chunks, embedding=embeddings
+        )
+
+        faiss_instructor_embeddings.save_local(index_path)
+        return faiss_instructor_embeddings
+    else:
+        chromadb_instructor_embeddings = Chroma.from_documents(
+            documents=chunks, embedding=embeddings, persist_directory=index_path
+        )
 
-    chromadb_instructor_embeddings.persist()
-    return chromadb_instructor_embeddings
+        chromadb_instructor_embeddings.persist()
+        return chromadb_instructor_embeddings
 
 
 # Constants
@@ -40,7 +52,8 @@ device_type, hf_pipeline_device_type = get_device_types()
 hf_embeddings_model_name = (
     os.environ.get("HF_EMBEDDINGS_MODEL_NAME") or "hkunlp/instructor-xl"
 )
-index_path = os.environ.get("CHROMADB_INDEX_PATH")
+index_path = os.environ.get("FAISS_INDEX_PATH") or os.environ.get("CHROMADB_INDEX_PATH")
+using_faiss = os.environ.get("FAISS_INDEX_PATH") is not None
 source_pdfs_path = os.environ.get("SOURCE_PDFS_PATH")
 chunk_size = os.environ.get("CHUNCK_SIZE")
 chunk_overlap = os.environ.get("CHUNK_OVERLAP")
@@ -71,7 +84,11 @@ if not os.path.isdir(index_path):
     index = generate_index(chunks, embeddings)
 else:
     print("The index persist directory is present. Loading index ...")
-    index = Chroma(embedding_function=embeddings, persist_directory=index_path)
+    index = (
+        FAISS.load_local(index_path, embeddings)
+        if using_faiss
+        else Chroma(embedding_function=embeddings, persist_directory=index_path)
+    )
 
 end = timer()