app.py

import gradio as gr
# from dotenv import load_dotenv
# load_dotenv()

import warnings
warnings.filterwarnings("ignore")

import os, requests, shutil
from collections import defaultdict
from itertools import chain

from langchain.document_loaders import TextLoader
from langchain.text_splitter import CharacterTextSplitter, RecursiveCharacterTextSplitter
from langchain.embeddings import HuggingFaceEmbeddings
from langchain.vectorstores import Chroma  
from langchain.llms import HuggingFaceEndpoint
from langchain.storage import InMemoryStore
from langchain.chains import LLMChain
from langchain.prompts import PromptTemplate
from langchain.retrievers import ParentDocumentRetriever, BM25Retriever
from langchain.retrievers.document_compressors import LLMChainExtractor, LLMChainFilter, EmbeddingsFilter
from langchain_community.document_loaders import PyMuPDFLoader
from langchain.prompts import PromptTemplate


HF_READ_API_KEY = os.environ["HF_READ_API_KEY"]

def get_text(docs):
    return ['Result ' + str(i+1) + '\n' + d.page_content + '\n' for i, d in enumerate(docs)]

def load_pdf(path):
    loader = PyMuPDFLoader(path)
    docs = loader.load()

    return docs, 'PDF loaded successfully'


def multi_query_retrieval(query, llm, retriever):
    DEFAULT_QUERY_PROMPT = PromptTemplate(
        input_variables=["question"],
        template="""You are an AI assistant. Generate 3 different versions of the given question to retrieve relevant docs. 
        Provide these alternative questions separated by newlines. 
        Original question: {question}""",
    )
    mq_llm_chain = LLMChain(llm=llm, prompt=DEFAULT_QUERY_PROMPT)
    
    generated_queries = mq_llm_chain.invoke(query)['text'].split("\n")
    all_queries = [query] + generated_queries
    
    all_retrieved_docs = []
    for q in all_queries:
        retrieved_docs = retriever.get_relevant_documents(q)
        all_retrieved_docs.extend(retrieved_docs)
    
    unique_retrieved_docs = [doc for i, doc in enumerate(all_retrieved_docs) if doc not in all_retrieved_docs[:i]]
    
    return get_text(unique_retrieved_docs)

def compressed_retrieval(query, llm, retriever, extractor_type='chain', embedding_model=None):
    retrieved_docs = retriever.get_relevant_documents(query)
    if extractor_type == 'chain':
        extractor = LLMChainExtractor.from_llm(llm)
    elif extractor_type == 'filter':
        extractor = LLMChainFilter.from_llm(llm)
    elif extractor_type == 'embeddings':
        if embedding_model is None:
            raise ValueError("Embeddings model must be provided for embeddings extractor.")
        extractor = EmbeddingsFilter(embeddings=embedding_model, similarity_threshold=0.5)
    else:
        raise ValueError("Invalid extractor_type. Options are 'chain', 'filter', or 'embeddings'.")
    compressed_docs = extractor.compress_documents(retrieved_docs, query)
    return get_text(compressed_docs)

def unique_by_key(iterable, key_func):
    seen = set()
    for element in iterable:
        key = key_func(element)
        if key not in seen:
            seen.add(key)
            yield element

def ensemble_retrieval(query, retrievers_list, c=60):
    retrieved_docs_by_retriever = [retriever.get_relevant_documents(query) for retriever in retrievers_list]
    weights = [1 / len(retrievers_list)] * len(retrievers_list)
    rrf_score = defaultdict(float)
    for doc_list, weight in zip(retrieved_docs_by_retriever, weights):
        for rank, doc in enumerate(doc_list, start=1):
            rrf_score[doc.page_content] += weight / (rank + c)
            
    all_docs = chain.from_iterable(retrieved_docs_by_retriever)
    sorted_docs = sorted(
        unique_by_key(all_docs, lambda doc: doc.page_content),
        key=lambda doc: rrf_score[doc.page_content],
        reverse=True
    )
    return get_text(sorted_docs)

def long_context_reorder_retrieval(query, retriever):
    retrieved_docs = retriever.get_relevant_documents(query)
    retrieved_docs.reverse() 
    reordered_results = []
    for i, doc in enumerate(retrieved_docs):
        if i % 2 == 1:
            reordered_results.append(doc) 
        else:
            reordered_results.insert(0, doc)
    return get_text(reordered_results)

