Update app.py

app.py

@@ -1,28 +1,26 @@
 import os
 import multiprocessing
 import concurrent.futures
-# from langchain.document_loaders import TextLoader, DirectoryLoader
 from langchain_community.document_loaders import TextLoader, DirectoryLoader
 from langchain.text_splitter import RecursiveCharacterTextSplitter
-# from langchain.vectorstores import FAISS
 from langchain_community.vectorstores import FAISS 
 from sentence_transformers import SentenceTransformer
 import faiss
-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 openai import OpenAI
+
 class MultiAgentRAG:
-    def __init__(self, embedding_model_name, lm_model_id, data_folder):
+    def __init__(self, embedding_model_name, openai_model_id, data_folder, api_key=None):
         self.all_splits = self.load_documents(data_folder)
         self.embeddings = SentenceTransformer(embedding_model_name)
         self.gpu_index = self.create_faiss_index()
-        self.tokenizer, self.model = self.initialize_llm(lm_model_id)
+        self.openai_client = OpenAI(api_key=api_key or os.environ.get("OPENAI_API_KEY"))
+        self.openai_model_id = openai_model_id
 
     def load_documents(self, folder_path):
         loader = DirectoryLoader(folder_path, loader_cls=TextLoader)
@@ -31,7 +29,7 @@ class MultiAgentRAG:
         all_splits = text_splitter.split_documents(documents)
         print('Length of documents:', len(documents))
         print("LEN of all_splits", len(all_splits))
-        for i in range(3):
+        for i in range(min(3, len(all_splits))):
             print(all_splits[i].page_content)
         return all_splits
 
@@ -44,47 +42,20 @@ class MultiAgentRAG:
         gpu_index = faiss.index_cpu_to_gpu(gpu_resource, 0, index)
         return gpu_index
 
-    def initialize_llm(self, model_id):
-        quantization_config = BitsAndBytesConfig(
-            load_in_4bit=True,
-            bnb_4bit_use_double_quant=True,
-            bnb_4bit_quant_type="nf4",
-            bnb_4bit_compute_dtype=torch.bfloat16
-        )
-        tokenizer = AutoTokenizer.from_pretrained(model_id)
-        model = AutoModelForCausalLM.from_pretrained(
-            model_id,
-            torch_dtype=torch.bfloat16,
-            device_map="auto",
-            quantization_config=quantization_config
-        )
-        return tokenizer, model
-
-    def generate_response_with_timeout(self, input_ids, max_new_tokens=1000):
+    def generate_openai_response(self, messages, max_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,
+            response = self.openai_client.chat.completions.create(
+                model=self.openai_model_id,
+                messages=messages,
+                max_tokens=max_tokens,
                 temperature=0.8,
-                repetition_penalty=1.2,
-                eos_token_id=[128001, 128008, 128009],
-                streamer=streamer,
+                top_p=1.0,
+                frequency_penalty=0,
+                presence_penalty=0
             )
-            
-            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
+            return response.choices[0].message.content
         except Exception as e:
-            print(f"Error in generate_response_with_timeout: {str(e)}")
+            print(f"Error in generate_openai_response: {str(e)}")
             return "Text generation process encountered an error"
 
     def retrieval_agent(self, query):
@@ -97,44 +68,51 @@ class MultiAgentRAG:
         return content
 
     def grading_agent(self, query, retrieved_content):
-        grading_prompt = f"""
-        Evaluate the relevance of the following retrieved content to the given query:
-        
-        Query: {query}
-        
-        Retrieved Content:
-        {retrieved_content}
+        messages = [
+            {"role": "system", "content": "You are an expert at evaluating the relevance of retrieved content to a query."},
+            {"role": "user", "content": f"""
+            Evaluate the relevance of the following retrieved content to the given query:
+            
+            Query: {query}
+            
+            Retrieved Content:
+            {retrieved_content}
+            
+            Rate the relevance on a scale of 1-10 and explain your rating:
+            """}
+        ]
         
-        Rate the relevance on a scale of 1-10 and explain your rating:
-        """
-        input_ids = self.tokenizer.encode(grading_prompt, return_tensors="pt").to(self.model.device)
-        grading_response = self.generate_response_with_timeout(input_ids)
+        grading_response = self.generate_openai_response(messages)
         
         # Extract the numerical rating from the response
-        rating = int(re.search(r'\d+', grading_response).group())
+        match = re.search(r'\b([1-9]|10)\b', grading_response)
+        rating = int(match.group()) if match else 5  # Default to 5 if no rating found
         return rating, grading_response
 
     def query_rewrite_agent(self, original_query):
-        rewrite_prompt = f"""
-        The following query did not yield relevant results. Please rewrite it to potentially improve retrieval:
-        
-        Original Query: {original_query}
+        messages = [
+            {"role": "system", "content": "You are an expert at rewriting queries to improve information retrieval results."},
+            {"role": "user", "content": f"""
+            The following query did not yield relevant results. Please rewrite it to potentially improve retrieval:
+            
+            Original Query: {original_query}
+            
+            Rewritten Query:
+            """}
+        ]
         
-        Rewritten Query:
-        """
-        input_ids = self.tokenizer.encode(rewrite_prompt, return_tensors="pt").to(self.model.device)
-        rewritten_query = self.generate_response_with_timeout(input_ids)
+        rewritten_query = self.generate_openai_response(messages)
         return rewritten_query.strip()
 
     def generation_agent(self, query, retrieved_content):
-        conversation = [
+        messages = [
             {"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 {retrieved_content}, 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.
+            DO NOT GIVE REFRENCE TO ANY CHUNKS OR TICKETS, BE ON POINT.
             
             Here's my question:
             Query: {query}
@@ -142,24 +120,19 @@ class MultiAgentRAG:
             """}
         ]
         
-        input_ids = self.tokenizer.encode(self.tokenizer.apply_chat_template(conversation, tokenize=False), return_tensors="pt").to(self.model.device)
-        return self.generate_response_with_timeout(input_ids)
+        return self.generate_openai_response(messages)
 
     def run_multi_agent_rag(self, query):
         max_iterations = 3
         for i in range(max_iterations):
-            # Retrieval step
             retrieved_content = self.retrieval_agent(query)
             
-            # Grading step
             relevance_score, grading_explanation = self.grading_agent(query, retrieved_content)
             
-            if relevance_score >= 7:  # Assuming 7 out of 10 is the threshold for relevance
-                # Generation step
+            if relevance_score >= 7:                 
                 answer = self.generation_agent(query, retrieved_content)
                 return answer, retrieved_content, grading_explanation
             else:
-                # Query rewrite step
                 query = self.query_rewrite_agent(query)
         
         return "Unable to find a relevant answer after multiple attempts.", "", "Low relevance across all attempts."
@@ -204,8 +177,8 @@ def launch_interface(doc_retrieval_gen):
 
 if __name__ == "__main__":
     embedding_model_name = 'flax-sentence-embeddings/all_datasets_v3_MiniLM-L12'
-    lm_model_id = "meta-llama/Meta-Llama-3.1-8B-Instruct"
+    openai_model_id = "gpt-4-turbo" 
     data_folder = 'sample_embedding_folder2'
 
-    multi_agent_rag = MultiAgentRAG(embedding_model_name, lm_model_id, data_folder)
+    multi_agent_rag = MultiAgentRAG(embedding_model_name, openai_model_id, data_folder)
     launch_interface(multi_agent_rag)