VARAG initial commit - Adithya s K

app.py

@@ -0,0 +1,475 @@
+import gradio as gr
+import os
+import lancedb
+from sentence_transformers import SentenceTransformer
+from dotenv import load_dotenv
+from typing import List
+from PIL import Image
+import base64
+import io
+import time
+from collections import namedtuple
+import pandas as pd
+import concurrent.futures
+from varag.rag import SimpleRAG, VisionRAG, ColpaliRAG, HybridColpaliRAG
+from transformers import Qwen2VLForConditionalGeneration, AutoTokenizer, AutoProcessor
+from qwen_vl_utils import process_vision_info
+from varag.chunking import FixedTokenChunker
+from varag.utils import get_model_colpali
+import argparse
+import spaces
+import torch
+
+load_dotenv()
+
+# Initialize shared database
+shared_db = lancedb.connect("~/rag_demo_db")
+
+
+@spaces.GPU
+def get_all_model():
+    # Initialize embedding models
+    # text_embedding_model = SentenceTransformer("all-MiniLM-L6-v2", trust_remote_code=True)
+    text_embedding_model = SentenceTransformer(
+        "BAAI/bge-base-en", trust_remote_code=True
+    )
+    # text_embedding_model = SentenceTransformer("BAAI/bge-large-en-v1.5", trust_remote_code=True)
+    # text_embedding_model = SentenceTransformer("BAAI/bge-small-en-v1.5", trust_remote_code=True)
+    image_embedding_model = SentenceTransformer(
+        "jinaai/jina-clip-v1", trust_remote_code=True
+    )
+    colpali_model, colpali_processor = get_model_colpali("vidore/colpali-v1.2")
+
+    return text_embedding_model, image_embedding_model, colpali_model, colpali_processor
+
+
+text_embedding_model, image_embedding_model, colpali_model, colpali_processor = (
+    get_all_model()
+)
+
+# Initialize RAG instances
+simple_rag = SimpleRAG(
+    text_embedding_model=text_embedding_model, db=shared_db, table_name="simpleDemo"
+)
+vision_rag = VisionRAG(
+    image_embedding_model=image_embedding_model, db=shared_db, table_name="visionDemo"
+)
+colpali_rag = ColpaliRAG(
+    colpali_model=colpali_model,
+    colpali_processor=colpali_processor,
+    db=shared_db,
+    table_name="colpaliDemo",
+)
+hybrid_rag = HybridColpaliRAG(
+    colpali_model=colpali_model,
+    colpali_processor=colpali_processor,
+    image_embedding_model=image_embedding_model,
+    db=shared_db,
+    table_name="hybridDemo",
+)
+
+
+IngestResult = namedtuple("IngestResult", ["status_text", "progress_table"])
+
+
+@spaces.GPU
+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),
+    )
+
+
+@spaces.GPU
+def retrieve_data(query, top_k, sequential=False):
+    results = {}
+    timings = {}
+
+    def retrieve_simple():
+        start_time = time.time()
+        simple_results = simple_rag.search(query, k=top_k)
+
+        print(simple_results)
+
+        simple_context = []
+        for i, r in enumerate(simple_results, 1):
+            context_piece = f"Result {i}:\n"
+            context_piece += f"Source: {r.get('document_name', 'Unknown')}\n"
+            context_piece += f"Chunk Index: {r.get('chunk_index', 'Unknown')}\n"
+
+            context_piece += f"Content:\n{r['text']}\n"
+            context_piece += "-" * 40 + "\n"  # Separator
+            simple_context.append(context_piece)
+
+        simple_context = "\n".join(simple_context)
+        end_time = time.time()
+        return "SimpleRAG", simple_context, end_time - start_time
+
+    def retrieve_vision():
+        start_time = time.time()
+        vision_results = vision_rag.search(query, k=top_k)
+        vision_images = [r["image"] for r in vision_results]
+        end_time = time.time()
+        return "VisionRAG", vision_images, end_time - start_time
+
+    def retrieve_colpali():
+        start_time = time.time()
+        colpali_results = colpali_rag.search(query, k=top_k)
+        colpali_images = [r["image"] for r in colpali_results]
+        end_time = time.time()
+        return "ColpaliRAG", colpali_images, end_time - start_time
+
+    def retrieve_hybrid():
+        start_time = time.time()
+        hybrid_results = hybrid_rag.search(query, k=top_k, use_image_search=True)
+        hybrid_images = [r["image"] for r in hybrid_results]
+        end_time = time.time()
+        return "HybridColpaliRAG", hybrid_images, end_time - start_time
+
+    retrieval_functions = [
+        retrieve_simple,
+        retrieve_vision,
+        retrieve_colpali,
+        retrieve_hybrid,
+    ]
+
+    if sequential:
+        for func in retrieval_functions:
+            rag_type, content, timing = func()
+            results[rag_type] = content
+            timings[rag_type] = timing
+    else:
+        with concurrent.futures.ThreadPoolExecutor() as executor:
+            future_results = [executor.submit(func) for func in retrieval_functions]
+            for future in concurrent.futures.as_completed(future_results):
+                rag_type, content, timing = future.result()
+                results[rag_type] = content
+                timings[rag_type] = timing
+
+    return results, timings
+
+
+# @spaces.GPU
+# def query_data(query, retrieved_results):
+#     results = {}
+
+#     # SimpleRAG
+#     simple_context = retrieved_results["SimpleRAG"]
+#     simple_response = llm.query(
+#         context=simple_context,
+#         system_prompt="Given the below information answer the questions",
+#         query=query,
+#     )
+#     results["SimpleRAG"] = {"response": simple_response, "context": simple_context}
+
+#     # VisionRAG
+#     vision_images = retrieved_results["VisionRAG"]
+#     vision_context = f"Query: {query}\n\nRelevant image information:\n" + "\n".join(
+#         [f"Image {i+1}" for i in range(len(vision_images))]
+#     )
+#     vision_response = vlm.query(vision_context, vision_images, max_tokens=500)
+#     results["VisionRAG"] = {
+#         "response": vision_response,
+#         "context": vision_context,
+#         "images": vision_images,
+#     }
+
+#     # ColpaliRAG
+#     colpali_images = retrieved_results["ColpaliRAG"]
+#     colpali_context = f"Query: {query}\n\nRelevant image information:\n" + "\n".join(
+#         [f"Image {i+1}" for i in range(len(colpali_images))]
+#     )
+#     colpali_response = vlm.query(colpali_context, colpali_images, max_tokens=500)
+#     results["ColpaliRAG"] = {
+#         "response": colpali_response,
+#         "context": colpali_context,
+#         "images": colpali_images,
+#     }
+
+#     # HybridColpaliRAG
+#     hybrid_images = retrieved_results["HybridColpaliRAG"]
+#     hybrid_context = f"Query: {query}\n\nRelevant image information:\n" + "\n".join(
+#         [f"Image {i+1}" for i in range(len(hybrid_images))]
+#     )
+#     hybrid_response = vlm.query(hybrid_context, hybrid_images, max_tokens=500)
+#     results["HybridColpaliRAG"] = {
+#         "response": hybrid_response,
+#         "context": hybrid_context,
+#         "images": hybrid_images,
+#     }
+
+#     return results
+
+
+def update_api_key(api_key):
+    os.environ["OPENAI_API_KEY"] = api_key
+    return "API key updated successfully."
+
+
+def change_table(simple_table, vision_table, colpali_table, hybrid_table):
+    simple_rag.change_table(simple_table)
+    vision_rag.change_table(vision_table)
+    colpali_rag.change_table(colpali_table)
+    hybrid_rag.change_table(hybrid_table)
+    return "Table names updated successfully."
+
+
+def gradio_interface():
+    with gr.Blocks(
+        theme=gr.themes.Monochrome(radius_size=gr.themes.sizes.radius_none)
+    ) as demo:
+        gr.Markdown(
+            """
+# 👁️👁️ Vision RAG Playground
+
+### Explore and Compare Vision-Augmented Retrieval Techniques
+Built on [VARAG](https://github.com/adithya-s-k/VARAG) - Vision-Augmented Retrieval and Generation
+
+**[⭐ Star the Repository](https://github.com/adithya-s-k/VARAG)** to support the project!
+
+1. **Simple RAG**: Text-based retrieval with OCR support for scanned documents.
+2. **Vision RAG**: Combines text and image retrieval using cross-modal embeddings.
+3. **ColPali RAG**: Embeds entire document pages as images for layout-aware retrieval.
+4. **Hybrid ColPali RAG**: Two-stage retrieval combining image embeddings and ColPali's token-level matching.
+
+            """
+        )
+
+        with gr.Tab("Ingest Data"):
+            pdf_input = gr.File(
+                label="Upload PDF(s)", file_count="multiple", file_types=["pdf"]
+            )
+            use_ocr = gr.Checkbox(label="Use OCR (for SimpleRAG)")
+            chunk_size = gr.Slider(
+                50, 5000, value=200, step=10, label="Chunk Size (for SimpleRAG)"
+            )
+            ingest_button = gr.Button("Ingest PDFs")
+            ingest_output = gr.Markdown(
+                label="Ingestion Status :",
+            )
+            progress_table = gr.DataFrame(
+                label="Ingestion Progress", headers=["Technique", "Time Taken (s)"]
+            )
+
+        with gr.Tab("Retrieve and Query Data"):
+            query_input = gr.Textbox(label="Enter your query")
+            top_k_slider = gr.Slider(1, 10, value=3, step=1, label="Top K Results")
+            sequential_checkbox = gr.Checkbox(label="Sequential Retrieval", value=False)
+            retrieve_button = gr.Button("Retrieve")
+            query_button = gr.Button("Query")
+
+            retrieval_timing = gr.DataFrame(
+                label="Retrieval Timings", headers=["RAG Type", "Time (s)"]
+            )
+
+            with gr.Row():
+                with gr.Column():
+                    with gr.Accordion("SimpleRAG", open=True):
+                        simple_content = gr.Textbox(
+                            label="SimpleRAG Content", lines=10, max_lines=10
+                        )
+                        simple_response = gr.Markdown(label="SimpleRAG Response")
+                with gr.Column():
+                    with gr.Accordion("VisionRAG", open=True):
+                        vision_gallery = gr.Gallery(label="VisionRAG Images")
+                        vision_response = gr.Markdown(label="VisionRAG Response")
+
+            with gr.Row():
+                with gr.Column():
+                    with gr.Accordion("ColpaliRAG", open=True):
+                        colpali_gallery = gr.Gallery(label="ColpaliRAG Images")
+                        colpali_response = gr.Markdown(label="ColpaliRAG Response")
+                with gr.Column():
+                    with gr.Accordion("HybridColpaliRAG", open=True):
+                        hybrid_gallery = gr.Gallery(label="HybridColpaliRAG Images")
+                        hybrid_response = gr.Markdown(label="HybridColpaliRAG Response")
+
+        with gr.Tab("Settings"):
+            api_key_input = gr.Textbox(label="OpenAI API Key", type="password")
+            update_api_button = gr.Button("Update API Key")
+            api_update_status = gr.Textbox(label="API Update Status")
+
+            simple_table_input = gr.Textbox(
+                label="SimpleRAG Table Name", value="simpleDemo"
+            )
+            vision_table_input = gr.Textbox(
+                label="VisionRAG Table Name", value="visionDemo"
+            )
+            colpali_table_input = gr.Textbox(
+                label="ColpaliRAG Table Name", value="colpaliDemo"
+            )
+            hybrid_table_input = gr.Textbox(
+                label="HybridColpaliRAG Table Name", value="hybridDemo"
+            )
+            update_table_button = gr.Button("Update Table Names")
+            table_update_status = gr.Textbox(label="Table Update Status")
+
+        retrieved_results = gr.State({})
+
+        def update_retrieval_results(query, top_k, sequential):
+            results, timings = retrieve_data(query, top_k, sequential)
+            timing_df = pd.DataFrame(
+                list(timings.items()), columns=["RAG Type", "Time (s)"]
+            )
+            return (
+                results["SimpleRAG"],
+                results["VisionRAG"],
+                results["ColpaliRAG"],
+                results["HybridColpaliRAG"],
+                timing_df,
+                results,
+            )
+
+        retrieve_button.click(
+            update_retrieval_results,
+            inputs=[query_input, top_k_slider, sequential_checkbox],
+            outputs=[
+                simple_content,
+                vision_gallery,
+                colpali_gallery,
+                hybrid_gallery,
+                retrieval_timing,
+                retrieved_results,
+            ],
+        )
+
+        # def update_query_results(query, retrieved_results):
+        #     results = query_data(query, retrieved_results)
+        #     return (
+        #         results["SimpleRAG"]["response"],
+        #         results["VisionRAG"]["response"],
+        #         results["ColpaliRAG"]["response"],
+        #         results["HybridColpaliRAG"]["response"],
+        #     )
+
+        # query_button.click(
+        #     update_query_results,
+        #     inputs=[query_input, retrieved_results],
+        #     outputs=[
+        #         simple_response,
+        #         vision_response,
+        #         colpali_response,
+        #         hybrid_response,
+        #     ],
+        # )
+
+        ingest_button.click(
+            ingest_data,
+            inputs=[pdf_input, use_ocr, chunk_size],
+            outputs=[ingest_output, progress_table],
+        )
+
+        update_api_button.click(
+            update_api_key, inputs=[api_key_input], outputs=api_update_status
+        )
+
+        update_table_button.click(
+            change_table,
+            inputs=[
+                simple_table_input,
+                vision_table_input,
+                colpali_table_input,
+                hybrid_table_input,
+            ],
+            outputs=table_update_status,
+        )
+
+    return demo
+
+
+# Parse command-line arguments
+def parse_args():
+    parser = argparse.ArgumentParser(description="VisionRAG Gradio App")
+    parser.add_argument(
+        "--share", action="store_true", help="Enable Gradio share feature"
+    )
+    return parser.parse_args()
+
+
+# Launch the app
+if __name__ == "__main__":
+    args = parse_args()
+    app = gradio_interface()
+    app.launch(share=args.share)

packages.txt

@@ -0,0 +1 @@
+poppler-utils

requirements.txt

@@ -0,0 +1,20 @@
+git+https://github.com/huggingface/transformers
+torch
+lancedb
+colpali-engine
+pdf2image
+pypdf
+pymupdf
+timm
+einops
+sentence-transformers
+tiktoken
+docling
+pdf2image
+GPUtil
+accelerate==0.30.1
+mteb>=1.12.22
+qwen-vl-utils
+torchvision
+fastapi<0.113.0
+git+https://github.com/adithya-s-k/VARAG