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"]
             )