general refactoring

rag_app/knowledge_base/build_vector_store.py

@@ -1,18 +1,13 @@
 # vectorization functions
 from langchain_community.vectorstores import FAISS
 from langchain_community.vectorstores import Chroma
-#from langchain_community.document_loaders import DirectoryLoader
-#from langchain_text_splitters import RecursiveCharacterTextSplitter
-#from langchain_community.embeddings.sentence_transformer import (
-#    SentenceTransformerEmbeddings,
-#)
-#from langchain_huggingface import HuggingFaceEmbeddings
 from langchain_community.retrievers import BM25Retriever
+
 from rag_app.knowledge_base.create_embedding import create_embeddings
 from rag_app.utils.generate_summary import generate_description, generate_keywords
+
 import time
 import os
-#from dotenv import load_dotenv
 
 def build_vector_store(
         docs: list, 
@@ -32,11 +27,6 @@ def build_vector_store(
         FAISS_INDEX_PATH = db_path
 
     embeddings,chunks = create_embeddings(docs, chunk_size, chunk_overlap, embedding_model)
-    # for chunk in chunks:
-    #     keywords=generate_keywords(chunk)
-    #     description=generate_description(chunk)
-    #     chunk.metadata['keywords']=keywords
-    #     chunk.metadata['description']=description
 
     #load chunks into vector store
     print(f'Loading chunks into faiss vector store ...')

rag_app/knowledge_base/create_embedding.py

@@ -1,21 +1,16 @@
 # embeddings functions
-#from langchain_community.vectorstores import FAISS
-#from langchain_community.document_loaders import ReadTheDocsLoader
-#from langchain_community.vectorstores.utils import filter_complex_metadata
 from langchain_text_splitters import RecursiveCharacterTextSplitter
-# from langchain_huggingface import HuggingFaceEmbeddings
 from langchain_community.embeddings.sentence_transformer import (
     SentenceTransformerEmbeddings,
 )
 import time
 from langchain_core.documents import Document
-
+from config import EMBEDDING_MODEL
 
 def create_embeddings(
         docs: list[Document], 
         chunk_size:int = 500, 
         chunk_overlap:int = 50,
-        embedding_model: str = "sentence-transformers/multi-qa-mpnet-base-dot-v1", 
         ):
     """given a sequence of `Document` objects this fucntion will
     generate embeddings for it.
@@ -47,8 +42,7 @@ def create_embeddings(
     print(f'Time taken to chunk {len(docs)} documents: {et} seconds.')
 
     #Stage two: embed the docs.
-    #embeddings = HuggingFaceEmbeddings(model_name=embedding_model)
-    embeddings = SentenceTransformerEmbeddings(model_name=embedding_model)
+    embeddings = SentenceTransformerEmbeddings(model_name=EMBEDDING_MODEL)
     print(f"created a total of {len(chunks)} chunks")
 
     return embeddings,chunks

rag_app/knowledge_base/reranking.py

@@ -1,4 +1,3 @@
-# from get_db_retriever import get_db_retriever
 from pathlib import Path
 from langchain_community.vectorstores import FAISS
 from dotenv import load_dotenv
@@ -11,29 +10,36 @@ from langchain_community.vectorstores import Chroma
 load_dotenv()
 
 
-def get_reranked_docs_faiss(query:str, 
-                      path_to_db:str, 
-                      embedding_model:str,
-                      hf_api_key:str, 
-                      num_docs:int=5) -> list:
+def get_reranked_docs_faiss(
+    query:str, 
+    path_to_db:str, 
+    embedding_model:str,
+    hf_api_key:str, 
+    num_docs:int=5
+    ) -> list:
     """ Re-ranks the similarity search results and returns top-k highest ranked docs
 
-        Args:
-            query (str): The search query
-            path_to_db (str): Path to the vectorstore database
-            embedding_model (str): Embedding model used in the vector store
-            num_docs (int): Number of documents to return
-        
-        Returns: A list of documents with the highest rank
+    Args:
+        query (str): The search query
+        path_to_db (str): Path to the vectorstore database
+        embedding_model (str): Embedding model used in the vector store
+        num_docs (int): Number of documents to return
+    
+    Returns: A list of documents with the highest rank
     """
     assert num_docs <= 10, "num_docs should be less than similarity search results"
     
-    embeddings = HuggingFaceInferenceAPIEmbeddings(api_key=hf_api_key,
-                                                   model_name=embedding_model)
+    embeddings = HuggingFaceInferenceAPIEmbeddings(
+        api_key=hf_api_key,
+        model_name=embedding_model
+        )
+    
     # Load the vectorstore database
-    db = FAISS.load_local(folder_path=path_to_db,
-                          embeddings=embeddings,
-                          allow_dangerous_deserialization=True)
+    db = FAISS.load_local(
+        folder_path=path_to_db,
+        embeddings=embeddings,
+        allow_dangerous_deserialization=True
+        )
     
     # Get 10 documents based on similarity search
     docs =  db.similarity_search(query=query, k=10)