Upload 5 files

ChatEngine.py

@@ -0,0 +1,55 @@
+from llama_index.llms.huggingface import HuggingFaceLLM
+from llama_index.core.base.llms.types import ChatMessage, MessageRole
+from configs import MODEL_NAME, CONTEXT_WINDOW, TEMPERATURE, SYSTEM_PROMPT, DEVICE
+
+class ChatEngine:
+    def __init__(self, retriever, model_name=MODEL_NAME, context_window=CONTEXT_WINDOW, temperature=TEMPERATURE):
+        """
+        Initializes the ChatEngine with a retriever and a language model.
+
+        Args:
+            retriever (HybridRetriever): An instance of a retriever to fetch relevant documents.
+            model_name (str): The name of the language model to be used.
+            context_window (int, optional): The maximum context window size for the language model. Defaults to 32000.
+            temperature (float, optional): The temperature setting for the language model. Defaults to 0.
+        """
+
+        self.retriever = retriever
+
+        self.llm = HuggingFaceLLM(model_name=model_name, 
+                                tokenizer_name=model_name, 
+                                system_prompt=SYSTEM_PROMPT, 
+                                context_window=context_window, 
+                                generate_kwargs={"temperature": temperature},
+                                device_map=DEVICE)
+        self.chat_history = []
+
+    def ask_question(self, question):
+        """
+        Asks a question to the language model, using the retriever to fetch relevant documents.
+
+        Args:
+            question (str): The question to be asked.
+
+        Returns:
+            str: The response from the language model in markdown format.
+        """
+        
+        question = "[INST]" + question + "[/INST]"
+        
+        results = self.retriever.best_docs(question)
+        document = [doc.text for doc, sc in results]
+
+        self.chat_history.append(ChatMessage(role=MessageRole.USER, content=f"Question: {question}"))
+        
+        self.chat_history.append(ChatMessage(role=MessageRole.SYSTEM, content=f"Document: {document}"))
+
+        response = self.llm.chat(self.chat_history)
+        response_content = response.content if hasattr(response, 'content') else str(response)
+
+        if response_content.lower().startswith("assistant:"):
+            response_content = response_content[len("assistant:"):].strip()
+            
+        self.chat_history.append(ChatMessage(role=MessageRole.SYSTEM, content=response_content))
+        
+        return response.message.content

HybridRetriever.py

@@ -0,0 +1,62 @@
+from llama_index.retrievers.bm25 import BM25Retriever 
+from llama_index.core.retrievers import VectorIndexRetriever
+from llama_index.core import Document
+
+class HybridRetriever:
+    def __init__(self, bm25_retriever: BM25Retriever, vector_retriever: VectorIndexRetriever):
+        """
+        Inıtializes a Hybrid Retriever with BM25Retriever and VectorIndexRetriever.
+
+        Args:
+            bm25_retriever (BM25Retriever): An instance of BM25Retriever for keyword-based retrieval.
+            vector_retriever (VectorIndexRetriever): An instance of VectorIndexRetriever for vector-based retrieval.
+        """
+        
+        self.bm25_retriever = bm25_retriever
+        self.vector_retriever = vector_retriever
+        self.top_k = vector_retriever._similarity_top_k + bm25_retriever._similarity_top_k
+
+    def retrieve(self, query: str):
+        """
+        Retrieves documents relevant to the query using both BM25 and vector retrieval methods.
+        
+        Args:
+            query (str): The query string for which relevant documents are to be retrieved.
+        
+        Returns:
+            list: A list of tuples, each containing the document text and its combined score.
+        """
+        query = "[INST] " + " [/INST]"
+        # Perform keyword search using BM25 retriever
+        bm25_results = self.bm25_retriever.retrieve(query)
+        
+        # Perform vector search using VectorIndexRetriever
+        vector_results = self.vector_retriever.retrieve(query)
+        
+        # Combine results, filter duplicates, and calculate combined scores
+        combined_results = {}
+        for result in bm25_results:
+            combined_results[result.node.text] = {'score': result.score}
+        
+        for result in vector_results:
+            if result.node.text in combined_results:
+                combined_results[result.node.text]['score'] += result.score
+            else:
+                combined_results[result.node.text] = {'score': result.score}
+
+        # Convert combined results to a list of tuples and sort by score
+        combined_results_list = sorted(combined_results.items(), key=lambda item: item[1]['score'], reverse=True)
+        return combined_results_list # {score, document}
+        
+    def best_docs(self, query: str):
+        """
+        Retrieves the most relevant documents to the query as Document objects with their scores.
+
+        Args:
+            query (str): The query string for which the most relevant documents are to be retrieved.
+
+        Returns:
+            list: A list of tuples, each containing a Document object and its score.
+        """
+        top_results = self.retrieve(query)
+        return [(Document(text=text), score) for text, score in top_results]

configs.py

@@ -0,0 +1,14 @@
+import torch
+
+MODEL_NAME = "mistralai/Mistral-7B-Instruct-v0.3"
+EMBEDDING_NAME = "intfloat/e5-mistral-7b-instruct"
+TEMPERATURE = 0.0
+CONTEXT_WINDOW = 32_000
+TOP_K = 5
+CHUNK_SIZE = 512
+CHUNK_OVERLAP = 10
+SYSTEM_PROMPT = """[INST] You are a helpful assistant that answers 
+                user questions using the documents provided. 
+                Your answer MUST be in markdown format without 
+                any prefixes like 'assistant:' [/INST]"""
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"

helper_functions.py

@@ -0,0 +1,19 @@
+def print_results(query: str, results):
+    """
+        Prints the retrieved documents and their scores.
+        
+        Args:
+            query (str): The query string for which documents were retrieved.
+            results (list): A list of tuples, each containing a Document object and its score.
+
+        Example usage:
+        ```python
+        query = "Fee"
+        hybrid_retriever = HybridRetriever(bm25_retriever=bm25_retriever, vector_retriever=vector_retriever)
+        results = hybrid_retriever.best_docs(query)
+        print_results(query, results)
+        ```
+    """
+    print(f"\n\t\tQuery: {query}")
+    for doc, score in results:
+        print(f"Document: {doc} | Score: {score['score']:.4f}\n")

requirements.txt

@@ -0,0 +1,6 @@
+llama-index==0.10.43
+llama-index-retrievers-bm25==0.1.3
+llama-index-llms-huggingface==0.2.3
+llama-index-embeddings-huggingface==0.2.1
+llama-index-embeddings-instructor==0.1.3
+docx2txt==0.8