Update app.py

app.py

@@ -67,87 +67,182 @@ hybrid_rag = HybridColpaliRAG(
 IngestResult = namedtuple("IngestResult", ["status_text", "progress_table"])
 
 
-@spaces.GPU(duration=120)
+# @spaces.GPU(duration=120)
+# def ingest_data(pdf_files, use_ocr, chunk_size, progress=gr.Progress()):
+# file_paths = [pdf_file.name for pdf_file in pdf_files]
+# total_start_time = time.time()
+# progress_data = []
+
+# # SimpleRAG
+# yield IngestResult(
+# status_text="Starting SimpleRAG ingestion...\n",
+# progress_table=pd.DataFrame(progress_data),
+# )
+# start_time = time.time()
+# simple_rag.index(
+# file_paths,
+# recursive=False,
+# chunking_strategy=FixedTokenChunker(chunk_size=chunk_size),
+# metadata={"source": "gradio_upload"},
+# overwrite=True,
+# verbose=True,
+# ocr=use_ocr,
+# )
+# simple_time = time.time() - start_time
+# progress_data.append(
+# {"Technique": "SimpleRAG", "Time Taken (s)": f"{simple_time:.2f}"}
+# )
+# yield IngestResult(
+# status_text=f"SimpleRAG ingestion complete. Time taken: {simple_time:.2f} seconds\n\n",
+# progress_table=pd.DataFrame(progress_data),
+# )
+# # progress(0.25, desc="SimpleRAG complete")
+
+# # VisionRAG
+# yield IngestResult(
+# status_text="Starting VisionRAG ingestion...\n",
+# progress_table=pd.DataFrame(progress_data),
+# )
+# start_time = time.time()
+# vision_rag.index(file_paths, overwrite=False, recursive=False, verbose=True)
+# vision_time = time.time() - start_time
+# progress_data.append(
+# {"Technique": "VisionRAG", "Time Taken (s)": f"{vision_time:.2f}"}
+# )
+# yield IngestResult(
+# status_text=f"VisionRAG ingestion complete. Time taken: {vision_time:.2f} seconds\n\n",
+# progress_table=pd.DataFrame(progress_data),
+# )
+# # progress(0.5, desc="VisionRAG complete")
+
+# # ColpaliRAG
+# yield IngestResult(
+# status_text="Starting ColpaliRAG ingestion...\n",
+# progress_table=pd.DataFrame(progress_data),
+# )
+# start_time = time.time()
+# colpali_rag.index(file_paths, overwrite=False, recursive=False, verbose=True)
+# colpali_time = time.time() - start_time
+# progress_data.append(
+# {"Technique": "ColpaliRAG", "Time Taken (s)": f"{colpali_time:.2f}"}
+# )
+# yield IngestResult(
+# status_text=f"ColpaliRAG ingestion complete. Time taken: {colpali_time:.2f} seconds\n\n",
+# progress_table=pd.DataFrame(progress_data),
+# )
+# # progress(0.75, desc="ColpaliRAG complete")
+
+# # HybridColpaliRAG
+# yield IngestResult(
+# status_text="Starting HybridColpaliRAG ingestion...\n",
+# progress_table=pd.DataFrame(progress_data),
+# )
+# start_time = time.time()
+# hybrid_rag.index(file_paths, overwrite=False, recursive=False, verbose=True)
+# hybrid_time = time.time() - start_time
+# progress_data.append(
+# {"Technique": "HybridColpaliRAG", "Time Taken (s)": f"{hybrid_time:.2f}"}
+# )
+# yield IngestResult(
+# status_text=f"HybridColpaliRAG ingestion complete. Time taken: {hybrid_time:.2f} seconds\n\n",
+# progress_table=pd.DataFrame(progress_data),
+# )
+# # progress(1.0, desc="HybridColpaliRAG complete")
+
+# total_time = time.time() - total_start_time
+# progress_data.append({"Technique": "Total", "Time Taken (s)": f"{total_time:.2f}"})
+# yield IngestResult(
+# status_text=f"Total ingestion time: {total_time:.2f} seconds",
+# progress_table=pd.DataFrame(progress_data),
+# )
+
+
 def ingest_data(pdf_files, use_ocr, chunk_size, progress=gr.Progress()):
  file_paths = [pdf_file.name for pdf_file in pdf_files]
  total_start_time = time.time()
  progress_data = []
 
- # SimpleRAG
- yield IngestResult(
- status_text="Starting SimpleRAG ingestion...\n",
- progress_table=pd.DataFrame(progress_data),
- )
- start_time = time.time()
- simple_rag.index(
- file_paths,
- recursive=False,
- chunking_strategy=FixedTokenChunker(chunk_size=chunk_size),
- metadata={"source": "gradio_upload"},
- overwrite=True,
- verbose=True,
- ocr=use_ocr,
- )
- simple_time = time.time() - start_time
- progress_data.append(
- {"Technique": "SimpleRAG", "Time Taken (s)": f"{simple_time:.2f}"}
- )
- yield IngestResult(
- status_text=f"SimpleRAG ingestion complete. Time taken: {simple_time:.2f} seconds\n\n",
- progress_table=pd.DataFrame(progress_data),
- )
- # progress(0.25, desc="SimpleRAG complete")
+ @spaces.GPU(duration=120)
+ def ingest_simple_rag():
+ yield IngestResult(
+ status_text="Starting SimpleRAG ingestion...\n",
+ progress_table=pd.DataFrame(progress_data),
+ )
+ start_time = time.time()
+ simple_rag.index(
+ file_paths,
+ recursive=False,
+ chunking_strategy=FixedTokenChunker(chunk_size=chunk_size),
+ metadata={"source": "gradio_upload"},
+ overwrite=True,
+ verbose=True,
+ ocr=use_ocr,
+ )
+ simple_time = time.time() - start_time
+ progress_data.append(
+ {"Technique": "SimpleRAG", "Time Taken (s)": f"{simple_time:.2f}"}
+ )
+ yield IngestResult(
+ status_text=f"SimpleRAG ingestion complete. Time taken: {simple_time:.2f} seconds\n\n",
+ progress_table=pd.DataFrame(progress_data),
+ )
 
- # VisionRAG
- yield IngestResult(
- status_text="Starting VisionRAG ingestion...\n",
- progress_table=pd.DataFrame(progress_data),
- )
- start_time = time.time()
- vision_rag.index(file_paths, overwrite=False, recursive=False, verbose=True)
- vision_time = time.time() - start_time
- progress_data.append(
- {"Technique": "VisionRAG", "Time Taken (s)": f"{vision_time:.2f}"}
- )
- yield IngestResult(
- status_text=f"VisionRAG ingestion complete. Time taken: {vision_time:.2f} seconds\n\n",
- progress_table=pd.DataFrame(progress_data),
- )
- # progress(0.5, desc="VisionRAG complete")
+ @spaces.GPU(duration=120)
+ def ingest_vision_rag():
+ yield IngestResult(
+ status_text="Starting VisionRAG ingestion...\n",
+ progress_table=pd.DataFrame(progress_data),
+ )
+ start_time = time.time()
+ vision_rag.index(file_paths, overwrite=False, recursive=False, verbose=True)
+ vision_time = time.time() - start_time
+ progress_data.append(
+ {"Technique": "VisionRAG", "Time Taken (s)": f"{vision_time:.2f}"}
+ )
+ yield IngestResult(
+ status_text=f"VisionRAG ingestion complete. Time taken: {vision_time:.2f} seconds\n\n",
+ progress_table=pd.DataFrame(progress_data),
+ )
 
- # ColpaliRAG
- yield IngestResult(
- status_text="Starting ColpaliRAG ingestion...\n",
- progress_table=pd.DataFrame(progress_data),
- )
- start_time = time.time()
- colpali_rag.index(file_paths, overwrite=False, recursive=False, verbose=True)
- colpali_time = time.time() - start_time
- progress_data.append(
- {"Technique": "ColpaliRAG", "Time Taken (s)": f"{colpali_time:.2f}"}
- )
- yield IngestResult(
- status_text=f"ColpaliRAG ingestion complete. Time taken: {colpali_time:.2f} seconds\n\n",
- progress_table=pd.DataFrame(progress_data),
- )
- # progress(0.75, desc="ColpaliRAG complete")
+ @spaces.GPU(duration=120)
+ def ingest_colpali_rag():
+ yield IngestResult(
+ status_text="Starting ColpaliRAG ingestion...\n",
+ progress_table=pd.DataFrame(progress_data),
+ )
+ start_time = time.time()
+ colpali_rag.index(file_paths, overwrite=False, recursive=False, verbose=True)
+ colpali_time = time.time() - start_time
+ progress_data.append(
+ {"Technique": "ColpaliRAG", "Time Taken (s)": f"{colpali_time:.2f}"}
+ )
+ yield IngestResult(
+ status_text=f"ColpaliRAG ingestion complete. Time taken: {colpali_time:.2f} seconds\n\n",
+ progress_table=pd.DataFrame(progress_data),
+ )
 
- # HybridColpaliRAG
- yield IngestResult(
- status_text="Starting HybridColpaliRAG ingestion...\n",
- progress_table=pd.DataFrame(progress_data),
- )
- start_time = time.time()
- hybrid_rag.index(file_paths, overwrite=False, recursive=False, verbose=True)
- hybrid_time = time.time() - start_time
- progress_data.append(
- {"Technique": "HybridColpaliRAG", "Time Taken (s)": f"{hybrid_time:.2f}"}
- )
- yield IngestResult(
- status_text=f"HybridColpaliRAG ingestion complete. Time taken: {hybrid_time:.2f} seconds\n\n",
- progress_table=pd.DataFrame(progress_data),
- )
- # progress(1.0, desc="HybridColpaliRAG complete")
+ @spaces.GPU(duration=120)
+ def ingest_hybrid_rag():
+ yield IngestResult(
+ status_text="Starting HybridColpaliRAG ingestion...\n",
+ progress_table=pd.DataFrame(progress_data),
+ )
+ start_time = time.time()
+ hybrid_rag.index(file_paths, overwrite=False, recursive=False, verbose=True)
+ hybrid_time = time.time() - start_time
+ progress_data.append(
+ {"Technique": "HybridColpaliRAG", "Time Taken (s)": f"{hybrid_time:.2f}"}
+ )
+ yield IngestResult(
+ status_text=f"HybridColpaliRAG ingestion complete. Time taken: {hybrid_time:.2f} seconds\n\n",
+ progress_table=pd.DataFrame(progress_data),
+ )
+
+ # Call each ingestion function
+ yield from ingest_simple_rag()
+ yield from ingest_vision_rag()
+ yield from ingest_colpali_rag()
+ yield from ingest_hybrid_rag()
 
  total_time = time.time() - total_start_time
  progress_data.append({"Technique": "Total", "Time Taken (s)": f"{total_time:.2f}"})
@@ -313,6 +408,18 @@ Built on [VARAG](https://github.com/adithya-s-k/VARAG) - Vision-Augmented Retrie
  )
 
  with gr.Tab("Ingest Data"):
+ gr.Markdown(
+ """
+## ⚠️ Important Note on Data Ingestion
+
+This Space has a maximum GPU-enabled time of 120 seconds. It's recommended to try ingesting only 1 or 2 pdfs at a time.
+
+If you want to ingest a larger amount of data, please try it out in a Google Colab notebook:
+
+[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/adithya-s-k/VARAG/blob/main/docs/demo.ipynb)
+
+ """
+ )
  pdf_input = gr.File(
  label="Upload PDF(s)", file_count="multiple", file_types=["pdf"]
  )