support multimodal pdfchat and switch to marker pdf

Author: CodingWithTim

fastchat/conversation.py

@@ -362,18 +362,39 @@ def update_last_message(self, message: str):
     def to_gradio_chatbot(self):
         """Convert the conversation to gradio chatbot format."""
         from fastchat.serve.vision.image import ImageFormat
+        import re
 
         ret = []
         for i, (role, msg) in enumerate(self.messages[self.offset :]):
             if i % 2 == 0:
                 if type(msg) is tuple:
                     msg, images = msg
-                    image = images[0]  # Only one image on gradio at one time
-                    if image.image_format == ImageFormat.URL:
-                        img_str = f'<img src="{image.url}" alt="user upload image" />'
-                    elif image.image_format == ImageFormat.BYTES:
-                        img_str = f'<img src="data:image/{image.filetype};base64,{image.base64_str}" alt="user upload image" />'
-                    msg = img_str + msg.replace("<image>\n", "").strip()
+
+                    pattern = re.compile("!\[\]\(_page_\d_Figure_\d\.jpeg\)")
+                    embed_locations = pattern.findall(msg)
+
+                    pdfchat = False
+                    for i, embed_str in enumerate(embed_locations):
+                        if i >= len(images):
+                            break
+
+                        image = images[i]
+                        msg = msg.replace(
+                            embed_str,
+                            f'<img src="data:image/{image.filetype};base64,{image.base64_str}" alt="document image" />',
+                        )
+                        pdfchat = True
+
+                    if not pdfchat:
+                        # vision arena only supports one image on gradio at one time
+                        image = images[0]
+                        if image.image_format == ImageFormat.URL:
+                            img_str = (
+                                f'<img src="{image.url}" alt="user upload image" />'
+                            )
+                        elif image.image_format == ImageFormat.BYTES:
+                            img_str = f'<img src="data:image/{image.filetype};base64,{image.base64_str}" alt="user upload image" />'
+                        msg = img_str + msg.replace("<image>\n", "").strip()
 
                 ret.append([msg, None])
             else:

fastchat/serve/gradio_block_arena_vision.py

@@ -232,10 +232,51 @@ def wrap_pdfchat_query(query, document):
     return reformatted_query_context
 
 
