Support doc loading, detector and retriever

app.py

@@ -1,15 +1,22 @@
 import base64
 import hashlib
 import io
+import json
 import os
+import tempfile
+from collections import OrderedDict as CollectionsOrderedDict
 from pathlib import Path
 from threading import Thread
-from typing import Iterator, Optional, List, Union
+from typing import Iterator, Optional, List, Union, OrderedDict
 
+import fitz
 import gradio as gr
+import requests
 import spaces
 import torch
 from PIL import Image
+from colpali_engine import ColPali, ColPaliProcessor
+from huggingface_hub import hf_hub_download
 from pydantic import BaseModel
 from qwen_vl_utils import process_vision_info
 from swift.llm import (
@@ -20,6 +27,7 @@ from swift.llm import (
     inference,
     inference_stream,
 )
+from tqdm import tqdm
 from transformers import (
     Qwen2VLForConditionalGeneration,
     PreTrainedTokenizer,
@@ -27,6 +35,8 @@ from transformers import (
     TextIteratorStreamer,
     AutoTokenizer,
 )
+from ultralytics import YOLO
+from ultralytics.engine.results import Results
 
 MAX_MAX_NEW_TOKENS = 2048
 DEFAULT_MAX_NEW_TOKENS = 1024
@@ -49,6 +59,264 @@ this demo is governed by the original [license](https://huggingface.co/meta-llam
 """
 
 
+class MultimodalSample(BaseModel):
+    question: str
+    answer: str
+    category: str
+    evidence_pages: List[int] = []
+    raw_output: str = ""
+    pred: str = ""
+    source: str = ""
+    annotator: str = ""
+    generator: str = ""
+    retrieved_pages: List[int] = []
+
+
+class MultimodalObject(BaseModel):
+    id: str = ""
+    page: int = 0
+    text: str = ""
+    image_string: str = ""
+    snippet: str = ""
+    score: float = 0.0
+    source: str = ""
+    category: str = ""
+
+    def get_image(self) -> Optional[Image.Image]:
+        if self.image_string:
+            return convert_text_to_image(self.image_string)
+
+    @classmethod
+    def from_image(cls, image: Image.Image, **kwargs):
+        return cls(image_string=convert_image_to_text(image), **kwargs)
+
+
+class ObjectDetector(BaseModel, arbitrary_types_allowed=True):
+    def run(self, image: Image.Image) -> List[MultimodalObject]:
+        raise NotImplementedError()
+
+
+class YoloDetector(ObjectDetector):
+    repo_id: str = "DILHTWD/documentlayoutsegmentation_YOLOv8_ondoclaynet"
+    filename: str = "yolov8x-doclaynet-epoch64-imgsz640-initiallr1e-4-finallr1e-5.pt"
+    local_dir: str = "data/yolo"
+    client: Optional[YOLO] = None
+
+    def load(self):
+        if self.client is None:
+            if not Path(self.local_dir, self.filename).exists():
+                hf_hub_download(
+                    repo_id=self.repo_id,
+                    filename=self.filename,
+                    local_dir=self.local_dir,
+                )
+            self.client = YOLO(Path(self.local_dir, self.filename))
+
+    def save_image(self, image: Image.Image) -> str:
+        text = convert_image_to_text(image)
+        hash_id = hashlib.md5(text.encode()).hexdigest()
+        path = Path(self.local_dir, f"{hash_id}.png")
+        image.save(path)
+        return str(path)
+
+    @staticmethod
+    def extract_subimage(image: Image.Image, box: List[float]) -> Image.Image:
+        return image.crop((round(box[0]), round(box[1]), round(box[2]), round(box[3])))
+
+    def run(self, image: Image.Image) -> List[MultimodalObject]:
+        self.load()
+        path = self.save_image(image)
+        results: List[Results] = self.client(source=[path])
+        assert len(results) == 1
+        objects = []
+
+        for i, label_id in enumerate(results[0].boxes.cls):
+            label = results[0].names[label_id.item()]
+            score = results[0].boxes.conf[i].item()
+            box: List[float] = results[0].boxes.xyxy[i].tolist()
+            subimage = self.extract_subimage(image, box)
+            objects.append(
+                MultimodalObject(
+                    image_string=convert_image_to_text(subimage),
+                    category=label,
+                    score=score,
+                )
+            )
+
+        return objects
+
+
+class MultimodalPage(BaseModel):
+    number: int
+    objects: List[MultimodalObject]
+    text: str
+    image_string: str
+    source: str
+    score: float = 0.0
+
+    def get_tables_and_figures(self) -> List[MultimodalObject]:
+        return [o for o in self.objects if o.category in ["Table", "Picture"]]
+
+    def get_full_image(self) -> Image.Image:
+        return convert_text_to_image(self.image_string)
+
+    @classmethod
+    def from_text(cls, text: str):
+        return MultimodalPage(
+            text=text, number=0, objects=[], image_string="", source=""
+        )
+
+    @classmethod
+    def from_image(cls, image: Image.Image):
+        return MultimodalPage(
+            image_string=convert_image_to_text(image),
+            number=0,
+            objects=[],
+            text="",
+            source="",
+        )
+
+
+class MultimodalDocument(BaseModel):
+    pages: List[MultimodalPage]
+
+    def get_page(self, i: int) -> MultimodalPage:
+        pages = [p for p in self.pages if p.number == i]
+        assert len(pages) == 1
+        return pages[0]
+
+    @classmethod
+    def load_from_pdf(cls, path: str, dpi: int = 150, detector: ObjectDetector = None):
+        # Each page as an image (with optional extracted text)
+        doc = fitz.open(path)
+        pages = []
+
+        for i, page in enumerate(tqdm(doc.pages(), desc=path)):
+            text = page.get_text()
+            zoom = dpi / 72  # 72 is the default DPI
+            matrix = fitz.Matrix(zoom, zoom)
+            pix = page.get_pixmap(matrix=matrix)
+            image = Image.frombytes("RGB", (pix.width, pix.height), pix.samples)
+
+            objects = []
+            if detector:
+                objects = detector.run(image)
+            for o in objects:
+                o.page, o.source = i + 1, path
+
+            pages.append(
+                MultimodalPage(
+                    number=i + 1,
+                    objects=objects,
+                    text=text,
+                    image_string=convert_image_to_text(image),
+                    source=path,
+                )
+            )
+
+        return cls(pages=pages)
+
+    @classmethod
+    def load(cls, path: str):
+        pages = []
+        with open(path) as f:
+            for line in f:
+                pages.append(MultimodalPage(**json.loads(line)))
+        return cls(pages=pages)
+
+    def save(self, path: str):
+        Path(path).parent.mkdir(exist_ok=True, parents=True)
+        with open(path, "w") as f:
+            for o in self.pages:
+                print(o.model_dump_json(), file=f)
+
+    def get_domain(self) -> str:
+        filename = Path(self.pages[0].source).name
+        if filename.startswith("NYSE"):
+            return "Financial<br>Report"
+        elif filename[:4].isdigit() and filename[4] == "." and filename[5].isdigit():
+            return "Academic<br>Paper"
+        else:
+            return "Technical<br>Manuals"
+
+
+class MultimodalRetriever(BaseModel, arbitrary_types_allowed=True):
+    def run(self, query: str, doc: MultimodalDocument) -> MultimodalDocument:
+        raise NotImplementedError
+
+    @staticmethod
+    def get_top_pages(doc: MultimodalDocument, k: int) -> List[int]:
+        # Get top-k in terms of score but maintain the original order
+        doc = doc.copy(deep=True)
+        pages = sorted(doc.pages, key=lambda x: x.score, reverse=True)
+        threshold = pages[:k][-1].score
+        return [p.number for p in doc.pages if p.score >= threshold]
+
+
+class ColpaliRetriever(MultimodalRetriever):
+    path: str = "vidore/colpali-v1.2"
+    model: Optional[ColPali] = None
+    processor: Optional[ColPaliProcessor] = None
+    device: str = "cuda"
+    cache: OrderedDict[str, torch.Tensor] = CollectionsOrderedDict()
+
+    def load(self):
+        if self.model is None:
+            self.model = ColPali.from_pretrained(
+                self.path, torch_dtype=torch.bfloat16, device_map=self.device
+            )
+            self.model = self.model.eval()
+            self.processor = ColPaliProcessor.from_pretrained(self.path)
+
+    def encode_document(self, doc: MultimodalDocument) -> torch.Tensor:
+        hash_id = hashlib.md5(doc.json().encode()).hexdigest()
+        if len(self.cache) > 100:
+            self.cache.popitem(last=False)
+        if hash_id not in self.cache:
+            images = [page.get_full_image() for page in doc.pages]
+            batch_size = 8
+
+            ds: List[torch.Tensor] = []
+            for i in tqdm(range(0, len(images), batch_size), desc="Encoding document"):
+                batch = self.processor.process_images(images[i : i + batch_size])
+                with torch.no_grad():
+                    # noinspection PyTypeChecker
+                    ds.append(self.model(**batch.to(self.device)).cpu())
+
+            lengths = [x.shape[1] for x in ds]
+            if len(set(lengths)) != 1:
+                print("Warning: Inconsistent lengths from colqwen", set(lengths))
+                assert "colqwen" in self.path
+                for i, x in enumerate(ds):
+                    ds[i] = x[:, : min(lengths), :]
+            self.cache[hash_id] = torch.cat(ds)
+        return self.cache[hash_id]
+
+    def run(self, query: str, doc: MultimodalDocument) -> MultimodalDocument:
+        doc = doc.copy(deep=True)
+        self.load()
+        ds = self.encode_document(doc)
+        with torch.no_grad():
+            # noinspection PyTypeChecker
+            qs = self.model(**self.processor.process_queries([query]).to(self.device))
+
+        # noinspection PyTypeChecker
+        scores = self.processor.score_multi_vector(qs.cpu(), ds).squeeze()
+        assert len(scores) == len(doc.pages)
+        for i, page in enumerate(doc.pages):
+            page.score = scores[i].item()
+
+        return doc
+
+
+class DummyRetriever(MultimodalRetriever):
+    def run(self, query: str, doc: MultimodalDocument) -> MultimodalDocument:
+        doc = doc.copy(deep=True)
+        for i, page in enumerate(doc.pages):
+            page.score = i
+        return doc
+
+
 def convert_image_to_text(image: Image) -> str:
     # This is also how OpenAI encodes images: https://platform.openai.com/docs/guides/vision
     with io.BytesIO() as output:
@@ -272,6 +540,24 @@ class QwenModel(EvalModel):
             yield "".join(outputs)
 
 
+class DummyModel(EvalModel):
+    engine: str = ""
+
+    def run(self, inputs: List[Union[str, Image.Image]]) -> str:
+        return " ".join(inputs)
+
+    def run_stream(self, inputs: List[Union[str, Image.Image]]) -> Iterator[str]:
+        assert self is not None
+        text = " ".join([x for x in inputs if isinstance(x, str)])
+        num_images = sum(1 for x in inputs if isinstance(x, Image.Image))
+        tokens = f"Hello this is your message: {text}, images: {num_images}".split()
+        for i in range(len(tokens)):
+            yield " ".join(tokens[: i + 1])
+            import time
+
+            time.sleep(0.05)
+
+
 if not torch.cuda.is_available():
     DESCRIPTION += "\n<p>Running on CPU 🥶 This demo does not work on CPU.</p>"
 
@@ -279,86 +565,90 @@ if not torch.cuda.is_available():
 if torch.cuda.is_available():
     model = QwenModel()
     model.load()
+    detect_model = YoloDetector()
+    detect_model.load()
+    retriever = ColpaliRetriever()
+    retriever.load()
+else:
+    model = DummyModel()
+    detect_model = None
+    retriever = DummyRetriever()
+
+
+def get_file_path(file: gr.File = None, url: str = None) -> Optional[str]:
+    if file is not None:
+        # noinspection PyUnresolvedReferences
+        return file.name
+
+    if url is not None:
+        response = requests.get(url)
+        response.raise_for_status()
+        save_path = Path(tempfile.mkdtemp(), url.split("/")[-1])
+
+        if "application/pdf" in response.headers.get("Content-Type", ""):
+            # Open the file in binary write mode
+            with open(save_path, "wb") as file:
+                file.write(response.content)
+        return str(save_path)
 
 
 @spaces.GPU
 def generate(
-    message: str,
-    chat_history: list[dict],
-    system_prompt: str = "",
-    max_new_tokens: int = 1024,
-    temperature: float = 0.6,
-    top_p: float = 0.9,
-    top_k: int = 50,
-    repetition_penalty: float = 1.2,
+    query: str, file: gr.File = None, url: str = None, top_k=5
 ) -> Iterator[str]:
-    for text in model.run_stream([message]):
-        yield text
+    sample = MultimodalSample(question=query, answer="", category="")
+    path = get_file_path(file, url)
+
+    if path is not None:
+        doc = MultimodalDocument.load_from_pdf(path, detector=detect_model)
+        output = retriever.run(sample.question, doc)
+        sorted_pages = sorted(output.pages, key=lambda p: p.score, reverse=True)
+        sample.retrieved_pages = sorted([p.number for p in sorted_pages][:top_k])
+
+        context = []
+        for p in doc.pages:
+            if p.number in sample.retrieved_pages:
+                if p.text:
+                    context.append(p.text)
+                context.extend(o.get_image() for o in p.get_tables_and_figures())
+
+        inputs = [
+            "Context:",
+            *context,
+            f"Answer the following question in 200 words or less: {sample.question}",
+        ]
+    else:
+        inputs = [
+            "Context:",
+            f"Answer the following question in 200 words or less: {sample.question}",
+        ]
 
+    for text in model.run_stream(inputs):
+        yield text
 
-chat_interface = gr.ChatInterface(
-    fn=generate,
-    additional_inputs=[
-        gr.Textbox(label="System prompt", lines=6),
-        gr.Slider(
-            label="Max new tokens",
-            minimum=1,
-            maximum=MAX_MAX_NEW_TOKENS,
-            step=1,
-            value=DEFAULT_MAX_NEW_TOKENS,
-        ),
-        gr.Slider(
-            label="Temperature",
-            minimum=0.1,
-            maximum=4.0,
-            step=0.1,
-            value=0.6,
-        ),
-        gr.Slider(
-            label="Top-p (nucleus sampling)",
-            minimum=0.05,
-            maximum=1.0,
-            step=0.05,
-            value=0.9,
-        ),
-        gr.Slider(
-            label="Top-k",
-            minimum=1,
-            maximum=1000,
-            step=1,
-            value=50,
-        ),
-        gr.Slider(
-            label="Repetition penalty",
-            minimum=1.0,
-            maximum=2.0,
-            step=0.05,
-            value=1.2,
-        ),
-    ],
-    stop_btn=None,
-    examples=[
-        [
-            "Roger has 5 tennis balls. He buys 2 more cans of tennis balls. Each can has 3 tennis balls. How many tennis balls does he have now?"
-        ],
-        [
-            "Natalia sold clips to 48 of her friends in April, and then she sold half as many clips in May. How many clips did Natalia sell altogether in April and May?"
-        ],
-        [
-            "Weng earns $12 an hour for babysitting. Yesterday, she just did 50 minutes of babysitting. How much did she earn?"
-        ],
-    ],
-    cache_examples=False,
-    type="messages",
-)
 
 with gr.Blocks(css_paths="style.css", fill_height=True) as demo:
     gr.Markdown(DESCRIPTION)
     gr.DuplicateButton(
         value="Duplicate Space for private use", elem_id="duplicate-button"
     )
-    chat_interface.render()
+
+    with gr.Row():
+        pdf_upload = gr.File(label="Upload PDF (optional)", file_types=[".pdf"])
+        with gr.Column():
+            url_input = gr.Textbox(label="Enter PDF URL (optional)")
+            text_input = gr.Textbox(label="Enter your message", lines=3)
+
+    submit_button = gr.Button("Submit")
+    result = gr.Textbox(label="Response", lines=10)
+
+    submit_button.click(
+        generate, inputs=[text_input, pdf_upload, url_input], outputs=result
+    )
+
     gr.Markdown(LICENSE)
 
+demo.launch()
+
 if __name__ == "__main__":
     demo.queue(max_size=20).launch()