Spaces:
Sleeping
Sleeping
| from utils.settings import configure_settings | |
| from utils.index import load_index | |
| from utils.constant import INDEX_PATH, TOP_K_RETRIEVAL, TOP_N_RERANKER | |
| from llama_index.core import PromptTemplate | |
| from llama_index.retrievers.bm25 import BM25Retriever | |
| from llama_index.core.postprocessor import LLMRerank | |
| from llama_index.core.query_engine import RetrieverQueryEngine | |
| import Stemmer | |
| class QueryEngineManager: | |
| _instance = None | |
| def __new__(cls): | |
| if cls._instance is None: | |
| cls._instance = super(QueryEngineManager, cls).__new__(cls) | |
| cls._instance._initialized = False | |
| return cls._instance | |
| def __init__(self): | |
| if not self._initialized: | |
| self._initialized = True | |
| self.index = None | |
| self.retriever = None | |
| self.reranker = None | |
| self.query_engine = None | |
| self._configure() | |
| def initialize_index(self): | |
| self.index = load_index(path=INDEX_PATH) | |
| self.nodes = list(self.index.docstore.docs.values()) | |
| def _configure(self): | |
| configure_settings() | |
| self.initialize_index() | |
| self.reranker = LLMRerank(top_n=TOP_N_RERANKER) | |
| def get_engine(self, bm25: bool = False, semantic: bool = False): | |
| if bm25: | |
| self.retriever = BM25Retriever.from_defaults( | |
| nodes=self.nodes, | |
| stemmer=Stemmer.Stemmer("english"), | |
| similarity_top_k=TOP_K_RETRIEVAL, | |
| language="english" | |
| ) | |
| elif semantic: | |
| self.retriever = self.index.as_retriever(similarity_top_k=TOP_K_RETRIEVAL) | |
| qa_template = PromptTemplate( | |
| """Given the following context and question, provide a detailed response. | |
| Context: {context_str} | |
| Question: {query_str} | |
| Let me explain this in detail:""", | |
| prompt_type="text_qa" | |
| ) | |
| self.query_engine = RetrieverQueryEngine.from_args( | |
| retriever=self.retriever, | |
| text_qa_template=qa_template, | |
| #node_postprocessors=[self.reranker] | |
| ) | |
| return self.query_engine | |
| def get_engine(bm25: bool = False, semantic: bool = False): | |
| engine_manager = QueryEngineManager() | |
| return engine_manager.get_engine(bm25, semantic), engine_manager |