Add cache

benchmark.py

@@ -7,6 +7,7 @@ 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 = []
@@ -21,13 +22,13 @@ def get_benchmark_result(path, retriever):
         t0 = time.time()
         results = retriever.get_relevant_documents(query)
         t = time.time() - t0
-        retrieval_latency.append(t)
+        retrieval_latency.append(str(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 
+            query_result[i].append(text) 
 
         if target in result_content:
             retrieval_result.append("Success")
@@ -37,10 +38,10 @@ def get_benchmark_result(path, retriever):
         #     break
 
     df["retrieval_result"] = retrieval_result 
-    df["retrieval_latency"] = retrieval_latency 
+    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')
+        df.to_csv('benchmark_result.csv')
     print(df['retrieval_result'].value_counts())
     print(df['retrieval_result'].value_counts()/ len(df))
 

main.py

@@ -1,10 +1,13 @@
 from dotenv import load_dotenv
 import json
-import os
+import os, time
 import uuid
 
 from retrieval_pipeline import get_retriever, get_compression_retriever
 import benchmark
+from retrieval_pipeline.hybrid_search import store
+
+from retrieval_pipeline.cache import SemanticCache
 
 load_dotenv()
 ELASTICSEARCH_URL = os.getenv('ELASTICSEARCH_URL')
@@ -16,18 +19,26 @@ print(ELASTICSEARCH_URL)
 if __name__ == "__main__":
     retriever = get_retriever(index='masa.ai', elasticsearch_url=ELASTICSEARCH_URL)
     compression_retriever = get_compression_retriever(retriever)
+
+    semantic_cache_retriever = SemanticCache(compression_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()
+
+    # benchmark.get_benchmark_result("benchmark-reranker.csv", retriever=compression_retriever)
+
+    for i in range(100):
+        query = input("query: ")
+        t0 = time.time()
+        # retrieved_chunks = compression_retriever.get_relevant_documents(query)
+        retrieved_chunks = semantic_cache_retriever.get_relevant_documents(query)
+
+        t = time.time() - t0
+
+        print(list(store.yield_keys()))
+        print('time:', t)
+
+        print("Result:")
+        for r in retrieved_chunks:
+            print(r.page_content[:50])
+        print()

requirements old.txt

@@ -0,0 +1,91 @@
+aiohttp==3.9.5
+aiolimiter==1.1.0
+aiosignal==1.3.1
+altair==5.3.0
+annotated-types==0.6.0
+async-timeout==4.0.3
+attrs==23.2.0
+blinker==1.8.2
+cachetools==5.3.3
+certifi==2024.2.2
+charset-normalizer==3.3.2
+click==8.1.7
+colorama==0.4.6
+Cython==3.0.10
+dataclasses-json==0.6.6
+elastic-transport==8.13.0
+elasticsearch==8.13.1
+filelock==3.14.0
+frozenlist==1.4.1
+fsspec==2024.3.1
+gitdb==4.0.11
+GitPython==3.1.43
+greenlet==3.0.3
+huggingface-hub==0.23.0
+idna==3.7
+intel-openmp==2021.4.0
+Jinja2==3.1.4
+joblib==1.4.2
+jsonpatch==1.33
+jsonpointer==2.4
+jsonschema==4.22.0
+jsonschema-specifications==2023.12.1
+langchain==0.1.20
+langchain-community==0.0.38
+langchain-core==0.1.52
+langchain-text-splitters==0.0.1
+langsmith==0.1.57
+markdown-it-py==3.0.0
+MarkupSafe==2.1.5
+marshmallow==3.21.2
+mdurl==0.1.2
+mkl==2021.4.0
+mpmath==1.3.0
+multidict==6.0.5
+mypy-extensions==1.0.0
+networkx==3.2.1
+numpy==1.26.4
+orjson==3.10.3
+packaging==23.2
+pandas==2.2.2
+pillow==10.3.0
+protobuf==4.25.3
+pyarrow==16.1.0
+pydantic==2.7.1
+pydantic_core==2.18.2
+pydeck==0.9.1
+Pygments==2.18.0
+python-dateutil==2.9.0.post0
+python-dotenv==1.0.1
+pytz==2024.1
+PyYAML==6.0.1
+referencing==0.35.1
+regex==2024.5.10
+requests==2.31.0
+rich==13.7.1
+rpds-py==0.18.1
+safetensors==0.4.3
+scikit-learn==1.4.2
+scipy==1.13.0
+sentence-transformers==2.7.0
+six==1.16.0
+smmap==5.0.1
+SQLAlchemy==2.0.30
+streamlit==1.34.0
+sympy==1.12
+tbb==2021.12.0
+tenacity==8.3.0
+threadpoolctl==3.5.0
+tokenizers==0.19.1
+toml==0.10.2
+toolz==0.12.1
+torch==2.3.0
+tornado==6.4
+tqdm==4.66.4
+transformers==4.40.2
+typing-inspect==0.9.0
+typing_extensions==4.11.0
+tzdata==2024.1
+urllib3==2.2.1
+watchdog==4.0.0
+yarl==1.9.4

requirements.txt

@@ -3,6 +3,7 @@ aiolimiter==1.1.0
 aiosignal==1.3.1
 altair==5.3.0
 annotated-types==0.6.0
+anyio==4.3.0
 async-timeout==4.0.3
 attrs==23.2.0
 blinker==1.8.2
@@ -11,17 +12,31 @@ certifi==2024.2.2
 charset-normalizer==3.3.2
 click==8.1.7
 colorama==0.4.6
+coloredlogs==15.0.1
 Cython==3.0.10
 dataclasses-json==0.6.6
 elastic-transport==8.13.0
 elasticsearch==8.13.1
+exceptiongroup==1.2.1
+faiss-cpu==1.8.0
+fastembed==0.2.6
 filelock==3.14.0
+flatbuffers==24.3.25
 frozenlist==1.4.1
 fsspec==2024.3.1
 gitdb==4.0.11
 GitPython==3.1.43
 greenlet==3.0.3
-huggingface-hub==0.23.0
+grpcio==1.63.0
+grpcio-tools==1.63.0
+h11==0.14.0
+h2==4.1.0
+hpack==4.0.0
+httpcore==1.0.5
+httpx==0.27.0
+huggingface-hub==0.20.3
+humanfriendly==10.0
+hyperframe==6.0.1
 idna==3.7
 intel-openmp==2021.4.0
 Jinja2==3.1.4
@@ -35,6 +50,8 @@ langchain-community==0.0.38
 langchain-core==0.1.52
 langchain-text-splitters==0.0.1
 langsmith==0.1.57
+llvmlite==0.42.0
+loguru==0.7.2
 markdown-it-py==3.0.0
 MarkupSafe==2.1.5
 marshmallow==3.21.2
@@ -44,21 +61,29 @@ mpmath==1.3.0
 multidict==6.0.5
 mypy-extensions==1.0.0
 networkx==3.2.1
+numba==0.59.1
 numpy==1.26.4
+onnx==1.16.0
+onnxruntime==1.17.3
 orjson==3.10.3
 packaging==23.2
 pandas==2.2.2
 pillow==10.3.0
-protobuf==4.25.3
+portalocker==2.8.2
+protobuf==5.26.1
 pyarrow==16.1.0
 pydantic==2.7.1
 pydantic_core==2.18.2
 pydeck==0.9.1
 Pygments==2.18.0
+pyreadline3==3.4.1
 python-dateutil==2.9.0.post0
 python-dotenv==1.0.1
 pytz==2024.1
+pywin32==306
 PyYAML==6.0.1
+qdrant-client==1.9.1
+rankerEval==0.2.0
 referencing==0.35.1
 regex==2024.5.10
 requests==2.31.0
@@ -67,9 +92,11 @@ rpds-py==0.18.1
 safetensors==0.4.3
 scikit-learn==1.4.2
 scipy==1.13.0
+semantic-cache==0.1.1
 sentence-transformers==2.7.0
 six==1.16.0
 smmap==5.0.1
+sniffio==1.3.1
 SQLAlchemy==2.0.30
 streamlit==1.34.0
 sympy==1.12
@@ -88,4 +115,5 @@ typing_extensions==4.11.0
 tzdata==2024.1
 urllib3==2.2.1
 watchdog==4.0.0
+win32-setctime==1.1.0
 yarl==1.9.4

retrieval_pipeline/cache.py

@@ -0,0 +1,94 @@
+import faiss
+from sentence_transformers import SentenceTransformer
+import time
+import json
+
+from langchain_core.documents import Document
+
+def init_cache():
+    index = faiss.IndexFlatL2(1024)
+    if index.is_trained:
+        print("Index trained")
+
+    # Initialize Sentence Transformer model
+    encoder = SentenceTransformer("multilingual-e5-large")
+
+    return index, encoder
+
+def retrieve_cache(json_file):
+    try:
+        with open(json_file, "r") as file:
+            cache = json.load(file)
+    except FileNotFoundError:
+        cache = {"query": [], "embeddings": [], "answers": [], "response_text": []}
+
+    return cache
+
+def store_cache(json_file, cache):
+    with open(json_file, "w") as file:
+        json.dump(cache, file)
+
+class SemanticCache:
+    def __init__(self, retriever, json_file="cache_file.json", thresold=0.35):
+        # Initialize Faiss index with Euclidean distance
+        self.retriever = retriever
+        self.index, self.encoder = init_cache()
+
+        # Set Euclidean distance threshold
+        # a distance of 0 means identicals sentences
+        # We only return from cache sentences under this thresold
+        self.euclidean_threshold = thresold
+
+        self.json_file = json_file
+        self.cache = retrieve_cache(self.json_file)
+
+    def query_database(self, query_text):
+        results = self.retriever.get_relevant_documents(query_text)
+        return results
+
+    def get_relevant_documents(self, query: str) -> str:
+        # Method to retrieve an answer from the cache or generate a new one
+        start_time = time.time()
+    # try:
+        # First we obtain the embeddings corresponding to the user query
+        embedding = self.encoder.encode([query])
+
+        # Search for the nearest neighbor in the index
+        self.index.nprobe = 8
+        D, I = self.index.search(embedding, 1)
+
+        if D[0] >= 0:
+            if I[0][0] >= 0 and D[0][0] <= self.euclidean_threshold:
+                row_id = int(I[0][0])
+
+                print("Answer recovered from Cache. ")
+                print(f"{D[0][0]:.3f} smaller than {self.euclidean_threshold}")
+                print(f"Found cache in row: {row_id} with score {D[0][0]:.3f}")
+
+                end_time = time.time()
+                elapsed_time = end_time - start_time
+                print(f"Time taken: {elapsed_time:.3f} seconds")
+                return [Document(**doc[k]) for doc in self.cache["answers"][row_id]]
+
+        # Handle the case when there are not enough results
+        # or Euclidean distance is not met, asking to chromaDB.
+        answer = self.query_database(query)
+        # response_text = answer["documents"][0][0]
+
+        self.cache["query"].append(query)
+        self.cache["embeddings"].append(embedding[0].tolist())
+        self.cache["answers"].append([doc.__dict__ for doc in answer])
+        # self.cache["response_text"].append(response_text)
+
+        print("Answer recovered from ChromaDB. ")
+        # print(f"response_text: {response_text}")
+
+        self.index.add(embedding)
+        store_cache(self.json_file, self.cache)
+        end_time = time.time()
+        elapsed_time = end_time - start_time
+        print(f"Time taken: {elapsed_time:.3f} seconds")
+
+        return answer 
+        # except Exception as e:
+        #     raise RuntimeError(f"Error during 'get_relevant_documents' method: {e}")

retrieval_pipeline/hybrid_search.py

@@ -9,6 +9,10 @@ import elasticsearch
 
 from typing import Optional, List
 
+from langchain.storage import LocalFileStore
+from langchain.embeddings import CacheBackedEmbeddings
+
+store = LocalFileStore("cache")
 
 class HybridRetriever(BaseRetriever):
     dense_db: ElasticVectorSearch
@@ -68,10 +72,6 @@ class HybridRetriever(BaseRetriever):
 
         # 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]
@@ -82,10 +82,21 @@ class HybridRetriever(BaseRetriever):
         raise NotImplementedError 
 
 def get_dense_db(elasticsearch_url, index_dense, embeddings):
+    # retriever cache
+    cached_embedder = CacheBackedEmbeddings.from_bytes_store(
+        embeddings, store, 
+        namespace='sentence-transformer',
+        # query_embedding_store=store,
+        # query_embedding_cache=True
+    )
+
+    cached_embedder.query_embedding_store = store
+
     dense_db = ElasticVectorSearch(
         elasticsearch_url=elasticsearch_url,
         index_name=index_dense,
         embedding=embeddings,
+        # embedding=cached_embedder,
     )
     return dense_db