Initial commit

.gitattributes

@@ -33,3 +33,5 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+multilingual-e5-large filter=lfs diff=lfs merge=lfs -text
+multilingual-e5-large/* filter=lfs diff=lfs merge=lfs -text

app.py

@@ -0,0 +1,80 @@
+import streamlit as st
+from dotenv import load_dotenv
+import json
+import os, time
+import uuid
+
+from retrieval_pipeline import get_retriever, get_compression_retriever
+import benchmark
+
+
+def get_result(query, compression_retriever):
+    t0 = time.time()
+    retrieved_chunks = compression_retriever.get_relevant_documents(query)
+    latency = time.time() - t0
+    return retrieved_chunks, latency
+ 
+st.set_page_config(
+    layout="wide",
+    page_title="Retrieval Demo"
+)
+
+def setup():
+    load_dotenv()
+    ELASTICSEARCH_URL = os.getenv('ELASTICSEARCH_URL')
+
+    retriever = get_retriever(index='masa.ai', elasticsearch_url=ELASTICSEARCH_URL)
+    compression_retriever = get_compression_retriever(retriever)
+    return compression_retriever
+
+
+def main():
+    st.title("Part 3: Search")
+    # st.sidebar.write("According to the Model Size 👇")
+    # menu = ["Nano", "Small", "Medium", "Large"]
+    # choice = st.sidebar.selectbox("Choose", menu)
+
+    st.sidebar.info("""
+**Model Options:**
+- **Nano**: ~4MB, blazing fast model with competitive performance (ranking precision).
+- **Small**: ~34MB, slightly slower with the best performance (ranking precision).
+- **Medium**: ~110MB, slower model with the best zero-shot performance (ranking precision).
+- **Large**: ~150MB, slower model with competitive performance (ranking precision) for 100+ languages.
+""")
+
+    with st.spinner('Setting up...'):
+        compression_retriever = setup()
+
+    with st.expander("Tech Stack Used"):
+        st.markdown("""
+        **Flash Rank**: Ultra-lite & Super-fast Python library for search & retrieval re-ranking.
+
+        - **Ultra-lite**: No heavy dependencies. Runs on CPU with a tiny ~4MB reranking model.
+        - **Super-fast**: Speed depends on the number of tokens in passages and query, plus model depth.
+        - **Cost-efficient**: Ideal for serverless deployments with low memory and time requirements.
+        - **Based on State-of-the-Art Cross-encoders**: Includes models like ms-marco-TinyBERT-L-2-v2 (default), ms-marco-MiniLM-L-12-v2, rank-T5-flan, and ms-marco-MultiBERT-L-12.
+        - **Sleek Models for Efficiency**: Designed for minimal overhead in user-facing scenarios.
+
+        _Flash Rank is tailored for scenarios requiring efficient and effective reranking, balancing performance with resource usage._
+        """)
+
+
+    with st.form(key='input_form'):
+        query_input = st.text_area("Query Input")
+        # context_input = st.text_area("Context Input")
+        submit_button = st.form_submit_button(label='Retrieve')
+    
+    if submit_button:
+        st.session_state.submitted = True
+
+    if 'submitted' in st.session_state:
+        with st.spinner('Processing...'):
+            result, latency = get_result(query_input, compression_retriever)
+            st.subheader("Please find the retrieved documents below 👇")
+            st.write("latency:", latency, " ms")
+            st.json(result)
+
+
+
+if __name__ == "__main__":
+    main()

benchmark.py

@@ -0,0 +1,47 @@
+import pandas as pd 
+from retrieval_pipeline import get_relevant_documents 
+import tqdm, time
+
+
+TOP_N = 3
+
+def get_benchmark_result(path, retriever):
+    df = pd.read_csv(path)
+    retrieval_result = []
+    query_result = [[] for i in range(TOP_N)] 
+    retrieval_latency = []
+
+    # j = 0
+    for i, row in tqdm.tqdm(df.iterrows()):
+        # j+=1
+        query = row['query']
+        target = row['body']
+
+
+        t0 = time.time()
+        results = retriever.get_relevant_documents(query)
+        t = time.time() - t0
+        retrieval_latency.append(t)
+
+        result_content = [result.page_content for result in results]
+        # results_content = get_relevant_documents(query, retriever, top_k=5)
+
+        for i, text in enumerate(result_content):
+            query_result[i] = text 
+
+        if target in result_content:
+            retrieval_result.append("Success")
+        else:
+            retrieval_result.append("Failed")
+        # if j>20:
+        #     break
+
+    df["retrieval_result"] = retrieval_result 
+    df["retrieval_latency"] = retrieval_latency 
+    for i in range(TOP_N):
+        df[f'q{i+1}'] = query_result[i]
+    df.to_csv('benchmark_result q3 topk 5.csv')
+    print(df['retrieval_result'].value_counts())
+    print(df['retrieval_result'].value_counts()/ len(df))
+
+

main.py

@@ -0,0 +1,33 @@
+from dotenv import load_dotenv
+import json
+import os
+import uuid
+
+from retrieval_pipeline import get_retriever, get_compression_retriever
+import benchmark
+
+load_dotenv()
+ELASTICSEARCH_URL = os.getenv('ELASTICSEARCH_URL')
+# HUGGINGFACE_KEY = os.getenv('HUGGINGFACE_KEY')
+
+os.environ["ES_ENDPOINT"] = ELASTICSEARCH_URL
+print(ELASTICSEARCH_URL)
+
+if __name__ == "__main__":
+    retriever = get_retriever(index='masa.ai', elasticsearch_url=ELASTICSEARCH_URL)
+    compression_retriever = get_compression_retriever(retriever)
+    retrieved_chunks = compression_retriever.get_relevant_documents('Gunung Semeru')
+    print(retrieved_chunks)
+    
+    # retrieved_chunks = retriever.get_relevant_documents('Gunung Semeru')
+    # print(retrieved_chunks)
+
+    benchmark.get_benchmark_result("benchmark-reranker.csv", retriever=compression_retriever)
+
+    # for i in range(100):
+    #     query = input("query: ")
+    #     retrieved_chunks = retriever.get_relevant_documents(query)
+    #     print("Result:")
+    #     for r in retrieved_chunks:
+    #         print(r.page_content[:50])
+    #     print()

retrieval_pipeline/__init__.py

@@ -0,0 +1,2 @@
+from retrieval_pipeline.main import *
+from retrieval_pipeline.hybrid_search import *

retrieval_pipeline/hybrid_search.py

@@ -0,0 +1,95 @@
+from langchain_core.vectorstores import VectorStoreRetriever
+from langchain_community.retrievers import ElasticSearchBM25Retriever
+from langchain_community.vectorstores.elastic_vector_search import ElasticVectorSearch
+from langchain_core.callbacks.manager import AsyncCallbackManagerForRetrieverRun
+from langchain_core.retrievers import BaseRetriever
+from langchain_core.documents import Document
+import elasticsearch
+
+
+from typing import Optional, List
+
+
+class HybridRetriever(BaseRetriever):
+    dense_db: ElasticVectorSearch
+    dense_retriever: VectorStoreRetriever
+    sparse_retriever: ElasticSearchBM25Retriever
+    index_dense: str
+    index_sparse: str
+    top_k_dense: int
+    top_k_sparse: int
+
+    is_training: bool = False
+    
+    @classmethod
+    def create(
+        cls, dense_db, dense_retriever, sparse_retriever, index_dense, index_sparse, top_k_dense, top_k_sparse
+        ):
+
+        return cls(
+                dense_db=dense_db,
+                dense_retriever=dense_retriever,
+                sparse_retriever=sparse_retriever,
+                index_dense=index_dense,
+                index_sparse=index_sparse,
+                top_k_dense=top_k_dense,
+                top_k_sparse=top_k_sparse,
+            )
+
+    def reset_indices(self):
+        result = self.dense_db.client.indices.delete(
+            index=self.index_dense,
+            ignore_unavailable=True,
+            allow_no_indices=True,
+        )
+
+        
+        logging.info('dense_db delete:', result)
+
+        result = self.sparse_retriever.client.indices.delete(
+            index=self.index_sparse,
+            ignore_unavailable=True,
+            allow_no_indices=True,
+        )
+
+        logging.info('sparse_retriever delete:', result)
+
+    def add_documents(self, documents, batch_size=25):
+        for i in range(0, len(documents), batch_size):
+            print('batch', i) 
+            dense_batch = documents[i:i + batch_size]
+            sparse_batch = [doc.page_content for doc in dense_batch]
+            self.dense_retriever.add_documents(dense_batch)
+            self.sparse_retriever.add_texts(sparse_batch)
+        
+    def _get_relevant_documents(self, query: str, **kwargs):
+        dense_results = self.dense_retriever.get_relevant_documents(query)[:self.top_k_dense]
+        sparse_results = self.sparse_retriever.get_relevant_documents(query)[:self.top_k_sparse]
+
+        # Combine results (you'll need a strategy here) 
+        combined_results = dense_results + sparse_results
+        # result_text = [doc.page_content for doc in combined_results]
+
+        # reranked_result = rerank.rerank(query, documents=result_text, model="rerank-lite-1", top_k=self.top_k_dense+self.top_k_sparse)
+        # reranked_result = sorted(reranked_result.results, key=lambda result: result.index)
+
+        # Create LangChain Documents
+        documents = [Document(page_content=doc.page_content, metadata=doc.metadata) for doc in combined_results]
+        # documents = [Document(page_content=doc.page_content, metadata=doc.metadata, relevance_score=result.relevance_score) for result, doc in zip(reranked_result, combined_results)]
+        return documents
+    
+    async def aget_relevant_documents(self, query: str): 
+        raise NotImplementedError 
+
+def get_dense_db(elasticsearch_url, index_dense, embeddings):
+    dense_db = ElasticVectorSearch(
+        elasticsearch_url=elasticsearch_url,
+        index_name=index_dense,
+        embedding=embeddings,
+    )
+    return dense_db
+    
+def get_sparse_retriever(elasticsearch_url, index_sparse):
+    sparse_retriever = ElasticSearchBM25Retriever(client=elasticsearch.Elasticsearch(elasticsearch_url), 
+                                                  index_name=index_sparse)
+    return sparse_retriever

retrieval_pipeline/main.py

@@ -0,0 +1,98 @@
+from langchain.vectorstores import ElasticVectorSearch
+from langchain.llms import OpenAI, HuggingFaceHub
+from langchain_community.embeddings import HuggingFaceEmbeddings
+from retrieval_pipeline.hybrid_search import HybridRetriever, get_dense_db, get_sparse_retriever
+from retrieval_pipeline.utils import get_hybrid_indexes
+
+from langchain.retrievers import ContextualCompressionRetriever
+from langchain.retrievers.document_compressors import CrossEncoderReranker
+from langchain_community.cross_encoders import HuggingFaceCrossEncoder
+
+import logging
+import tqdm
+
+
+def get_compression_retriever(retriever):
+    model = HuggingFaceCrossEncoder(model_name="BAAI/bge-reranker-base")
+    compressor = CrossEncoderReranker(model=model, top_n=3)
+    compression_retriever = ContextualCompressionRetriever(
+        base_compressor=compressor, base_retriever=retriever
+    )
+    return compression_retriever
+
+# Embedding Models Loader
+def get_huggingface_embeddings(model_name):
+    logging.info(f"Loading Huggingface Embedding")
+    embeddings = HuggingFaceEmbeddings(model_name=model_name)
+    return embeddings
+
+def get_vectorstore(index_name, embeddings, elasticsearch_url=None):
+    logging.info(f"Loading vectorstore")
+
+    index_dense, index_sparse = get_hybrid_indexes(index_name)
+
+    dense_db = get_dense_db(elasticsearch_url, index_dense, embeddings)
+    dense_retriever = dense_db.as_retriever()
+
+    sparse_retriever = get_sparse_retriever(elasticsearch_url, index_sparse)
+
+    hybrid_retriever = HybridRetriever(
+        dense_db=dense_db,
+        dense_retriever=dense_retriever,
+        sparse_retriever=sparse_retriever,
+        index_dense=index_dense,
+        index_sparse=index_sparse,
+        top_k_dense=2,
+        top_k_sparse=3
+    )
+
+    # db = ElasticVectorSearch(
+    #     elasticsearch_url=elasticsearch_url,
+    #     index_name=index_name,
+    #     embedding=embeddings,
+    # )
+    return hybrid_retriever
+
+def get_retriever(index, elasticsearch_url):
+    # cache.init(pre_embedding_func=get_msg_func)
+    # cache.set_openai_key(openai_api_key)
+
+    embeddings = get_huggingface_embeddings(model_name="multilingual-e5-large")
+
+    # llm = get_openai_llm(
+    #     model_name=model_name, temperature=0, api_key=model_api_key
+    # )
+    # embeddings = get_openai_embeddings(embedding_api_key, embedding_name)
+
+    # question_generator = load_question_generator(llm)
+    # answer_generator = load_answer_generator(llm, company=model_config['company_name'], tone=model_config['tone'], additional_instructions=model_config['additional_instructions'])
+
+    retriever = get_vectorstore(
+        index,
+        embeddings=embeddings,
+        elasticsearch_url=elasticsearch_url,
+    )
+
+    # if history:
+    #     qa = get_conversational_chain(retriever, question_generator, answer_generator)
+    # else:
+    #     qa = get_retrieval_chain(retriever, answer_generator)
+
+    # chain = CustomLLMChain(
+    #     chain=qa,
+    #     model_name=llm.model_name,
+    #     use_history=history
+    # )
+    # 
+    # 
+    return retriever
+
+def get_relevant_documents(query, retriever, top_k):
+    results = retriever.get_relevant_documents(query)
+    passages = [{
+        "id": i,
+        "text": result.page_content
+    } for i, result in enumerate(results)]
+
+    reranked_result = ranker.rerank(RerankRequest(query=query, passages=passages))
+    return reranked_result 

retrieval_pipeline/utils.py

@@ -0,0 +1,5 @@
+def get_hybrid_indexes(index_name):
+    index_dense = f'{index_name}-dense'
+    index_sparse = f'{index_name}-sparse'
+
+    return index_dense, index_sparse