Update app.py

app.py

@@ -3,11 +3,8 @@ import multiprocessing
 import concurrent.futures
 from langchain.document_loaders import TextLoader, DirectoryLoader
 from langchain.text_splitter import RecursiveCharacterTextSplitter
-from langchain.vectorstores import FAISS
-from sentence_transformers import SentenceTransformer
-import faiss
-import torch
-import numpy as np
+from langchain_community.vectorstores import FAISS
+from langchain_community.embeddings import HuggingFaceEmbeddings
 from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer, BitsAndBytesConfig
 from datetime import datetime
 import json
@@ -24,7 +21,7 @@ logger = logging.getLogger(__name__)
 class DocumentRetrievalAndGeneration:
     def __init__(self, embedding_model_name, lm_model_id, data_folder):
         self.all_splits = self.load_documents(data_folder)
-        self.embeddings = SentenceTransformer(embedding_model_name)
+        self.embeddings = HuggingFaceEmbeddings(model_name=embedding_model_name)
         self.vectordb = self.create_faiss_index()
         self.tokenizer, self.model = self.initialize_llm(lm_model_id)
         self.retriever_tool = self.create_retriever_tool()
@@ -40,24 +37,7 @@ class DocumentRetrievalAndGeneration:
         return all_splits
 
     def create_faiss_index(self):
-        all_texts = [split.page_content for split in self.all_splits]
-        embeddings = self.embeddings.encode(all_texts)
-        
-        # Create FAISS index
-        vector_dimension = embeddings.shape[1]
-        index = faiss.IndexFlatL2(vector_dimension)
-        index.add(embeddings)
-        
-        # Create docstore
-        docstore = {i: doc for i, doc in enumerate(self.all_splits)}
-        
-        # Create and return FAISS object
-        return FAISS(
-            embedding_function=self.embeddings.encode,
-            index=index,
-            docstore=docstore,
-            index_to_docstore_id={i: i for i in range(len(self.all_splits))}
-        )
+        return FAISS.from_documents(self.all_splits, self.embeddings)
 
     def initialize_llm(self, model_id):
         quantization_config = BitsAndBytesConfig(
@@ -145,6 +125,12 @@ Question:
         response = self.query_and_generate_response(query)
         return response
 
+    def save_index(self, path):
+        self.vectordb.save_local(path)
+
+    def load_index(self, path):
+        self.vectordb = FAISS.load_local(path, self.embeddings)
+
 if __name__ == "__main__":
     embedding_model_name = 'thenlper/gte-small'
     lm_model_id = "meta-llama/Meta-Llama-3.1-8B-Instruct"
@@ -152,6 +138,9 @@ if __name__ == "__main__":
 
     doc_retrieval_gen = DocumentRetrievalAndGeneration(embedding_model_name, lm_model_id, data_folder)
 
+    # Save the index for future use
+    doc_retrieval_gen.save_index("faiss_index")
+
     def launch_interface():
         css_code = """
             .gradio-container {