app.py

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 varag.chunking import FixedTokenChunker
from varag.utils import get_model_colpali
import argparse
import spaces
import torch
from docling.document_converter import DocumentConverter

load_dotenv()

# Initialize shared database
shared_db = lancedb.connect("~/rag_demo_db")




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

converter = DocumentConverter()

# 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(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 = []

    @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),
        )

    @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),
        )

    @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),
        )

    @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}"})
    yield IngestResult(
        status_text=f"Total ingestion time: {total_time:.2f} seconds",
        progress_table=pd.DataFrame(progress_data),
    )


@spaces.GPU(duration=120)
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"):
            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"]
            )
            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)