Improve performance with contextual compression, a technique where retrieved documents are compressed, and irrelevant information is filtered out.

app.py

@@ -1,12 +1,11 @@
 import streamlit as st
-from langchain.chains import ConversationalRetrievalChain
 from langchain.memory import ConversationBufferMemory
 from langchain.memory.chat_message_histories import StreamlitChatMessageHistory
 from langchain_community.chat_models import ChatOpenAI
-
 from calback_handler import PrintRetrievalHandler, StreamHandler
 from chat_profile import ChatProfileRoleEnum
 from document_retriever import configure_retriever
+from langchain.chains import ConversationalRetrievalChain
 
 st.set_page_config(
     page_title="InkChatGPT: Chat with Documents",
@@ -79,6 +78,7 @@ with chat_tab:
             retriever=result_retriever,
             memory=memory,
             verbose=False,
+            max_tokens_limit=4000,
         )
 
         avatars = {

document_retriever.py

@@ -2,19 +2,16 @@ import os
 import tempfile
 
 import streamlit as st
-from langchain_community.document_loaders import (
-    Docx2txtLoader,
-    PyPDFLoader,
-    TextLoader,
-    UnstructuredEPubLoader,
-)
+from langchain.retrievers import ContextualCompressionRetriever
+from langchain.retrievers.document_compressors import EmbeddingsFilter
+from langchain_community.document_loaders import Docx2txtLoader, PyPDFLoader, TextLoader
 from langchain_community.embeddings import HuggingFaceEmbeddings
 from langchain_community.vectorstores import DocArrayInMemorySearch
 from langchain_text_splitters import RecursiveCharacterTextSplitter
 
 
 @st.cache_resource(ttl="1h")
-def configure_retriever(files):
+def configure_retriever(files, use_compression=False):
     # Read documents
     docs = []
     temp_dir = tempfile.TemporaryDirectory()
@@ -32,8 +29,6 @@ def configure_retriever(files):
             loader = Docx2txtLoader(temp_filepath)
         elif extension == ".txt":
             loader = TextLoader(temp_filepath)
-        elif extension == ".epub":
-            loader = UnstructuredEPubLoader(temp_filepath)
         else:
             st.write("This document format is not supported!")
             return None
@@ -45,7 +40,7 @@ def configure_retriever(files):
     splits = text_splitter.split_documents(docs)
 
     # Create embeddings and store in vectordb
-    embeddings = HuggingFaceEmbeddings(model_name="all-mpnet-base-v2")
+    embeddings = HuggingFaceEmbeddings(model_name="all-MiniLM-L6-v2")
     vectordb = DocArrayInMemorySearch.from_documents(splits, embeddings)
 
     # Define retriever
@@ -53,4 +48,13 @@ def configure_retriever(files):
         search_type="mmr", search_kwargs={"k": 2, "fetch_k": 4}
     )
 
-    return retriever
+    if not use_compression:
+        return retriever
+
+    embeddings_filter = EmbeddingsFilter(
+        embeddings=embeddings, similarity_threshold=0.76
+    )
+
+    return ContextualCompressionRetriever(
+        base_compressor=embeddings_filter, base_retriever=retriever
+    )

requirements.txt

@@ -7,4 +7,5 @@ streamlit_chat
 streamlit-extras
 pypdf
 docx2txt
-unstructured
+unstructured
+tiktoken