Update app.py

app.py

@@ -9,12 +9,16 @@ import gradio as gr
 import os
 import pytesseract
 from PIL import Image
-
+import pickle 
 model = SentenceTransformer('sentence-transformers/all-MiniLM-L6-v2')
 index = faiss.read_index('IPC_index.faiss')
 index2 = faiss.read_index('CrpC_index.faiss')
-
-
+flattened_data = pickle.load('IPC_F')
+pdf_filenames = pickle.load('IPC_N')
+chunk_indices = pickle.load('IPC_C')
+flattened_data2 = pickle.load('CrPC_F')
+pdf_filenames2 = pickle.load('CrPC_N')
+chunk_indices2 = pickle.load('CrPC_C')
 # Step 3: Retrieval with Citations using PDF filename
 def retrieve_info_with_citation(query, top_k=5):
     query_embedding = model.encode([query])
@@ -24,7 +28,10 @@ def retrieve_info_with_citation(query, top_k=5):
     for i in range(min(top_k, len(I[0]))):
         if D[0][i] < 1.0:  # Relevance threshold
             chunk_index = I[0][i]
-            citation = f"Source: IPC"
+            pdf_filename = pdf_filenames[chunk_index]
+            chunk_number = chunk_indices[chunk_index] + 1 
+            match = flattened_data[chunk_index]
+            citation = f"Source: {pdf_filename}, Chunk: {chunk_number}"
             results.append((match, citation))
         else:
             break
@@ -37,13 +44,16 @@ def retrieve_info_with_citation(query, top_k=5):
 
 def retrieve_info_with_citation2(query, top_k=5):
     query_embedding = model.encode([query])
-    D, I = index.search(query_embedding, k=top_k)
+    D, I = index2.search(query_embedding, k=top_k)
 
     results = []
     for i in range(min(top_k, len(I[0]))):
         if D[0][i] < 1.0:  # Relevance threshold
             chunk_index = I[0][i]
-            citation = f"Source: CrPC"
+            pdf_filename = pdf_filenames2[chunk_index]
+            chunk_number = chunk_indices2[chunk_index] + 1 
+            match = flattened_data2[chunk_index]
+            citation = f"Source: {pdf_filename}, Chunk: {chunk_number}"
             results.append((match, citation))
         else:
             break
@@ -63,6 +73,7 @@ def retrieve_info2(query):
     formatted_results = "\n\n".join([f"{i+1}. {match}\n{citation}" for i, (match, citation) in enumerate(results)])
     return formatted_results
 
+
 ipc_tool = Tool(
     name="IPC Information Retrieval",
     func=retrieve_info,