def process_query(docs, query, embedding_model, inference_model, retrieval_method, chunk_size, chunk_overlap, max_new_tokens, temperature, top_p):

    
    chunking_parameters = {'chunk_size': chunk_size, 'chunk_overlap': chunk_overlap}
    inference_model_params = {'max_new_tokens': max_new_tokens, 'temperature': temperature, 'top_p': top_p}

    text_splitter = RecursiveCharacterTextSplitter(chunk_size=chunking_parameters['chunk_size'], chunk_overlap=chunking_parameters['chunk_overlap'])

    texts = text_splitter.split_documents(docs)

    hf = HuggingFaceEmbeddings(model_name=embedding_model)
    vector_db_from_docs = Chroma.from_documents(texts, hf)
    simple_retriever = vector_db_from_docs.as_retriever(search_kwargs={"k": 5})

    llm_model = HuggingFaceEndpoint(repo_id=inference_model,                     
                            max_new_tokens=inference_model_params['max_new_tokens'], 
                            temperature=inference_model_params['temperature'], 
                            top_p=inference_model_params['top_p'],
                            huggingfacehub_api_token=HF_READ_API_KEY)

    if retrieval_method == "Simple":
        retrieved_docs = simple_retriever.get_relevant_documents(query)
        result = get_text(retrieved_docs)
    elif retrieval_method == "Parent & Child":
        parent_text_splitter = child_text_splitter = text_splitter
        vector_db = Chroma(collection_name="parent_child", embedding_function=hf)
        store = InMemoryStore()
        pr_retriever = ParentDocumentRetriever(
            vectorstore=vector_db,
            docstore=store,
            child_splitter=child_text_splitter,
            parent_splitter=parent_text_splitter,
        )
        pr_retriever.add_documents(docs)
        retrieved_docs = pr_retriever.get_relevant_documents(query)
        result = get_text(retrieved_docs)
    elif retrieval_method == "Multi Query":
        result = multi_query_retrieval(query, llm_model, simple_retriever)
    elif retrieval_method == "Contextual Compression (chain extraction)":
        result = compressed_retrieval(query, llm_model, simple_retriever, extractor_type='chain')
    elif retrieval_method == "Contextual Compression (query filter)":
        result = compressed_retrieval(query, llm_model, simple_retriever, extractor_type='filter')
    elif retrieval_method == "Contextual Compression (embeddings filter)":
        result = compressed_retrieval(query, llm_model, simple_retriever, extractor_type='embeddings', embedding_model=hf)
    elif retrieval_method == "Ensemble":
        bm25_retriever = BM25Retriever.from_documents(docs)
        all_retrievers = [simple_retriever, bm25_retriever]
        result = ensemble_retrieval(query, all_retrievers)
    elif retrieval_method == "Long Context Reorder":
        result = long_context_reorder_retrieval(query, simple_retriever)
    else:
        raise ValueError(f"Unknown retrieval method: {retrieval_method}")
    
    
    prompt_template = PromptTemplate.from_template(
        "Answer the query {query} with the following context:\n {context}. If you cannot use the context to answer the query, say 'I cannot answer the query with the provided context.'"
    )
    
    answer = llm_model.invoke(prompt_template.format(query=query, context=result))

    return "\n".join(result), answer.strip()

embedding_model_list = ['sentence-transformers/all-MiniLM-L6-v2', 'BAAI/bge-small-en-v1.5', 'BAAI/bge-large-en-v1.5'] 
inference_model_list = ['google/gemma-2b-it', 'google/gemma-7b-it', 'microsoft/phi-2', 'mistralai/Mixtral-8x7B-Instruct-v0.1', 'mistralai/Mistral-7B-Instruct-v0.2']
retrieval_method_list = ["Simple", "Parent & Child", "Multi Query", 
                         "Contextual Compression (chain extraction)", "Contextual Compression (query filter)",
                         "Contextual Compression (embeddings filter)", "Ensemble", "Long Context Reorder"]


with gr.Blocks() as demo:
    gr.Markdown("## Compare Retrieval Methods for PDFs")
    with gr.Row():
        with gr.Column():
            pdf_url = gr.Textbox(label="Enter URL to PDF", value="https://www.berkshirehathaway.com/letters/2023ltr.pdf")
            load_button = gr.Button("Load and process PDF")
            status = gr.Textbox(label="Status")
            docs = gr.State()
            load_button.click(load_pdf, inputs=[pdf_url], outputs=[docs, status])
            
            query = gr.Textbox(label="Enter your query", value="What does Warren Buffet think about Coca Cola?")
            with gr.Row():
                embedding_model = gr.Dropdown(embedding_model_list, label="Select Embedding Model", value=embedding_model_list[0])
                inference_model = gr.Dropdown(inference_model_list, label="Select Inference Model", value=inference_model_list[3])
            retrieval_method = gr.Dropdown(retrieval_method_list, label="Select Retrieval Method", value=retrieval_method_list[0])
            
            with gr.Row():
                chunk_size = gr.Number(label="Chunk Size", value=1000)
                chunk_overlap = gr.Number(label="Chunk Overlap", value=200)
            
            with gr.Row():
                max_new_tokens = gr.Number(label="Max New Tokens", value=100)
                temperature = gr.Slider(minimum=0.0, maximum=1.0, step=0.1, label="Temperature", value=0.7)
                top_p = gr.Slider(minimum=0.0, maximum=1.0, step=0.1, label="Top P", value=0.9)
            
            search_button = gr.Button("Retrieval")
        with gr.Column():
            answer = gr.Textbox(label="Answer")
            retrieval_output = gr.Textbox(label="Retrieval Results")
            
    search_button.click(process_query, inputs=[docs, query, embedding_model, inference_model, retrieval_method, chunk_size, chunk_overlap, max_new_tokens, temperature, top_p], outputs=[retrieval_output, answer])

if __name__ == "__main__":
    demo.launch()