first

app.py

@@ -3,20 +3,33 @@ import faiss
 import numpy as np
 import requests
 import torch
+import pickle
 from transformers import AutoTokenizer, AutoModelForCausalLM
 from sentence_transformers import SentenceTransformer
 
 class CustomRetriever:
-    def __init__(self, faiss_index_path: str):
-        """Initializes the retriever by loading the FAISS index and setting up the embedding model."""
+    def __init__(self, faiss_index_path: str, metadata_path: str):
+        """Initializes the retriever by loading the FAISS index and document metadata."""
+        # Load the FAISS index
         self.index = faiss.read_index(faiss_index_path)
+        
+        # Load the document metadata (mapping FAISS indices to document content)
+        with open(metadata_path, 'rb') as file:
+            self.doc_metadata = pickle.load(file)
+        
+        # Load the SentenceTransformer for embedding queries
         self.embedder = SentenceTransformer('sentence-transformers/all-MiniLM-L6-v2')
     
     def retrieve(self, query: str, top_k: int = 5):
         """Retrieve top-k relevant documents based on the query."""
         query_embedding = self.embedder.encode([query])
+        
+        # Search the FAISS index for top-k similar embeddings
         distances, indices = self.index.search(np.array(query_embedding).astype('float32'), top_k)
-        return [(index, distance) for index, distance in zip(indices[0], distances[0])]
+        
+        # Retrieve the actual document content using the indices
+        retrieved_docs = [(self.doc_metadata[idx], distance) for idx, distance in zip(indices[0], distances[0])]
+        return retrieved_docs
 
 
 class CustomGenerator:
@@ -29,24 +42,32 @@ class CustomGenerator:
         """Generate a response using the retrieved documents and the user input."""
         context = "\n".join([f"Doc {i+1}: {doc}" for i, (doc, _) in enumerate(retrieved_docs)])
         prompt = f"Context:\n{context}\n\nUser: {user_input}\nBot:"
+        
         inputs = self.tokenizer(prompt, return_tensors="pt", truncation=True)
         with torch.no_grad():
             outputs = self.model.generate(inputs.input_ids, max_length=max_length, pad_token_id=self.tokenizer.eos_token_id)
+        
         response = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
         return response.split("Bot:")[-1].strip()
 
 
 def rag_chatbot(user_input):
     """The main RAG chatbot function to retrieve documents and generate a response."""
+    # Step 1: Retrieve relevant documents
     top_k = 5  # Number of documents to retrieve
-    retrieved_doc_ids = retriever.retrieve(user_input, top_k)
-    retrieved_docs = [(f"Dummy content for doc {doc_id}", distance) for doc_id, distance in retrieved_doc_ids]
+    retrieved_docs = retriever.retrieve(user_input, top_k)
+    
+    # Step 2: Generate a response using the documents
     response = generator.generate(user_input, retrieved_docs)
     return response
 
 
-FAISS_INDEX_PATH = ""
-retriever = CustomRetriever(faiss_index_path=FAISS_INDEX_PATH)
+# Paths to your FAISS index and metadata files
+FAISS_INDEX_PATH = r"C:\Users\schandrappa\Downloads\Banking_Regulations_Compliance_ChatBOT\faiss\index.faiss"
+METADATA_PATH = r"C:\Users\schandrappa\Downloads\Banking_Regulations_Compliance_ChatBOT\faiss\index.pkl"
+
+# Initialize retriever and generator
+retriever = CustomRetriever(faiss_index_path=FAISS_INDEX_PATH, metadata_path=METADATA_PATH)
 generator = CustomGenerator()
 
 # Gradio UI

faiss/index.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:2361a54ccf2ea94476514663e7f3cf028ab93e37d32c72faba6e0e01d4b77781
+size 1423760