Update app.py

app.py

@@ -1,4 +1,6 @@
 import os
+import multiprocessing
+import concurrent.futures
 from langchain.document_loaders import TextLoader, DirectoryLoader
 from langchain.text_splitter import RecursiveCharacterTextSplitter
 from langchain.vectorstores import FAISS
@@ -8,61 +10,44 @@ import torch
 import numpy as np
 from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer, BitsAndBytesConfig
 from datetime import datetime
+import json
 import gradio as gr
 import re
 from threading import Thread
+from transformers.agents import Tool, HfEngine, ReactJsonAgent
+from huggingface_hub import InferenceClient
+import logging
 
-class MultiDocumentAgentSystem:
-    def __init__(self, documents_dict, model, tokenizer, embeddings):
-        self.model = model
-        self.tokenizer = tokenizer
-        self.embeddings = embeddings
-        self.document_vectors = self.create_document_vectors(documents_dict)
-
-    def create_document_vectors(self, documents_dict):
-        document_vectors = {}
-        for doc_name, content in documents_dict.items():
-            vectors = self.embeddings.encode(content, convert_to_tensor=True)
-            document_vectors[doc_name] = vectors
-        return document_vectors
-
-    def query(self, user_input):
-        query_vector = self.embeddings.encode(user_input, convert_to_tensor=True)
-        
-        # Find the most similar document
-        most_similar_doc = max(self.document_vectors.items(), 
-                               key=lambda x: torch.cosine_similarity(query_vector, x[1], dim=0))
-        
-        # Generate response using the most similar document as context
-        response = self.generate_response(user_input, most_similar_doc[0], most_similar_doc[1])
-        return response
-
-    def generate_response(self, query, doc_name, doc_vector):
-        prompt = f"Based on the document '{doc_name}', answer the following question: {query}"
-        input_ids = self.tokenizer.encode(prompt, return_tensors="pt").to(self.model.device)
-        
-        with torch.no_grad():
-            output = self.model.generate(input_ids, max_length=150, num_return_sequences=1)
-        
-        response = self.tokenizer.decode(output[0], skip_special_tokens=True)
-        return response
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
 
 class DocumentRetrievalAndGeneration:
     def __init__(self, embedding_model_name, lm_model_id, data_folder):
-        self.documents_dict = self.load_documents(data_folder)
+        self.all_splits = self.load_documents(data_folder)
         self.embeddings = SentenceTransformer(embedding_model_name)
+        self.vectordb = self.create_faiss_index()
         self.tokenizer, self.model = self.initialize_llm(lm_model_id)
-        self.multi_doc_system = MultiDocumentAgentSystem(self.documents_dict, self.model, self.tokenizer, self.embeddings)
+        self.retriever_tool = self.create_retriever_tool()
+        self.agent = self.create_agent()
 
     def load_documents(self, folder_path):
-        documents_dict = {}
-        for file_name in os.listdir(folder_path):
-            if file_name.endswith('.txt'):
-                file_path = os.path.join(folder_path, file_name)
-                with open(file_path, 'r', encoding='utf-8') as file:
-                    content = file.read()
-                    documents_dict[file_name[:-4]] = content
-        return documents_dict
+        loader = DirectoryLoader(folder_path, loader_cls=TextLoader)
+        documents = loader.load()
+        text_splitter = RecursiveCharacterTextSplitter(chunk_size=200, chunk_overlap=20)
+        all_splits = text_splitter.split_documents(documents)
+        logger.info(f'Loaded {len(documents)} documents')
+        logger.info(f"Split into {len(all_splits)} chunks")
+        return all_splits
+
+    def create_faiss_index(self):
+        all_texts = [split.page_content for split in self.all_splits]
+        embeddings = self.embeddings.encode(all_texts, convert_to_tensor=True).cpu().numpy()
+        vectordb = FAISS.from_embeddings(
+            embeddings,
+            self.embeddings,
+            metadatas=[{"source": f"doc_{i}"} for i in range(len(all_texts))]
+        )
+        return vectordb
 
     def initialize_llm(self, model_id):
         quantization_config = BitsAndBytesConfig(
@@ -80,44 +65,79 @@ class DocumentRetrievalAndGeneration:
         )
         return tokenizer, model
 
