#perf added hybrid search using bm25 + semantic, minor change to text, splitter, and retrieval hyperparameters

ai_generate.py

@@ -1,5 +1,11 @@
 import gc
 import os
+import time
+import re
+import numpy as np
+import torch
+import bm25s
+
 from langchain_community.document_loaders import PyMuPDFLoader
 from langchain_core.documents import Document
 from langchain_community.embeddings.sentence_transformer import (
@@ -15,14 +21,12 @@ from langchain_anthropic import ChatAnthropic
 from dotenv import load_dotenv
 from langchain_core.output_parsers import XMLOutputParser
 from langchain.prompts import ChatPromptTemplate
-import re
-import numpy as np
-import torch
-import bm25s
 from langchain_community.cross_encoders import HuggingFaceCrossEncoder
 from langchain.retrievers import ContextualCompressionRetriever
 from langchain.retrievers.document_compressors import CrossEncoderReranker
 from langchain_core.messages import HumanMessage
+from langchain.retrievers import EnsembleRetriever
+from langchain_community.retrievers import BM25Retriever
 
 load_dotenv()
 
@@ -33,8 +37,10 @@ os.environ["GLOG_minloglevel"] = "2"
 # RAG parameters
 CHUNK_SIZE = 1024
 CHUNK_OVERLAP = CHUNK_SIZE // 8
-K = 10
+K = 20 # number of chunks to retrieve from semantic search
 FETCH_K = 50
+N_BM25 = 20 # number of chunks to retrieve from keyword search
+TOP_N = 10 # final number of chunks to keep
 
 model_kwargs = {"device": "cuda:1"}
 print("Loading embedding and reranker models...")
@@ -44,7 +50,7 @@ embedding_function = SentenceTransformerEmbeddings(
 # "sentence-transformers/all-MiniLM-L6-v2"
 # "mixedbread-ai/mxbai-embed-large-v1"
 reranker = HuggingFaceCrossEncoder(model_name="BAAI/bge-reranker-base", model_kwargs=model_kwargs)
-compressor = CrossEncoderReranker(model=reranker, top_n=K)
+compressor = CrossEncoderReranker(model=reranker, top_n=TOP_N)
 
 llm_model_translation = {
     "LLaMA 3": "llama3-70b-8192",
@@ -212,7 +218,30 @@ def load_llm(model: str, api_key: str, temperature: float = 1.0, max_length: int
 
 def create_db_with_langchain(path: list[str], url_content: dict, query: str):
     all_docs = []
-    text_splitter = RecursiveCharacterTextSplitter(chunk_size=CHUNK_SIZE, chunk_overlap=CHUNK_OVERLAP)
+
+    text_splitter = RecursiveCharacterTextSplitter(
+        chunk_size=CHUNK_SIZE,
+        chunk_overlap=CHUNK_OVERLAP,
+        separators=[
+            "\n\n",
+            "\n",
+            ".",
+            "\uff0e",  # Fullwidth full stop
+            "\u3002",  # Ideographic full stop
+            "?",
+            "!", 
+            ",",
+            "\uff0c",  # Fullwidth comma
+            "\u3001",  # Ideographic comma
+            " ",
+            "\u200B",  # Zero-width space
+            "",
+        ],
+        keep_separator=True,
+        is_separator_regex=False,
+        length_function=len,
+        add_start_index=False,
+    )
     if path:
         for file in path:
             loader = PyMuPDFLoader(file)
@@ -244,17 +273,19 @@ def create_db_with_langchain(path: list[str], url_content: dict, query: str):
     for idx, doc in enumerate(all_docs):
         print(f"Doc: {idx} | Length = {len(doc.page_content)}")
 
+    bm25_retriever = BM25Retriever.from_documents(all_docs)
+    bm25_retriever.k =  N_BM25
+
     assert len(all_docs) > 0, "No PDFs or scrapped data provided"
     db = Chroma.from_documents(all_docs, embedding_function)
     torch.cuda.empty_cache()
     gc.collect()
-    return db
+    return db, bm25_retriever
 
 
 def pretty_print_docs(docs):
     print(f"\n{'-' * 100}\n".join([f"Document {i+1}:\n\n" + d.page_content for i, d in enumerate(docs)]))
 
-
 def generate_rag(
     prompt: str,
     input_role: str,
@@ -275,12 +306,14 @@ def generate_rag(
 
     query = llm_wrapper(input_role, topic, context, model="OpenAI GPT 4o", task_type="rag", temperature=0.7)
     print("### Query: ", query)
-    db = create_db_with_langchain(path, url_content, query)
+    db, bm25_retriever = create_db_with_langchain(path, url_content, query)
     retriever = db.as_retriever(search_type="mmr", search_kwargs={"k": K, "fetch_k": FETCH_K, "lambda_mult": 0.75})
-
-    # docs = retriever.get_relevant_documents(query)
-    compression_retriever = ContextualCompressionRetriever(base_compressor=compressor, base_retriever=retriever)
+    t0 = time.time()
+    ensemble_retriever = EnsembleRetriever(retrievers=[bm25_retriever, retriever], weights=[0.4, 0.6])
+    compression_retriever = ContextualCompressionRetriever(base_compressor=compressor, base_retriever=ensemble_retriever)
     docs = compression_retriever.invoke(query)
+    t1 = time.time()
+    print(f"Time for retrieval : {t1 - t0:.2f}s")
     print(pretty_print_docs(docs))
 
     formatted_docs = format_docs_xml(docs)