Uploading Mulitimodal Retrieval Augmented Generation System.

app.py

@@ -1,16 +1,18 @@
 """Gradio setup for the Multimodal RAG system."""
 import os
 import torch
+import shutil
 import gradio as gr
 # import gc
 
-from utils import load_faiss_index, load_cache
+from utils import save_cache, load_cache, save_faiss_index, load_faiss_index
 from model_setup import embedding_model, model, processor
 from main import preprocess_pdf, semantic_search, generate_answer_stream
 
 torch.set_num_threads(4)  # cpu thread limit
 
-CACHE_DIR = "cache"
+# Creating a cache directory for the retrieved chunks and index files
+CACHE_DIR = "cache_dir"
 os.makedirs(CACHE_DIR, exist_ok=True)
 
 INDEX_FILE = os.path.join(CACHE_DIR, "index.faiss")
@@ -21,73 +23,85 @@ state = {
     "index": None,
     "chunks": None,
     "pdf_path": None,
-    "image_dir": "extracted_images",
 }
 
-# Function to clear cache to prevent stale cache retrieval if new document is uploaded
-def clear_cache_files():
-  if os.path.exists(INDEX_FILE):
-      os.remove(INDEX_FILE)
-  if os.path.exists(CHUNKS_FILE):
-      os.remove(CHUNKS_FILE)
-  state["index"], state["chunks"] = None, None
-
 def handle_pdf_upload(file):
   if file is None:
-      return "[ERROR ⚠️] No file uploaded."
+      return "[ERROR] No file uploaded."
+
+  state["pdf_path"] = file.name
+  state["image_dir"] = os.path.join(CACHE_DIR, "extracted_images")  
+
+  if os.path.exists(INDEX_FILE) and os.path.exists(CHUNKS_FILE):
+    # Load from cache
+    state["index"] = load_faiss_index(INDEX_FILE)
+    state["chunks"] = load_cache(CHUNKS_FILE)
+    return "✅ Loaded from cache and ready for Q&A!"
+  else:
+    # Run your PDF preprocessing
+    index, chunks = preprocess_pdf(
+        state["pdf_path"], 
+        state["image_dir"], 
+        embedding_model=embedding_model,
+        index_file=INDEX_FILE,
+        chunks_file=CHUNKS_FILE,
+        use_cache=True)
+    state["index"] = index
+    state["chunks"] = chunks
+
+    # Save to cache
+    save_faiss_index(index, INDEX_FILE)
+    save_cache(chunks, CHUNKS_FILE)
+
+    return "✅ Document processed and ready for Q&A!"
 
-  # Save uploaded file to cache directory to ensure accessibility
-  pdf_path = os.path.join(CACHE_DIR, os.path.basename(file.name))
-  with open(pdf_path, "wb") as f_out:
-      f_out.write(file.file.read())
-
-  if state["pdf_path"] != pdf_path:
-      clear_cache_files()
+def chat_streaming(message, history):
+    if state["index"] is None or state["chunks"] is None:
+      yield "[ERROR] Please upload and process a PDF first."
 
-  state["pdf_path"] = pdf_path
+    # Perform semantic search
+    retrieved_chunks = semantic_search(message, embedding_model, state["index"], state["chunks"], top_k=10)
 
-  index, chunks = preprocess_pdf(
-      file_path=state["pdf_path"],
-      image_dir=state["image_dir"],
-      embedding_model=embedding_model,
-      index_file=INDEX_FILE,
-      chunks_file=CHUNKS_FILE,
-      use_cache=True
-  )
-  state["index"], state["chunks"] = index, chunks
-  return "✅ Document processed and ready for Q&A!"
+    # Stream the answer
+    for partial in generate_answer_stream(message, retrieved_chunks, model, processor):
+      yield partial
 
-def chat_streaming(message, history):
-  if state["index"] is None or state["chunks"] is None:
-    yield "[ERROR ⚠️] Please upload and process a PDF first."
-    return
-  retrieved_chunks = semantic_search(message, embedding_model, state["index"], state["chunks"], top_k=10)
-  for partial in generate_answer_stream(message, retrieved_chunks, model, processor):
-    yield partial
+# Function for clearing the cache files before uploading another document to prevent stale cache retrieval
+def manual_clear_cache():
+  if not os.path.exists(INDEX_FILE) and not os.path.exists(CHUNKS_FILE):
+    return "⚠️No cache files exists to clear."
+  if os.path.exists(CACHE_DIR):
+    shutil.rmtree(CACHE_DIR)
 
+  state["index"], state["chunks"] = None, None
+  return "✅ Cache cleared! You can upload a new document now."
 
 description = """
-Remember to be specific when querying for better response.
-📖🧐 
+ Remember to be specific when querying for better response.
+ 📖🧐 
 """
-# Gradio setup
-with gr.Blocks() as demo:
-  gr.Markdown("""## 📚Simple Multimodal RAG System
-                Upload a PDF (≤50 pages recommended) and ask questions about it.""")
-  with gr.Row():
-    file_input = gr.File(label="📂Upload PDF")
-    upload_button = gr.Button("🔁Process PDF")
-  upload_status = gr.Textbox(label="Upload Status", interactive=False)
-
-  upload_button.click(handle_pdf_upload, inputs=file_input, outputs=upload_status)
-
-  chat = gr.ChatInterface(
-      fn=chat_streaming,
-      type="messages",
-      title="📄😃 Ask Questions on your PDF!",
-      description=description,
-      examples=[["What is this document about?"]]
-  )
-  chat.queue()
 
+with gr.Blocks() as demo:
+    gr.Markdown("## 📚Multimodal RAG System\nUpload a PDF (≤50 pages recommended) and ask questions about it.")
+
+    with gr.Row():
+      file_input = gr.File(label="📂Upload PDF")
+      upload_button = gr.Button("🔁Process PDF")
+    
+    with gr.Row():
+      clear_cache_button = gr.Button("🧹 Clear Cache")
+      clear_cache_status = gr.Textbox(label="Cache Clear Status", interactive=False)
+
+    upload_status = gr.Textbox(label="Upload Status", interactive=False)    
+    upload_button.click(handle_pdf_upload, inputs=file_input, outputs=upload_status)
+    clear_cache_button.click(manual_clear_cache, outputs=clear_cache_status)
+
+    chat = gr.ChatInterface(
+            fn=chat_streaming,
+            type="messages",
+            title="📄Ask Questions from PDF",
+            description=description,
+            examples=[["What is this document about?"]]
+        )
+    chat.queue()
 demo.launch()

utils.py

@@ -4,6 +4,7 @@ import os
 import gc
 import json
 import torch
+import shutil
 from typing import List, Dict
 import faiss
 import numpy as np
@@ -53,11 +54,8 @@ def load_faiss_index(filepath: str):
 # Deleting extracted images directory after captioning
 def cleanup_images(image_dir: str):
   try:
-    for filename in os.listdir(image_dir):
-      file_path = os.path.join(image_dir, filename)
-      if os.path.isfile(file_path):
-        os.remove(file_path)
-    print(f"[INFO] Cleaned up extracted images in: {image_dir}")
+    shutil.rmtree(image_dir)
+    print(f"[INFO] Cleaned up extracted images directory: {image_dir}")
   except Exception as e:
     print(f"[WARNING] Failed to delete some images in {image_dir}: {e}")