-    def generate_response_with_timeout(self, input_ids, max_new_tokens=1000):
-        try:
-            streamer = TextIteratorStreamer(self.tokenizer, timeout=60.0, skip_prompt=True, skip_special_tokens=True)
-            generate_kwargs = dict(
-                input_ids=input_ids,
-                max_new_tokens=max_new_tokens,
-                do_sample=True,
-                top_p=1.0,
-                top_k=20,
-                temperature=0.8,
-                repetition_penalty=1.2,
-                eos_token_id=self.tokenizer.eos_token_id,
-                streamer=streamer,
-            )
-            
-            thread = Thread(target=self.model.generate, kwargs=generate_kwargs)
-            thread.start()
-            
-            generated_text = ""
-            for new_text in streamer:
-                generated_text += new_text
-            
-            return generated_text
-        except Exception as e:
-            print(f"Error in generate_response_with_timeout: {str(e)}")
-            return "Text generation process encountered an error"
+    def create_retriever_tool(self):
+        class RetrieverTool(Tool):
+            name = "retriever"
+            description = "Retrieves documents from the knowledge base that are semantically similar to the input query."
+            inputs = {
+                "query": {
+                    "type": "text",
+                    "description": "The query to perform. Use affirmative form rather than a question.",
+                }
+            }
+            output_type = "text"
+
+            def __init__(self, vectordb, **kwargs):
+                super().__init__(**kwargs)
+                self.vectordb = vectordb
+
+            def forward(self, query: str) -> str:
+                docs = self.vectordb.similarity_search(query, k=3)
+                return "\nRetrieved documents:\n" + "".join(
+                    [f"===== Document {str(i)} =====\n" + doc.page_content for i, doc in enumerate(docs)]
+                )
+
+        return RetrieverTool(self.vectordb)
+
+    def create_agent(self):
+        llm_engine = HfEngine("meta-llama/Meta-Llama-3.1-8B-Instruct")
+        return ReactJsonAgent(tools=[self.retriever_tool], llm_engine=llm_engine, max_iterations=4, verbose=2)
+
+    def run_agentic_rag(self, question: str) -> str:
+        enhanced_question = f"""Using the information in your knowledge base, accessible with the 'retriever' tool,
+give a comprehensive answer to the question below.
+Respond only to the question asked, be concise and relevant.
+If you can't find information, try calling your retriever again with different arguments.
+Make sure to cover the question completely by calling the retriever tool several times with semantically different queries.
+Your queries should be in affirmative form, not questions.
+
+Question:
+{question}"""
+
+        return self.agent.run(enhanced_question)
+
+    def run_standard_rag(self, question: str) -> str:
+        context = self.retriever_tool(query=question)
+
+        prompt = f"""Given the question and supporting documents below, give a comprehensive answer to the question.
+Respond only to the question asked, be concise and relevant.
+Provide the number of the source document when relevant.
+
+Question:
+{question}
+
+{context}
+"""
+        messages = [{"role": "user", "content": prompt}]
+
+        reader_llm = InferenceClient("meta-llama/Meta-Llama-3.1-8B-Instruct")
+
+        return reader_llm.chat_completion(messages).choices[0].message.content
 
     def query_and_generate_response(self, query):
-        response = self.multi_doc_system.query(query)
-        return str(response), ""
+        agentic_answer = self.run_agentic_rag(query)
+        standard_answer = self.run_standard_rag(query)
+        
+        combined_answer = f"Agentic RAG Answer:\n{agentic_answer}\n\nStandard RAG Answer:\n{standard_answer}"
+        return combined_answer, ""  # Return empty string for 'content' as it's not used in this implementation
 
     def qa_infer_gradio(self, query):
-        response, related_queries = self.query_and_generate_response(query)
-        return response, related_queries
+        response = self.query_and_generate_response(query)
+        return response
 
 if __name__ == "__main__":
-    embedding_model_name = 'sentence-transformers/all-MiniLM-L6-v2'
-    lm_model_id = "facebook/opt-350m"  # You can change this to a different open-source model
+    embedding_model_name = 'thenlper/gte-small'
+    lm_model_id = "meta-llama/Meta-Llama-3.1-8B-Instruct"
     data_folder = 'sample_embedding_folder2'
 
     doc_retrieval_gen = DocumentRetrievalAndGeneration(embedding_model_name, lm_model_id, data_folder)
@@ -151,7 +171,7 @@ if __name__ == "__main__":
             cache_examples=False,
             outputs=[gr.Textbox(label="RESPONSE"), gr.Textbox(label="RELATED QUERIES")],
             css=css_code,
-            title="TI E2E FORUM"
+            title="TI E2E FORUM Multi-Agent RAG"
         )
 
         interface.launch(debug=True)