File size: 1,367 Bytes
966108f 36623c8 966108f 36623c8 966108f 36623c8 966108f 36623c8 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 |
from dotenv import load_dotenv
import json
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')
# 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)
semantic_cache_retriever = SemanticCache(compression_retriever)
retrieved_chunks = compression_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: ")
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() |