init commit

Dockerfile

@@ -0,0 +1,60 @@
+FROM nvidia/cuda:11.7.1-cudnn8-devel-ubuntu22.04
+ENV DEBIAN_FRONTEND=noninteractive
+RUN apt-get update && \
+    apt-get upgrade -y && \
+    apt-get install -y --no-install-recommends \
+    git \
+    git-lfs \
+    wget \
+    curl \
+    # python build dependencies \
+    build-essential \
+    libssl-dev \
+    zlib1g-dev \
+    libbz2-dev \
+    libreadline-dev \
+    libsqlite3-dev \
+    libncursesw5-dev \
+    xz-utils \
+    tk-dev \
+    libxml2-dev \
+    libxmlsec1-dev \
+    libffi-dev \
+    liblzma-dev \
+    # gradio dependencies \
+    ffmpeg \
+    poppler-utils \
+    && apt-get clean \
+    && rm -rf /var/lib/apt/lists/*
+
+
+RUN useradd -m -u 1000 user
+USER user
+ENV HOME=/home/user \
+    PATH=/home/user/.local/bin:${PATH}
+WORKDIR ${HOME}/app
+
+RUN curl https://pyenv.run | bash
+ENV PATH=${HOME}/.pyenv/shims:${HOME}/.pyenv/bin:${PATH}
+ARG PYTHON_VERSION=3.10.12
+RUN pyenv install ${PYTHON_VERSION} && \
+    pyenv global ${PYTHON_VERSION} && \
+    pyenv rehash && \
+    pip install --no-cache-dir -U pip setuptools wheel && \
+    pip install packaging ninja
+
+COPY --chown=1000 ./requirements.txt /tmp/requirements.txt
+RUN pip install --no-cache-dir --upgrade -r /tmp/requirements.txt && \
+    pip install flash-attn --no-build-isolation
+
+
+
+COPY --chown=1000 . ${HOME}/app
+ENV PYTHONPATH=${HOME}/app \
+    PYTHONUNBUFFERED=1 \
+    GRADIO_ALLOW_FLAGGING=never \
+    GRADIO_NUM_PORTS=1 \
+    GRADIO_SERVER_NAME=0.0.0.0 \
+    GRADIO_THEME=huggingface \
+    SYSTEM=spaces
+CMD ["python", "app.py"]

app.py

@@ -0,0 +1,104 @@
+import gradio as gr
+
+from tool import VisualRAGTool
+
+tool = VisualRAGTool()
+
+def search(query, k, api_key):
+    """Searches for the most relevant pages based on the query."""
+    print("=============== SEARCHING ===============")
+
+    context, answer = tool.search(query, k, api_key)
+
+    context_gallery = [(page.image, page.caption) for page in context]
+
+    print("========================================")
+
+    return gr.Gallery(value=context_gallery, label="Retrieved Documents", height=400, show_label=True, visible=True), answer
+
+def index(files, contextualize_embeds, api_key):
+    """Indexes the uploaded files."""
+    print("=============== INDEXING ===============")
+    
+    indexed_files_num = tool.index(
+        files=files, 
+        contextualize=contextualize_embeds, 
+        api_key=api_key, 
+        overwrite_db=True
+    )
+
+    print("========================================")
+    return gr.Textbox(f"Uploaded and processed {indexed_files_num} pages!"),\
+            gr.Textbox(
+                lines=2,
+                label="Query",
+                show_label=False,
+                placeholder="Enter your prompt here and press Shift+Enter or press the button",
+                interactive=True,
+            )
+
+def show_processing_status():
+    """Updates the upload status."""
+    return gr.Textbox(label="Processing Status", interactive=False, visible=True),\
+            gr.Checkbox(label="Contextualize Embeddings", visible=False),\
+            gr.Textbox(
+                    lines=2,
+                    label="Query",
+                    show_label=False,
+                    placeholder="Enter your prompt here and press Shift+Enter or press the button",
+                    interactive=False,
+                )
+
+with gr.Blocks(
+    theme=gr.themes.Ocean(),
+    title="ColPali Tool Demo",
+) as demo:
+    gr.Markdown("""\
+    # ColPali: Efficient Document Retrieval with Vision Language Models (ColQwen2) 📚   
+    Demo to test the ColPali RAG Tool powered by ColQwen2 (ColPali) on PDF documents.
+    ColPali is implemented from the [ColPali paper](https://arxiv.org/abs/2407.01449).
+
+    This tool allows you to upload PDF files and search for the most relevant pages based on your query.
+    Refresh the page if you change documents!
+
+    ⚠️ This demo uses a model trained exclusively on A4 PDFs in portrait mode, containing english text. Performance is expected to drop for other page formats and languages.
+    Other models will be released with better robustness towards different languages and document formats!
+    """)
+
+    api_key = gr.Textbox(placeholder="Enter your OpenAI KEY here (optional)", label="API key")
+
+    stored_messages = gr.State(value=[])
+
+    gr.Markdown("## 1️⃣ Upload PDFs")
+    gr.Markdown("Upload PDF files to index and search.")
+    with gr.Group():
+        contextualize_embeds = gr.Checkbox(label="Contextualize Embeddings", info="Add images surrouding context as metadata. Generated using gpt-4o-mini. ⚠️ Indexing will be longer!", value=True)
+        upload_files = gr.File(file_types=[".pdf"], file_count="multiple", label="Upload files")
+        processing_status = gr.Textbox(label="Processing Status", interactive=False, visible=False)
+        
+    gr.Markdown("## 2️⃣ Search")
+    gr.Markdown("Ask a question relevant to the documents you uploaded.")
+    with gr.Group():
+        chatbot = gr.Textbox(label="AI Assistant", placeholder="Generated response based on retrieved documents.", lines=6)
+        output_gallery = gr.Gallery(label="Retrieved Documents", height=600, show_label=True, visible=False)
+        with gr.Row(equal_height=True):
+            with gr.Column(scale=4):
+                text_input = gr.Textbox(
+                    lines=2,
+                    label="Query",
+                    show_label=False,
+                    placeholder="Enter your prompt here and press Shift+Enter or press the button",
+                )
+            with gr.Column(scale=1):
+                k = gr.Slider(minimum=1, maximum=10, step=1, value=3, label="Pages to retrieve")
+        search_button = gr.Button("🔍 Search", variant="primary")
+
+    # Define the flow of the demo
+    # upload_files.change(index, inputs=[upload_files, api_key], outputs=[upload_status])
+    upload_files.change(show_processing_status, inputs=[], outputs=[processing_status, contextualize_embeds, text_input])\
+                    .then(index, inputs=[upload_files, contextualize_embeds, api_key], outputs=[processing_status, text_input])
+    text_input.submit(search, inputs=[text_input, k, api_key], outputs=[output_gallery, chatbot])
+    search_button.click(search, inputs=[text_input, k, api_key], outputs=[output_gallery, chatbot])
+
+if __name__ == "__main__":
+    demo.queue(max_size=5).launch(debug=True, server_name="0.0.0.0")

requirements.txt

@@ -0,0 +1,13 @@
+colpali-engine==0.3.8
+pdf2image
+GPUtil
+accelerate==0.30.1
+openai
+gradio
+gradio_client
+tqdm
+Pillow
+pqdm
+smolagents
+pyyaml
+python-dotenv

tool.py

@@ -0,0 +1,344 @@
+from dataclasses import dataclass
+from typing import List, Optional, Tuple
+
+import torch
+from torch.utils.data import DataLoader, Dataset
+
+import base64
+from io import BytesIO
+from PIL import Image
+from pdf2image import convert_from_path
+
+from tqdm import tqdm
+from pqdm.processes import pqdm
+
+from colpali_engine.models import ColQwen2, ColQwen2Processor
+
+from smolagents import Tool, ChatMessage
+
+from utils import query_openai
+
+from dotenv import load_dotenv
+load_dotenv()
+
+def encode_image_to_base64(image):
+    """Encodes a PIL image to a base64 string."""
+    buffered = BytesIO()
+    image.save(buffered, format="JPEG")
+    return base64.b64encode(buffered.getvalue()).decode("utf-8")
+
+DEFAULT_SYSTEM_PROMPT = \
+"""You are a smart assistant designed to answer questions about a PDF document.
+    You are given relevant information in the form of PDF pages preceded by their metadata: document title, page identifier, surrounding context. 
+    Use them to construct a short response to the question, and cite your sources in the following format: (document, page number).
+    If it is not possible to answer using the provided pages, do not attempt to provide an answer and simply say the answer is not present within the documents.
+    Give detailed and extensive answers, only containing info in the pages you are given.
+    You can answer using information contained in plots and figures if necessary.
+    Answer in the same language as the query."""
+
+def _build_query(query, pages):
+    messages = []
+    messages.append({"type": "text", "text": "PDF pages:\n"})
+    for page in pages:
+        capt = page.caption
+        if capt is not None:
+            messages.append({
+                    "type": "text",
+                    "text": capt
+                })
+        messages.append({
+                "type": "image_url",
+                "image_url": {
+                "url": f"data:image/jpeg;base64,{encode_image_to_base64(page.image)}"
+                },
+            })
+    messages.append({"type": "text", "text": f"Query:\n{query}"})
+
+    return messages
+
+def query_openai(query, pages, api_key=None, system_prompt=DEFAULT_SYSTEM_PROMPT, model="gpt-4o-mini") -> ChatMessage:
+    """Calls OpenAI's GPT-4o-mini with the query and image data."""
+    if api_key and api_key.startswith("sk"):
+        try:
+            from openai import OpenAI
+    
+            client = OpenAI(api_key=api_key.strip())
+
+            response = client.chat.completions.create(
+                model=model,
+                messages=[
+                    {
+                        "role": "system",
+                        "content": system_prompt
+                    },
+                    {
+                    "role": "user",
+                    "content": _build_query(query, pages)
+                    }
+                ],
+                max_tokens=500,
+            )
+
+            message = ChatMessage.from_dict(
+                response.choices[0].message.model_dump(include={"role", "content", "tool_calls"})
+            )
+            message.raw = response
+
+            return message
+        
+        except Exception as e:
+            return "OpenAI API connection failure. Verify the provided key is correct (sk-***)."
+        
+    return "Enter your OpenAI API key to get a custom response"
+
+DEFAULT_CONTEXT_PROMPT = \
+"""You are a smart assistant designed to extract context of PDF pages.
+    Give concise answers, only containing info in the pages you are given.
+    You can answer using information contained in plots and figures if necessary."""
+
+RAG_SYSTEM_PROMPT = \
+""" You are a smart assistant designed to answer questions about a PDF document.
+
+    You are given relevant information in the form of PDF pages preceded by their metadata: document title, page identifier, surrounding context.
+    Use them to construct a response to the question, and cite your sources.
+    Use the following citation format:
+    "Some information from a first document [1, p.Page Number]. Some information from the same first document but at a different page [1, p.Page Number]. Some more information from another document [2, p.Page Number].
+    ...
+    Sources:
+    [1] Document Title
+    [2] Another Document Title"
+
+    You can answer using information contained in plots and figures if necessary.
+    If it is not possible to answer using the provided pages, do not attempt to provide an answer and simply say the answer is not present within the documents.
+    Give detailed answers, only containing info in the pages you are given.
+    Answer in the same language as the query."""
+
+@dataclass
+class Metadata:
+    doc_title: str
+    page_id: int
+    context: Optional[str] = None
+
+    def __str__(self):
+        return f"Document: {self.doc_title}, Page ID: {self.page_id}, Context: {self.context}"
+    
+@dataclass
+class Page:
+    image: Image.Image
+    metadata: Optional[Metadata] = None
+
+    @property
+    def caption(self):
+        if self.metadata is None:
+            return None
+        return f"Document: {self.metadata.doc_title}, Context: {self.metadata.context}"
+
+class VisualRAGTool(Tool):
+    name = "visual_rag"
+    description = """Performs a RAG query on your internal PDF documents and returns the generated text response."""
+    inputs = {
+        "query": {
+            "type": "string",
+            "description": "The query to perform. This should be semantically close to your target documents.",
+        },
+        "k": {
+            "type": "number",
+            "description": "The number of documents to retrieve.",
+            "default": 1,
+            "nullable": True,
+        },
+        "api_key": {
+            "type": "string",
+            "description": "The OpenAI API key to use for the query. If not provided, the key will be taken from the OPENAI_KEY environment variable.",
+            "nullable": True,
+        }
+    }
+    output_type = "string"
+
+    def _init_models(self, model_name: str) -> None:
+        self.device = "cuda" if torch.cuda.is_available() else "cpu"
+        self.model = ColQwen2.from_pretrained(
+                model_name,
+                torch_dtype=torch.bfloat16,
+                device_map="auto",
+                attn_implementation="flash_attention_2"
+            ).eval()
+        self.processor = ColQwen2Processor.from_pretrained(model_name)
+
+    def __init__(self, model_name: str = "vidore/colqwen2-v1.0", api_key: str = None, files: List[str] = None, **kwargs):
+        super().__init__(**kwargs)
+        self.model_name = model_name
+        self.api_key = api_key
+
+        self.embds = []
+        self.pages = []
+
+        self.files = files
+
+        self._init_models(self.model_name)
+
+        self.is_initialized = False 
+    
+    def setup(self):
+        """
+        Overwrite this method here for any operation that is expensive and needs to be executed before you start using
+        your tool. Such as loading a big model.
+        """
+        if self.files:
+            _ = self.index(self.files, self.api_key)
+
+        self.is_initialized = True
+
+    def _extract_contexts(self, images, api_key, window=10) -> List[str]:
+        """Extracts context from images."""
+        try:
+            args = [
+                {
+                    'query': "Give the general context about these pages. Give the context in the same language as the documents.",
+                    'pages': [Page(image=im) for im in images[max(i-window+1, 0):i+1]],
+                    'api_key': api_key,
+                    'system_prompt': DEFAULT_CONTEXT_PROMPT
+                } for i in range(0, len(images), window)
+            ]
+            window_contexts = pqdm(args, query_openai, n_jobs=8, argument_type='kwargs')
+
+            # code sequentially ftm with tqdm
+            # query = "Give the general context about these pages. Give the context in the same language as the documents."
+            # window_contexts = [query_openai(query, [Page(image=im) for im in images[max(i-window+1, 0):i+1]], api_key, DEFAULT_CONTEXT_PROMPT)\
+            #                      for i in tqdm(range(0, len(images), window))]
+
+            contexts = []
+            for i in range(len(images)):
+                context = window_contexts[i//window].content
+                contexts.append(context)
+
+        except Exception as e:
+            print(f"Error extracting contexts: {e}")
+            contexts = [None for _ in range(len(images))]
+
+        # Ensure that the number of contexts is equal to the number of images
+        assert len(contexts) == len(images)
+        
+        return contexts
+
+    def _process_file(self, file: str, contextualize: bool = True, api_key: str = None, window: int = 10) -> List[Page]:
+        """Converts a file to images and extracts metadata."""
+        title = file.split("/")[-1]
+        images = convert_from_path(file, thread_count=4)
+        if contextualize and api_key:
+            contexts = self._extract_contexts(images, api_key, window=window)
+        else:
+            contexts = [None for _ in range(len(images))]
+        metadatas = [Metadata(doc_title=title, page_id=i, context=contexts[i]) for i in range(len(images))]
+
+        return [Page(image=img, metadata=metadata) for img, metadata in zip(images, metadatas)]
+
+    def preprocess(self, files: List[str], contextualize: bool = True, api_key: str = None, window: int = 10) -> List[Page]:
+        """Preprocesses the files and extracts metadata."""
+        pages = [page for file in files for page in self._process_file(file, contextualize=contextualize, api_key=api_key, window=window)]     
+
+        print(f"Example metadata:\n{pages[0].metadata.context}")
+
+        return pages
+    
+    def _embed_images(self, pages: List[Page]) -> List[torch.Tensor]:
+        """Embeds the images using the model."""
+        """Example script to run inference with ColPali (ColQwen2)"""
+        # run inference - docs
+        dataloader = DataLoader(
+            pages,
+            batch_size=4,
+            shuffle=False,
+            collate_fn=lambda x: self.processor.process_images([p.image for p in x]).to(self.device),
+        )
+
+        embds = []
+
+        for batch_doc in tqdm(dataloader):
+            with torch.no_grad():
+                batch_doc = {k: v.to(self.device) for k, v in batch_doc.items()}
+                embeddings_doc = self.model(**batch_doc)
+            embds.extend(list(torch.unbind(embeddings_doc.to("cpu"))))
+
+        return embds
+
+    def index(self, files: List[str], contextualize: bool = True, api_key: str = None, overwrite_db: bool = False) -> int:
+        print("Converting files...")
+        # Convert files to images and extract metadata
+        pgs = self.preprocess(files, contextualize=contextualize, api_key=api_key or self.api_key)
+
+        # Embed the images
+        embds = self._embed_images(pgs)
+
+        # Overwrite the database if necessary
+        if overwrite_db:
+            self.pages = []
+            self.embds = []
+        
+        # Extend the pages
+        self.pages.extend(pgs)
+
+        # Extend the datasets
+        self.embds.extend(embds)
+
+        print(f"Extracted and indexed {len(pgs)} images from {len(files)} files.")
+        
+        return len(embds)
+
+    def retrieve(self, query: str, k: int) -> List[Page]:
+        """Retrieve the top k documents based on the query."""
+        k = min(k, len(self.embds))
+
+        qs = []
+        with torch.no_grad():
+            batch_query = self.processor.process_queries([query]).to(self.model.device)
+            embeddings_query = self.model(**batch_query)
+            qs.extend(list(torch.unbind(embeddings_query.to("cpu"))))
+
+        # Run scoring
+        scores = self.processor.score(qs, self.embds, device=self.device)[0]
+        top_k_idx = scores.topk(k).indices.tolist()
+
+        print("Top Scores:")
+        [print(f'Page {self.pages[idx].metadata.page_id}: {scores[idx]}') for idx in top_k_idx]
+
+        # Get the top k results
+        results = [self.pages[idx] for idx in top_k_idx]
+
+        return results
+        
+    def generate_answer(self, query: str, docs: List[Page], api_key: str = None) -> ChatMessage:
+        result = query_openai(query, docs, api_key or self.api_key, system_prompt=RAG_SYSTEM_PROMPT)
+        return result
+
+    def search(self, query: str, k: int = 1, api_key: str = None) -> Tuple[list, str]:
+        print(f"Searching for query: {query}")
+    
+        # Retrieve the top k documents
+        context = self.retrieve(query, k)
+
+        # Generate response from GPT-4o-mini
+        rag_answer = self.generate_answer(
+            query=query, 
+            docs=context, 
+            api_key=api_key
+        )
+
+        return context, rag_answer.content
+
+    def forward(self, query: str, k: int = 1, api_key: str = None) -> str:
+        assert isinstance(query, str), "Your search query must be a string"
+
+        # Online indexing
+        # if files:
+        #     _ = self.index(files, api_key)
+
+        # Retrieve the top k documents and generate response
+        _, rag_answer = self.search(
+            query=query, 
+            files=None, 
+            k=k, 
+            api_key=api_key
+        )
+
+        return rag_answer