Update app.py

app.py

@@ -14,6 +14,8 @@ from transformers.agents import Tool, HfEngine, ReactJsonAgent
 from huggingface_hub import InferenceClient
 import logging
 import torch
+import numpy as np
+import faiss
 
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
@@ -29,11 +31,7 @@ class DocumentRetrievalAndGeneration:
     def __init__(self, embedding_model_name, lm_model_id, data_folder):
         self.all_splits = self.load_documents(data_folder)
         self.embeddings = HuggingFaceEmbeddings(model_name=embedding_model_name)
-        if FAISS is not None:
-            self.vectordb = self.create_faiss_index()
-        else:
-            logger.warning("FAISS is not available. Vector search functionality will be limited.")
-            self.vectordb = None
+        self.gpu_index = self.create_faiss_index()
         self.tokenizer, self.model = self.initialize_llm(lm_model_id)
         self.retriever_tool = self.create_retriever_tool()
         self.agent = self.create_agent()
@@ -41,17 +39,20 @@ class DocumentRetrievalAndGeneration:
     def load_documents(self, folder_path):
         loader = DirectoryLoader(folder_path, loader_cls=TextLoader)
         documents = loader.load()
-        text_splitter = RecursiveCharacterTextSplitter(chunk_size=200, chunk_overlap=20)
+        text_splitter = RecursiveCharacterTextSplitter(chunk_size=5000, chunk_overlap=250)
         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):
-        if FAISS is None:
-            logger.error("FAISS is not available. Cannot create index.")
-            return None
-        return FAISS.from_documents(self.all_splits, self.embeddings)
+        all_texts = [split.page_content for split in self.all_splits]
+        embeddings = self.embeddings.embed_documents(all_texts)
+        index = faiss.IndexFlatL2(len(embeddings[0]))
+        index.add(np.array(embeddings))
+        gpu_resource = faiss.StandardGpuResources()
+        gpu_index = faiss.index_cpu_to_gpu(gpu_resource, 0, index)
+        return gpu_index
 
     def initialize_llm(self, model_id):
         quantization_config = BitsAndBytesConfig(
@@ -81,24 +82,56 @@ class DocumentRetrievalAndGeneration:
             }
             output_type = "text"
 
-            def __init__(self, vectordb, **kwargs):
+            def __init__(self, parent, **kwargs):
                 super().__init__(**kwargs)
-                self.vectordb = vectordb
+                self.parent = parent
 
             def forward(self, query: str) -> str:
-                if self.vectordb is None:
-                    return "Vector database is not available. Cannot perform retrieval."
-                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)
+                similarityThreshold = 1
+                query_embedding = self.parent.embeddings.embed_query(query)
+                distances, indices = self.parent.gpu_index.search(np.array([query_embedding]), k=3)
+                content = ""
+                filtered_results = []
+                for idx, distance in zip(indices[0], distances[0]):
+                    if distance <= similarityThreshold:
+                        filtered_results.append(idx)
+                    content += "-" * 50 + "\n"
+                    content += self.parent.all_splits[idx].page_content + "\n"
+                return content
+
+        return RetrieverTool(self)
 
     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 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=[128001, 128008, 128009],
+                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:
+            logger.error(f"Error in generate_response_with_timeout: {str(e)}")
+            return "Text generation process encountered an error"
+
     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.
@@ -115,20 +148,23 @@ 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
+        conversation = [
+            {"role": "system", "content": "You are a knowledgeable assistant with access to a comprehensive database."},
+            {"role": "user", "content": f"""
+            I need you to answer my question and provide related information in a specific format.
+            I have provided five relatable json files {context}, choose the most suitable chunks for answering the query.
+            RETURN ONLY SOLUTION without additional comments, sign-offs, retrived chunks, refrence to any Ticket or extra phrases. Be direct and to the point.
+            IF THERE IS NO ANSWER RELATABLE IN RETRIEVED CHUNKS, RETURN "NO SOLUTION AVAILABLE".
+            DO NOT GIVE REFRENCE TO ANY CHUNKS OR TICKETS,BE ON POINT.
+            
+            Here's my question:
+            Query: {question}
+            Solution==>
+            """}
+        ]
+        input_ids = self.tokenizer.apply_chat_template(conversation, return_tensors="pt").to(self.model.device)
+        
+        return self.generate_response_with_timeout(input_ids)
 
     def query_and_generate_response(self, query):
         agentic_answer = self.run_agentic_rag(query)
@@ -141,29 +177,17 @@ Question:
         response = self.query_and_generate_response(query)
         return response
 
-    def save_index(self, path):
-        if self.vectordb is not None:
-            self.vectordb.save_local(path)
-        else:
-            logger.warning("Vector database is not available. Cannot save index.")
-
-    def load_index(self, path):
-        if FAISS is not None:
-            self.vectordb = FAISS.load_local(path, self.embeddings)
-        else:
-            logger.warning("FAISS is not available. Cannot load index.")
-
 if __name__ == "__main__":
-    embedding_model_name = 'thenlper/gte-small'
+    embedding_model_name = 'flax-sentence-embeddings/all_datasets_v3_MiniLM-L12'
     lm_model_id = "meta-llama/Meta-Llama-3.1-8B-Instruct"
     data_folder = 'sample_embedding_folder2'
 
+    # Set your HuggingFace token here
+    os.environ["HUGGINGFACE_TOKEN"] = "your_huggingface_token_here"
+
     try:
         doc_retrieval_gen = DocumentRetrievalAndGeneration(embedding_model_name, lm_model_id, data_folder)
 
-        # Save the index for future use
-        doc_retrieval_gen.save_index("faiss_index")
-
         def launch_interface():
             css_code = """
                 .gradio-container {