-LLAMA_PARSE_MAX_RETRY = 2
-LLAMAPARSE_SUPPORTED_LANGS = {
+# LLAMA_PARSE_MAX_RETRY = 2
+# LLAMAPARSE_SUPPORTED_LANGS = {
+#     "English": "en",
+#     "Chinese": "ch_sim",
+#     "Russian": "ru",
+#     "Spanish": "es",
+#     "Japanese": "ja",
+#     "Korean": "ko",
+#     "French": "fr",
+#     "German": "de",
+#     "Vietnamese": "vi",
+# }
+
+
+# def parse_pdf(file_path):
+#     from llama_parse import LlamaParse
+
+#     assert (
+#         "LLAMA_CLOUD_API_KEY" in os.environ
+#     ), "Make sure to specify LlamaParse API key."
+
+#     for _ in range(LLAMA_PARSE_MAX_RETRY):
+#         try:
+#             documents = LlamaParse(
+#                 result_type="markdown",
+#                 verbose=True,
+#                 languages=list(LLAMAPARSE_SUPPORTED_LANGS.values()),
+#                 accurate_mode=True,
+#             ).load_data(file_path)
+#             assert len(documents) > 0
+#             break
+#         except AssertionError as e:
+#             continue
+
+#     output = "\n".join(
+#         [f"Page {i+1}:\n{doc.text}\n" for i, doc in enumerate(documents)]
+#     )
+
+#     return output
+
+
+PDFPARSE_MAX_RETRY = 2
+PDFPARSE_SUPPORTED_LANGS = {
     "English": "en",
-    "Chinese": "ch_sim",
+    "Chinese": "zh",
     "Russian": "ru",
     "Spanish": "es",
     "Japanese": "ja",
@@ -244,33 +285,36 @@ def wrap_pdfchat_query(query, document):
     "German": "de",
     "Vietnamese": "vi",
 }
+MARKER_PDFPARSE_CONFIG = {
+    "output_format": "markdown",
+    "languages": ",".join(PDFPARSE_SUPPORTED_LANGS.values()),
+}
 
 
 def parse_pdf(file_path):
-    from llama_parse import LlamaParse
-
-    assert (
-        "LLAMA_CLOUD_API_KEY" in os.environ
-    ), "Make sure to specify LlamaParse API key."
+    from marker.config.parser import ConfigParser
+    from marker.models import create_model_dict
+    from marker.converters.pdf import PdfConverter
 
-    for _ in range(LLAMA_PARSE_MAX_RETRY):
+    output_md, output_images = None, None
+    for _ in range(PDFPARSE_MAX_RETRY):
         try:
-            documents = LlamaParse(
-                result_type="markdown",
-                verbose=True,
-                languages=list(LLAMAPARSE_SUPPORTED_LANGS.values()),
-                accurate_mode=True,
-            ).load_data(file_path)
-            assert len(documents) > 0
+            config_parser = ConfigParser(MARKER_PDFPARSE_CONFIG)
+
+            converter = PdfConverter(
+                config=config_parser.generate_config_dict(),
+                artifact_dict=create_model_dict(),
+                processor_list=config_parser.get_processors(),
+                renderer=config_parser.get_renderer(),
+            )
+            rendered = converter(file_path)
+            output_md = rendered.markdown
+            output_images = list(rendered.images.values())
             break
         except AssertionError as e:
             continue
 
-    output = "\n".join(
-        [f"Page {i+1}:\n{doc.text}\n" for i, doc in enumerate(documents)]
-    )
-
-    return output
+    return output_md, output_images
 
 
 def _prepare_text_with_image(state, text, images, csam_flag):
@@ -284,12 +328,26 @@ def _prepare_text_with_image(state, text, images, csam_flag):
     return text
 
 
+# def _prepare_text_with_pdf(text, pdfs):
+#     if len(pdfs) > 0:
+#         document_content = parse_pdf(pdfs[0])
+#         print("Document processed")
+#         text = wrap_pdfchat_query(text, document_content)
+
+#     return text
+
+
 def _prepare_text_with_pdf(text, pdfs):
     if len(pdfs) > 0:
-        document_content = parse_pdf(pdfs[0])
+        parsed_text, imgs = parse_pdf(pdfs[0])
         print("Document processed")
-        text = wrap_pdfchat_query(text, document_content)
+        wrapped_text = wrap_pdfchat_query(text, parsed_text)
 
+        imgs = convert_pdf_images_to_conversation_format(imgs)
+
+        if len(imgs) > 0:
+            return wrapped_text, imgs
+        return wrapped_text
     return text
 
 
@@ -307,6 +365,20 @@ def convert_images_to_conversation_format(images):
     return conv_images
 
 
+def convert_pdf_images_to_conversation_format(images):
+    MAX_NSFW_ENDPOINT_IMAGE_SIZE_IN_MB = 5 / 1.5
+    conv_images = []
+    if len(images) > 0:
+        for img in images:
+            # pdf parser returns a PIL image object instead of path
+            conv_images.append(
+                Image(url="").to_conversation_format(
+                    MAX_NSFW_ENDPOINT_IMAGE_SIZE_IN_MB, pil_img=img
+                )
+            )
+    return conv_images
+
+
 def moderate_input(state, text, all_conv_text, model_list, images, ip):
     text_flagged = moderation_filter(all_conv_text, model_list)
     # flagged = moderation_filter(text, [state.model_name])

fastchat/serve/vision/image.py

@@ -1,7 +1,6 @@
 import base64
 from enum import auto, IntEnum
 from io import BytesIO
-
 from pydantic import BaseModel
 
 
@@ -115,11 +114,25 @@ def convert_url_to_image_bytes(self, max_image_size_mb):
 
         return image_format, img_base64_str
 
-    def to_conversation_format(self, max_image_size_mb):
-        image_format, image_bytes = self.convert_url_to_image_bytes(
-            max_image_size_mb=max_image_size_mb
+    def convert_pil_image_to_image_bytes(self, pil_img, max_image_size_mb):
+        image_format, image_bytes = self.resize_image_and_return_image_in_bytes(
+            pil_img, max_image_size_mb
         )
 
+        img_base64_str = base64.b64encode(image_bytes.getvalue()).decode()
+
+        return image_format, img_base64_str
+
+    def to_conversation_format(self, max_image_size_mb, pil_img=None):
+        if pil_img:
+            image_format, image_bytes = self.convert_pil_image_to_image_bytes(
+                pil_img=pil_img, max_image_size_mb=max_image_size_mb
+            )
+        else:
+            image_format, image_bytes = self.convert_url_to_image_bytes(
+                max_image_size_mb=max_image_size_mb
+            )
+
         self.filetype = image_format
         self.image_format = ImageFormat.BYTES
         self.base64_str = image_bytes