Upload 24 files

.idea/.gitignore

@@ -0,0 +1,8 @@
+# Default ignored files
+/shelf/
+/workspace.xml
+# Editor-based HTTP Client requests
+/httpRequests/
+# Datasource local storage ignored files
+/dataSources/
+/dataSources.local.xml

README.md

@@ -0,0 +1,159 @@
+# Sparrow Parse
+
+## Description
+
+This module implements Sparrow Parse [library](https://pypi.org/project/sparrow-parse/) library with helpful methods for data pre-processing, parsing and extracting information. Library relies on Visual LLM functionality, Table Transformers and is part of Sparrow. Check main [README](https://github.com/katanaml/sparrow)
+
+## Install
+
+```
+pip install sparrow-parse
+```
+
+## Parsing and extraction
+
+### Sparrow Parse VL (vision-language model) extractor with local MLX or Hugging Face Cloud GPU infra
+
+```
+# run locally: python -m sparrow_parse.extractors.vllm_extractor
+
+from sparrow_parse.vllm.inference_factory import InferenceFactory
+from sparrow_parse.extractors.vllm_extractor import VLLMExtractor
+
+extractor = VLLMExtractor()
+
+config = {
+    "method": "mlx",  # Could be 'huggingface', 'mlx' or 'local_gpu'
+    "model_name": "mlx-community/Qwen2-VL-72B-Instruct-4bit",
+}
+
+# Use the factory to get the correct instance
+factory = InferenceFactory(config)
+model_inference_instance = factory.get_inference_instance()
+
+input_data = [
+    {
+        "file_path": "/Users/andrejb/Work/katana-git/sparrow/sparrow-ml/llm/data/bonds_table.jpg",
+        "text_input": "retrieve all data. return response in JSON format"
+    }
+]
+
+# Now you can run inference without knowing which implementation is used
+results_array, num_pages = extractor.run_inference(model_inference_instance, input_data, tables_only=False,
+                                 generic_query=False,
+                                 crop_size=80,
+                                 debug_dir=None,
+                                 debug=True,
+                                 mode=None)
+
+for i, result in enumerate(results_array):
+    print(f"Result for page {i + 1}:", result)
+print(f"Number of pages: {num_pages}")
+```
+
+Use `tables_only=True` if you want to extract only tables.
+
+Use `crop_size=N` (where `N` is an integer) to crop N pixels from all borders of the input images. This can be helpful for removing unwanted borders or frame artifacts from scanned documents.
+
+Use `mode="static"` if you want to simulate LLM call, without executing LLM backend.
+
+Method `run_inference` will return results and number of pages processed.
+
+To run with Hugging Face backend use these config values:
+
+```
+config = {
+    "method": "huggingface",  # Could be 'huggingface' or 'local_gpu'
+    "hf_space": "katanaml/sparrow-qwen2-vl-7b",
+    "hf_token": os.getenv('HF_TOKEN'),
+}
+```
+
+Note: GPU backend `katanaml/sparrow-qwen2-vl-7b` is private, to be able to run below command, you need to create your own backend on Hugging Face space using [code](https://github.com/katanaml/sparrow/tree/main/sparrow-data/parse/sparrow_parse/vllm/infra/qwen2_vl_7b) from Sparrow Parse.
+
+## PDF pre-processing
+
+```
+from sparrow_parse.extractor.pdf_optimizer import PDFOptimizer
+
+pdf_optimizer = PDFOptimizer()
+
+num_pages, output_files, temp_dir = pdf_optimizer.split_pdf_to_pages(file_path,
+                                                                     debug_dir,
+                                                                     convert_to_images)
+
+```
+
+Example:
+
+*file_path* - `/data/invoice_1.pdf`
+
+*debug_dir* - set to not `None`, for debug purposes only
+
+*convert_to_images* - default `False`, to split into PDF files
+
+## Image cropping
+
+```
+from sparrow_parse.helpers.image_optimizer import ImageOptimizer
+
+image_optimizer = ImageOptimizer()
+
+cropped_file_path = image_optimizer.crop_image_borders(file_path, temp_dir, debug_dir, crop_size)
+```
+
+Example:
+
+*file_path* - `/data/invoice_1.jpg`
+
+*temp_dir* - directory to store cropped files
+
+*debug_dir* - set to not `None`, for debug purposes only
+
+*crop_size* - Number of pixels to crop from each border
+
+## Library build
+
+Create Python virtual environment
+
+```
+python -m venv .env_sparrow_parse
+```
+
+Install Python libraries
+
+```
+pip install -r requirements.txt
+```
+
+Build package
+
+```
+pip install setuptools wheel
+python setup.py sdist bdist_wheel
+```
+
+Upload to PyPI
+
+```
+pip install twine
+twine upload dist/*
+```
+
+## Commercial usage
+
+Sparrow is available under the GPL 3.0 license, promoting freedom to use, modify, and distribute the software while ensuring any modifications remain open source under the same license. This aligns with our commitment to supporting the open-source community and fostering collaboration.
+
+Additionally, we recognize the diverse needs of organizations, including small to medium-sized enterprises (SMEs). Therefore, Sparrow is also offered for free commercial use to organizations with gross revenue below $5 million USD in the past 12 months, enabling them to leverage Sparrow without the financial burden often associated with high-quality software solutions.
+
+For businesses that exceed this revenue threshold or require usage terms not accommodated by the GPL 3.0 license—such as integrating Sparrow into proprietary software without the obligation to disclose source code modifications—we offer dual licensing options. Dual licensing allows Sparrow to be used under a separate proprietary license, offering greater flexibility for commercial applications and proprietary integrations. This model supports both the project's sustainability and the business's needs for confidentiality and customization.
+
+If your organization is seeking to utilize Sparrow under a proprietary license, or if you are interested in custom workflows, consulting services, or dedicated support and maintenance options, please contact us at abaranovskis@redsamuraiconsulting.com. We're here to provide tailored solutions that meet your unique requirements, ensuring you can maximize the benefits of Sparrow for your projects and workflows.
+
+## Author
+
+[Katana ML](https://katanaml.io), [Andrej Baranovskij](https://github.com/abaranovskis-redsamurai)
+
+## License
+
+Licensed under the GPL 3.0. Copyright 2020-2025 Katana ML, Andrej Baranovskij. [Copy of the license](https://github.com/katanaml/sparrow/blob/main/LICENSE).

requirements.txt

@@ -0,0 +1,21 @@
+rich
+# mlx-vlm==0.1.12 works with transformers from source only
+# git+https://github.com/huggingface/transformers.git
+# transformers==4.48.2
+torchvision==0.21.0
+torch==2.6.0
+sentence-transformers==3.3.1
+numpy==2.1.3
+pypdf==5.2.0
+gradio_client
+pdf2image
+# mlx==0.22.0; sys_platform == "darwin" and platform_machine == "arm64"
+mlx>=0.22.0; sys_platform == "darwin" and platform_machine == "arm64"
+mlx-vlm==0.1.12; sys_platform == "darwin" and platform_machine == "arm64"
+
+
+# Force reinstall:
+# pip install --force-reinstall -r requirements.txt
+
+# For pdf2image, additional step is required:
+# brew install poppler

setup.py

@@ -0,0 +1,37 @@
+from setuptools import setup, find_packages
+
+with open("README.md", "r", encoding="utf-8") as fh:
+    long_description = fh.read()
+
+with open("requirements.txt", "r", encoding="utf-8") as fh:
+    requirements = fh.read().splitlines()
+
+setup(
+    name="sparrow-parse",
+    version="0.5.4",
+    author="Andrej Baranovskij",
+    author_email="andrejus.baranovskis@gmail.com",
+    description="Sparrow Parse is a Python package (part of Sparrow) for parsing and extracting information from documents.",
+    long_description=long_description,
+    long_description_content_type="text/markdown",
+    url="https://github.com/katanaml/sparrow/tree/main/sparrow-data/parse",
+    project_urls={
+        "Homepage": "https://github.com/katanaml/sparrow/tree/main/sparrow-data/parse",
+        "Repository": "https://github.com/katanaml/sparrow",
+    },
+    classifiers=[
+        "Operating System :: OS Independent",
+        "License :: OSI Approved :: GNU General Public License v3 (GPLv3)",
+        "Topic :: Software Development",
+        "Programming Language :: Python :: 3.10",
+    ],
+    entry_points={
+        'console_scripts': [
+            'sparrow-parse=sparrow_parse:main',
+        ],
+    },
+    keywords="llm, vllm, ocr, vision",
+    packages=find_packages(),
+    python_requires='>=3.10',
+    install_requires=requirements,
+)

sparrow_parse/__init__.py

@@ -0,0 +1 @@
+__version__ = '0.5.4'

sparrow_parse/__main__.py

@@ -0,0 +1,6 @@
+def main():
+    print('Sparrow Parse is a Python package for parsing and extracting information from documents.')
+
+
+if __name__ == "__main__":
+    main()

sparrow_parse/extractors/vllm_extractor.py

@@ -0,0 +1,229 @@
+import json
+from sparrow_parse.vllm.inference_factory import InferenceFactory
+from sparrow_parse.helpers.pdf_optimizer import PDFOptimizer
+from sparrow_parse.helpers.image_optimizer import ImageOptimizer
+from sparrow_parse.processors.table_structure_processor import TableDetector
+from rich import print
+import os
+import tempfile
+import shutil
+
+
+class VLLMExtractor(object):
+    def __init__(self):
+        pass
+
+    def run_inference(self, model_inference_instance, input_data, tables_only=False,
+                      generic_query=False, crop_size=None, debug_dir=None, debug=False, mode=None):
+        """
+        Main entry point for processing input data using a model inference instance.
+        Handles generic queries, PDFs, and table extraction.
+        """
+        if generic_query:
+            input_data[0]["text_input"] = "retrieve document data. return response in JSON format"
+
+        if debug:
+            print("Input data:", input_data)
+
+        file_path = input_data[0]["file_path"]
+        if self.is_pdf(file_path):
+            return self._process_pdf(model_inference_instance, input_data, tables_only, crop_size, debug, debug_dir, mode)
+
+        return self._process_non_pdf(model_inference_instance, input_data, tables_only, crop_size, debug, debug_dir)
+
+
+    def _process_pdf(self, model_inference_instance, input_data, tables_only, crop_size, debug, debug_dir, mode):
+        """
+        Handles processing and inference for PDF files, including page splitting and optional table extraction.
+        """
+        pdf_optimizer = PDFOptimizer()
+        num_pages, output_files, temp_dir = pdf_optimizer.split_pdf_to_pages(input_data[0]["file_path"],
+                                                                             debug_dir, convert_to_images=True)
+
+        results = self._process_pages(model_inference_instance, output_files, input_data, tables_only, crop_size, debug, debug_dir)
+
+        # Clean up temporary directory
+        shutil.rmtree(temp_dir, ignore_errors=True)
+        return results, num_pages
+
+
+    def _process_non_pdf(self, model_inference_instance, input_data, tables_only, crop_size, debug, debug_dir):
+        """
+        Handles processing and inference for non-PDF files, with optional table extraction.
+        """
+        file_path = input_data[0]["file_path"]
+
+        if tables_only:
+            return self._extract_tables(model_inference_instance, file_path, input_data, debug, debug_dir), 1
+        else:
+            temp_dir = tempfile.mkdtemp()
+
+            if crop_size:
+                if debug:
+                    print(f"Cropping image borders by {crop_size} pixels.")
+                image_optimizer = ImageOptimizer()
+                cropped_file_path = image_optimizer.crop_image_borders(file_path, temp_dir, debug_dir, crop_size)
+                input_data[0]["file_path"] = cropped_file_path
+
+            file_path = input_data[0]["file_path"]
+            input_data[0]["file_path"] = [file_path]
+            results = model_inference_instance.inference(input_data)
+
+            shutil.rmtree(temp_dir, ignore_errors=True)
+
+            return results, 1
+
+    def _process_pages(self, model_inference_instance, output_files, input_data, tables_only, crop_size, debug, debug_dir):
+        """
+        Processes individual pages (PDF split) and handles table extraction or inference.
+
+        Args:
+            model_inference_instance: The model inference object.
+            output_files: List of file paths for the split PDF pages.
+            input_data: Input data for inference.
+            tables_only: Whether to only process tables.
+            crop_size: Size for cropping image borders.
+            debug: Debug flag for logging.
+            debug_dir: Directory for saving debug information.
+
+        Returns:
+            List of results from the processing or inference.
+        """
+        results_array = []
+
+        if tables_only:
+            if debug:
+                print(f"Processing {len(output_files)} pages for table extraction.")
+            # Process each page individually for table extraction
+            for i, file_path in enumerate(output_files):
+                tables_result = self._extract_tables(
+                    model_inference_instance, file_path, input_data, debug, debug_dir, page_index=i
+                )
+                # Since _extract_tables returns a list with one JSON string, unpack it
+                results_array.extend(tables_result)  # Unpack the single JSON string
+        else:
+            if debug:
+                print(f"Processing {len(output_files)} pages for inference at once.")
+
+            temp_dir = tempfile.mkdtemp()
+            cropped_files = []
+
+            if crop_size:
+                if debug:
+                    print(f"Cropping image borders by {crop_size} pixels from {len(output_files)} images.")
+
+                image_optimizer = ImageOptimizer()
+
+                # Process each file in the output_files array
+                for file_path in output_files:
+                    cropped_file_path = image_optimizer.crop_image_borders(
+                        file_path,
+                        temp_dir,
+                        debug_dir,
+                        crop_size
+                    )
+                    cropped_files.append(cropped_file_path)
+
+                # Use the cropped files for inference
+                input_data[0]["file_path"] = cropped_files
+            else:
+                # If no cropping needed, use original files directly
+                input_data[0]["file_path"] = output_files
+
+            # Process all files at once
+            results = model_inference_instance.inference(input_data)
+            results_array.extend(results)
+
+            # Clean up temporary directory
+            shutil.rmtree(temp_dir, ignore_errors=True)
+
+        return results_array
+
+
+    def _extract_tables(self, model_inference_instance, file_path, input_data, debug, debug_dir, page_index=None):
+        """
+        Detects and processes tables from an input file.
+        """
+        table_detector = TableDetector()
+        cropped_tables = table_detector.detect_tables(file_path, local=False, debug_dir=debug_dir, debug=debug)
+        results_array = []
+        temp_dir = tempfile.mkdtemp()
+
+        for i, table in enumerate(cropped_tables):
+            table_index = f"page_{page_index + 1}_table_{i + 1}" if page_index is not None else f"table_{i + 1}"
+            print(f"Processing {table_index} for document {file_path}")
+
+            output_filename = os.path.join(temp_dir, f"{table_index}.jpg")
+            table.save(output_filename, "JPEG")
+
+            input_data[0]["file_path"] = [output_filename]
+            result = self._run_model_inference(model_inference_instance, input_data)
+            results_array.append(result)
+
+        shutil.rmtree(temp_dir, ignore_errors=True)
+
+        # Merge results_array elements into a single JSON structure
+        merged_results = {"page_tables": results_array}
+
+        # Format the merged results as a JSON string with indentation
+        formatted_results = json.dumps(merged_results, indent=4)
+
+        # Return the formatted JSON string wrapped in a list
+        return [formatted_results]
+
+
+    @staticmethod
+    def _run_model_inference(model_inference_instance, input_data):
+        """
+        Runs model inference and handles JSON decoding.
+        """
+        result = model_inference_instance.inference(input_data)[0]
+        try:
+            return json.loads(result) if isinstance(result, str) else result
+        except json.JSONDecodeError:
+            return {"message": "Invalid JSON format in LLM output", "valid": "false"}
+
+
+    @staticmethod
+    def is_pdf(file_path):
+        """Checks if a file is a PDF based on its extension."""
+        return file_path.lower().endswith('.pdf')
+
+
+if __name__ == "__main__":
+    # run locally: python -m sparrow_parse.extractors.vllm_extractor
+
+    extractor = VLLMExtractor()
+
+    # # export HF_TOKEN="hf_"
+    # config = {
+    #     "method": "mlx",  # Could be 'huggingface', 'mlx' or 'local_gpu'
+    #     "model_name": "mlx-community/Qwen2.5-VL-7B-Instruct-8bit",
+    #     # "hf_space": "katanaml/sparrow-qwen2-vl-7b",
+    #     # "hf_token": os.getenv('HF_TOKEN'),
+    #     # Additional fields for local GPU inference
+    #     # "device": "cuda", "model_path": "model.pth"
+    # }
+    #
+    # # Use the factory to get the correct instance
+    # factory = InferenceFactory(config)
+    # model_inference_instance = factory.get_inference_instance()
+    #
+    # input_data = [
+    #     {
+    #         "file_path": "/Users/andrejb/Work/katana-git/sparrow/sparrow-ml/llm/data/bonds_table.png",
+    #         "text_input": "retrieve document data. return response in JSON format"
+    #     }
+    # ]
+    #
+    # # Now you can run inference without knowing which implementation is used
+    # results_array, num_pages = extractor.run_inference(model_inference_instance, input_data, tables_only=False,
+    #                                                    generic_query=False,
+    #                                                    crop_size=0,
+    #                                                    debug_dir="/Users/andrejb/Work/katana-git/sparrow/sparrow-ml/llm/data/",
+    #                                                    debug=True,
+    #                                                    mode=None)
+    #
+    # for i, result in enumerate(results_array):
+    #     print(f"Result for page {i + 1}:", result)
+    # print(f"Number of pages: {num_pages}")

sparrow_parse/helpers/image_optimizer.py

@@ -0,0 +1,59 @@
+from PIL import Image
+import os
+
+
+class ImageOptimizer(object):
+    def __init__(self):
+        pass
+
+    def crop_image_borders(self, file_path, temp_dir, debug_dir=None, crop_size=60):
+        """
+        Crops all four borders of an image by the specified size.
+
+        Args:
+            file_path (str): Path to the input image
+            temp_dir (str): Temporary directory to store the cropped image
+            debug_dir (str, optional): Directory to save a debug copy of the cropped image
+            crop_size (int): Number of pixels to crop from each border
+
+        Returns:
+            str: Path to the cropped image in temp_dir
+        """
+        try:
+            # Open the image
+            with Image.open(file_path) as img:
+                # Get image dimensions
+                width, height = img.size
+
+                # Calculate the crop box
+                left = crop_size
+                top = crop_size
+                right = width - crop_size
+                bottom = height - crop_size
+
+                # Ensure we're not trying to crop more than the image size
+                if right <= left or bottom <= top:
+                    raise ValueError("Crop size is too large for the image dimensions")
+
+                # Perform the crop
+                cropped_img = img.crop((left, top, right, bottom))
+
+                # Get original filename without path
+                filename = os.path.basename(file_path)
+                name, ext = os.path.splitext(filename)
+
+                # Save cropped image in temp_dir
+                output_path = os.path.join(temp_dir, f"{name}_cropped{ext}")
+                cropped_img.save(output_path)
+
+                # If debug_dir is provided, save a debug copy
+                if debug_dir:
+                    os.makedirs(debug_dir, exist_ok=True)
+                    debug_path = os.path.join(debug_dir, f"{name}_cropped_debug{ext}")
+                    cropped_img.save(debug_path)
+                    print(f"Debug cropped image saved to: {debug_path}")
+
+                return output_path
+
+        except Exception as e:
+            raise Exception(f"Error processing image: {str(e)}")

sparrow_parse/helpers/pdf_optimizer.py

@@ -0,0 +1,79 @@
+import pypdf
+from pdf2image import convert_from_path
+import os
+import tempfile
+import shutil
+
+
+class PDFOptimizer(object):
+    def __init__(self):
+        pass
+
+    def split_pdf_to_pages(self, file_path, debug_dir=None, convert_to_images=False):
+        # Create a temporary directory
+        temp_dir = tempfile.mkdtemp()
+        output_files = []
+
+        if not convert_to_images:
+            # Open the PDF file
+            with open(file_path, 'rb') as pdf_file:
+                reader = pypdf.PdfReader(pdf_file)
+                number_of_pages = len(reader.pages)
+
+                # Split the PDF into separate files per page
+                for page_num in range(number_of_pages):
+                    writer = pypdf.PdfWriter()
+                    writer.add_page(reader.pages[page_num])
+
+                    output_filename = os.path.join(temp_dir, f'page_{page_num + 1}.pdf')
+                    with open(output_filename, 'wb') as output_file:
+                        writer.write(output_file)
+                        output_files.append(output_filename)
+
+                    if debug_dir:
+                        # Save each page to the debug folder
+                        debug_output_filename = os.path.join(debug_dir, f'page_{page_num + 1}.pdf')
+                        with open(debug_output_filename, 'wb') as output_file:
+                            writer.write(output_file)
+
+            # Return the number of pages, the list of file paths, and the temporary directory
+            return number_of_pages, output_files, temp_dir
+        else:
+            # Convert the PDF to images
+            images = convert_from_path(file_path, dpi=300)
+            base_name = os.path.splitext(os.path.basename(file_path))[0]
+
+            # Save the images to the temporary directory
+            for i, image in enumerate(images):
+                output_filename = os.path.join(temp_dir, f'{base_name}_page_{i + 1}.jpg')
+                image.save(output_filename, 'JPEG')
+                output_files.append(output_filename)
+
+                if debug_dir:
+                    # Save each image to the debug folder
+                    os.makedirs(debug_dir, exist_ok=True)
+                    debug_output_filename = os.path.join(debug_dir, f'{base_name}_page_{i + 1}_debug.jpg')
+                    image.save(debug_output_filename, 'JPEG')
+                    print(f"Debug image saved to: {debug_output_filename}")
+
+            # Return the number of pages, the list of file paths, and the temporary directory
+            return len(images), output_files, temp_dir
+
+
+if __name__ == "__main__":
+    pdf_optimizer = PDFOptimizer()
+
+    # debug_dir = "/Users/andrejb/infra/shared/katana-git/sparrow/sparrow-ml/llm/data/"
+    # # Ensure the output directory exists
+    # os.makedirs(output_directory, exist_ok=True)
+    #
+    # # Split the optimized PDF into separate pages
+    # num_pages, output_files, temp_dir = pdf_optimizer.split_pdf_to_pages("/Users/andrejb/infra/shared/katana-git/sparrow/sparrow-ml/llm/data/oracle_10k_2014_q1_small.pdf",
+    #                                                                      debug_dir,
+    #                                                                      True)
+    #
+    # print(f"Number of pages: {num_pages}")
+    # print(f"Output files: {output_files}")
+    # print(f"Temporary directory: {temp_dir}")
+    #
+    # shutil.rmtree(temp_dir, ignore_errors=True)

sparrow_parse/processors/table_structure_processor.py

@@ -0,0 +1,275 @@
+from rich.progress import Progress, SpinnerColumn, TextColumn
+from rich import print
+from transformers import AutoModelForObjectDetection
+import torch
+from PIL import Image
+from torchvision import transforms
+import os
+
+
+class TableDetector(object):
+    _model = None  # Static variable to hold the table detection model
+    _device = None  # Static variable to hold the device information
+
+    def __init__(self):
+        pass
+
+    class MaxResize(object):
+        def __init__(self, max_size=800):
+            self.max_size = max_size
+
+        def __call__(self, image):
+            width, height = image.size
+            current_max_size = max(width, height)
+            scale = self.max_size / current_max_size
+            resized_image = image.resize((int(round(scale * width)), int(round(scale * height))))
+
+            return resized_image
+
+    @classmethod
+    def _initialize_model(cls, invoke_pipeline_step, local):
+        """
+        Static method to initialize the table detection model if not already initialized.
+        """
+        if cls._model is None:
+            # Use invoke_pipeline_step to load the model
+            cls._model, cls._device = invoke_pipeline_step(
+                lambda: cls.load_table_detection_model(),
+                "Loading table detection model...",
+                local
+            )
+            print("Table detection model initialized.")
+
+
+    def detect_tables(self, file_path, local=True, debug_dir=None, debug=False):
+        # Ensure the model is initialized using invoke_pipeline_step
+        self._initialize_model(self.invoke_pipeline_step, local)
+
+        # Use the static model and device
+        model, device = self._model, self._device
+
+        outputs, image = self.invoke_pipeline_step(
+            lambda: self.prepare_image(file_path, model, device),
+            "Preparing image for table detection...",
+            local
+        )
+
+        objects = self.invoke_pipeline_step(
+            lambda: self.identify_tables(model, outputs, image),
+            "Identifying tables in the image...",
+            local
+        )
+
+        cropped_tables = self.invoke_pipeline_step(
+            lambda: self.crop_tables(file_path, image, objects, debug, debug_dir),
+            "Cropping tables from the image...",
+            local
+        )
+
+        return cropped_tables
+
+
+    @staticmethod
+    def load_table_detection_model():
+        model = AutoModelForObjectDetection.from_pretrained("microsoft/table-transformer-detection", revision="no_timm")
+
+        device = "cuda" if torch.cuda.is_available() else "cpu"
+        model.to(device)
+
+        return model, device
+
+
+    def prepare_image(self, file_path, model, device):
+        image = Image.open(file_path).convert("RGB")
+
+        detection_transform = transforms.Compose([
+            self.MaxResize(800),
+            transforms.ToTensor(),
+            transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
+        ])
+
+        pixel_values = detection_transform(image).unsqueeze(0)
+        pixel_values = pixel_values.to(device)
+
+        with torch.no_grad():
+            outputs = model(pixel_values)
+
+        return outputs, image
+
+    def identify_tables(self, model, outputs, image):
+        id2label = model.config.id2label
+        id2label[len(model.config.id2label)] = "no object"
+
+        objects = self.outputs_to_objects(outputs, image.size, id2label)
+        return objects
+
+
+    def crop_tables(self, file_path, image, objects, debug, debug_dir):
+        tokens = []
+        detection_class_thresholds = {
+            "table": 0.5,
+            "table rotated": 0.5,
+            "no object": 10
+        }
+        crop_padding = 30
+
+        tables_crops = self.objects_to_crops(image, tokens, objects, detection_class_thresholds, padding=crop_padding)
+
+        cropped_tables = []
+
+        if len(tables_crops) == 0:
+            if debug:
+                print("No tables detected in: ", file_path)
+
+            return None
+        elif len(tables_crops) > 1:
+            for i, table_crop in enumerate(tables_crops):
+                if debug:
+                    print("Table detected in:", file_path, "-", i + 1)
+
+                cropped_table = table_crop['image'].convert("RGB")
+                cropped_tables.append(cropped_table)
+
+                if debug_dir:
+                    file_name_table = self.append_filename(file_path, debug_dir, f"table_cropped_{i + 1}")
+                    cropped_table.save(file_name_table)
+        else:
+            if debug:
+                print("Table detected in: ", file_path)
+
+            cropped_table = tables_crops[0]['image'].convert("RGB")
+            cropped_tables.append(cropped_table)
+
+            if debug_dir:
+                file_name_table = self.append_filename(file_path, debug_dir, "table_cropped")
+                cropped_table.save(file_name_table)
+
+        return cropped_tables
+
+    # for output bounding box post-processing
+    @staticmethod
+    def box_cxcywh_to_xyxy(x):
+        x_c, y_c, w, h = x.unbind(-1)
+        b = [(x_c - 0.5 * w), (y_c - 0.5 * h), (x_c + 0.5 * w), (y_c + 0.5 * h)]
+        return torch.stack(b, dim=1)
+
+    def rescale_bboxes(self, out_bbox, size):
+        img_w, img_h = size
+        b = self.box_cxcywh_to_xyxy(out_bbox)
+        b = b * torch.tensor([img_w, img_h, img_w, img_h], dtype=torch.float32)
+        return b
+
+    def outputs_to_objects(self, outputs, img_size, id2label):
+        m = outputs.logits.softmax(-1).max(-1)
+        pred_labels = list(m.indices.detach().cpu().numpy())[0]
+        pred_scores = list(m.values.detach().cpu().numpy())[0]
+        pred_bboxes = outputs['pred_boxes'].detach().cpu()[0]
+        pred_bboxes = [elem.tolist() for elem in self.rescale_bboxes(pred_bboxes, img_size)]
+
+        objects = []
+        for label, score, bbox in zip(pred_labels, pred_scores, pred_bboxes):
+            class_label = id2label[int(label)]
+            if not class_label == 'no object':
+                objects.append({'label': class_label, 'score': float(score),
+                                'bbox': [float(elem) for elem in bbox]})
+
+        return objects
+
+    def objects_to_crops(self, img, tokens, objects, class_thresholds, padding=10):
+        """
+        Process the bounding boxes produced by the table detection model into
+        cropped table images and cropped tokens.
+        """
+
+        table_crops = []
+        for obj in objects:
+            if obj['score'] < class_thresholds[obj['label']]:
+                continue
+
+            cropped_table = {}
+
+            bbox = obj['bbox']
+            bbox = [bbox[0] - padding, bbox[1] - padding, bbox[2] + padding, bbox[3] + padding]
+
+            cropped_img = img.crop(bbox)
+
+            table_tokens = [token for token in tokens if self.iob(token['bbox'], bbox) >= 0.5]
+            for token in table_tokens:
+                token['bbox'] = [token['bbox'][0] - bbox[0],
+                                 token['bbox'][1] - bbox[1],
+                                 token['bbox'][2] - bbox[0],
+                                 token['bbox'][3] - bbox[1]]
+
+            # If table is predicted to be rotated, rotate cropped image and tokens/words:
+            if obj['label'] == 'table rotated':
+                cropped_img = cropped_img.rotate(270, expand=True)
+                for token in table_tokens:
+                    bbox = token['bbox']
+                    bbox = [cropped_img.size[0] - bbox[3] - 1,
+                            bbox[0],
+                            cropped_img.size[0] - bbox[1] - 1,
+                            bbox[2]]
+                    token['bbox'] = bbox
+
+            cropped_table['image'] = cropped_img
+            cropped_table['tokens'] = table_tokens
+
+            table_crops.append(cropped_table)
+
+        return table_crops
+
+
+    @staticmethod
+    def append_filename(file_path, debug_dir, word):
+        directory, filename = os.path.split(file_path)
+        name, ext = os.path.splitext(filename)
+        new_filename = f"{name}_{word}{ext}"
+        return os.path.join(debug_dir, new_filename)
+
+    @staticmethod
+    def iob(boxA, boxB):
+        # Determine the coordinates of the intersection rectangle
+        xA = max(boxA[0], boxB[0])
+        yA = max(boxA[1], boxB[1])
+        xB = min(boxA[2], boxB[2])
+        yB = min(boxA[3], boxB[3])
+
+        # Compute the area of intersection rectangle
+        interArea = max(0, xB - xA + 1) * max(0, yB - yA + 1)
+
+        # Compute the area of both the prediction and ground-truth rectangles
+        boxAArea = (boxA[2] - boxA[0] + 1) * (boxA[3] - boxA[1] + 1)
+        boxBArea = (boxB[2] - boxB[0] + 1) * (boxB[3] - boxB[1] + 1)
+
+        # Compute the intersection over box (IoB)
+        iob = interArea / float(boxAArea)
+
+        return iob
+
+
+    @staticmethod
+    def invoke_pipeline_step(task_call, task_description, local):
+        if local:
+            with Progress(
+                    SpinnerColumn(),
+                    TextColumn("[progress.description]{task.description}"),
+                    transient=False,
+            ) as progress:
+                progress.add_task(description=task_description, total=None)
+                ret = task_call()
+        else:
+            print(task_description)
+            ret = task_call()
+
+        return ret
+
+
+if __name__ == "__main__":
+    table_detector = TableDetector()
+
+    # file_path = "/Users/andrejb/Work/katana-git/sparrow/sparrow-ml/llm/data/bonds_table.png"
+    # cropped_tables = table_detector.detect_tables(file_path, local=True, debug_dir="/Users/andrejb/Work/katana-git/sparrow/sparrow-ml/llm/data/", debug=True)
+
+    # for i, cropped_table in enumerate(cropped_tables):
+    #     file_name_table = table_detector.append_filename(file_path, "cropped_" + str(i))
+    #     cropped_table.save(file_name_table)

sparrow_parse/text_extraction.py

@@ -0,0 +1,30 @@
+from mlx_vlm import load, apply_chat_template, generate
+from mlx_vlm.utils import load_image
+
+# For test purposes, we will use a sample image
+
+# Load model and processor
+qwen_vl_model, qwen_vl_processor = load("mlx-community/Qwen2.5-VL-7B-Instruct-8bit")
+qwen_vl_config = qwen_vl_model.config
+
+image = load_image("images/graph.png")
+
+messages = [
+    {"role": "system", "content": "You are an expert at extracting text from images. Format your response in json."},
+    {"role": "user", "content": "Extract the names, labels and y coordinates from the image."}
+]
+
+# Apply chat template
+prompt = apply_chat_template(qwen_vl_processor, qwen_vl_config, messages)
+
+# Generate text
+qwen_vl_output = generate(
+    qwen_vl_model,
+    qwen_vl_processor,
+    prompt,
+    image,
+    max_tokens=1000,
+    temperature=0.7,
+)
+
+print(qwen_vl_output)

sparrow_parse/vllm/huggingface_inference.py

@@ -0,0 +1,60 @@
+from gradio_client import Client, handle_file
+from sparrow_parse.vllm.inference_base import ModelInference
+import json
+import os
+import ast
+
+
+class HuggingFaceInference(ModelInference):
+    def __init__(self, hf_space, hf_token):
+        self.hf_space = hf_space
+        self.hf_token = hf_token
+
+
+    def process_response(self, output_text):
+        json_string = output_text
+
+        json_string = json_string.strip("[]'")
+        json_string = json_string.replace("```json\n", "").replace("\n```", "")
+        json_string = json_string.replace("'", "")
+
+        try:
+            formatted_json = json.loads(json_string)
+            return json.dumps(formatted_json, indent=2)
+        except json.JSONDecodeError as e:
+            print("Failed to parse JSON:", e)
+            return output_text
+
+
+    def inference(self, input_data, mode=None):
+        if mode == "static":
+            simple_json = self.get_simple_json()
+            return [simple_json]
+
+        client = Client(self.hf_space, hf_token=self.hf_token)
+
+        # Extract and prepare the absolute paths for all file paths in input_data
+        file_paths = [
+            os.path.abspath(file_path)
+            for data in input_data
+            for file_path in data["file_path"]
+        ]
+
+        # Validate file existence and prepare files for the Gradio client
+        image_files = [handle_file(path) for path in file_paths if os.path.exists(path)]
+
+        results = client.predict(
+            input_imgs=image_files,
+            text_input=input_data[0]["text_input"],  # Single shared text input for all images
+            api_name="/run_inference"  # Specify the Gradio API endpoint
+        )
+
+        # Convert the string into a Python list
+        parsed_results = ast.literal_eval(results)
+
+        results_array = []
+        for page_output in parsed_results:
+            page_result = self.process_response(page_output)
+            results_array.append(page_result)
+
+        return results_array

sparrow_parse/vllm/inference_base.py

@@ -0,0 +1,30 @@
+from abc import ABC, abstractmethod
+import json
+
+
+class ModelInference(ABC):
+    @abstractmethod
+    def inference(self, input_data, mode=None):
+        """This method should be implemented by subclasses."""
+        pass
+
+    def get_simple_json(self):
+        # Define a simple data structure
+        data = {
+          "table": [
+                {
+                  "description": "Revenues",
+                  "latest_amount": 12453,
+                  "previous_amount": 11445
+                },
+                {
+                  "description": "Operating expenses",
+                  "latest_amount": 9157,
+                  "previous_amount": 8822
+                }
+            ]
+        }
+
+        # Convert the dictionary to a JSON string
+        json_data = json.dumps(data, indent=4)
+        return json_data

sparrow_parse/vllm/inference_factory.py

@@ -0,0 +1,25 @@
+from sparrow_parse.vllm.huggingface_inference import HuggingFaceInference
+from sparrow_parse.vllm.local_gpu_inference import LocalGPUInference
+from sparrow_parse.vllm.mlx_inference import MLXInference
+
+
+class InferenceFactory:
+    def __init__(self, config):
+        self.config = config
+
+    def get_inference_instance(self):
+        if self.config["method"] == "huggingface":
+            return HuggingFaceInference(hf_space=self.config["hf_space"], hf_token=self.config["hf_token"])
+        elif self.config["method"] == "local_gpu":
+            model = self._load_local_model()  # Replace with actual model loading logic
+            return LocalGPUInference(model=model, device=self.config.get("device", "cuda"))
+        elif self.config["method"] == "mlx":
+            return MLXInference(model_name=self.config["model_name"])
+        else:
+            raise ValueError(f"Unknown method: {self.config['method']}")
+
+    def _load_local_model(self):
+        # Example: Load a PyTorch model (replace with actual loading code)
+        # model = torch.load('model.pth')
+        # return model
+        raise NotImplementedError("Model loading logic not implemented")

sparrow_parse/vllm/infra/qwen2_vl_7b/app.py

@@ -0,0 +1,155 @@
+import gradio as gr
+import spaces
+from transformers import Qwen2VLForConditionalGeneration, AutoProcessor
+from qwen_vl_utils import process_vision_info
+from PIL import Image
+from datetime import datetime
+import os
+
+# subprocess.run('pip install flash-attn --no-build-isolation', env={'FLASH_ATTENTION_SKIP_CUDA_BUILD': "TRUE"}, shell=True)
+
+DESCRIPTION = "[Sparrow Qwen2-VL-7B Backend](https://github.com/katanaml/sparrow)"
+
+
+def array_to_image_path(image_filepath, max_width=1250, max_height=1750):
+    if image_filepath is None:
+        raise ValueError("No image provided. Please upload an image before submitting.")
+
+    # Open the uploaded image using its filepath
+    img = Image.open(image_filepath)
+
+    # Extract the file extension from the uploaded file
+    input_image_extension = image_filepath.split('.')[-1].lower()  # Extract extension from filepath
+
+    # Set file extension based on the original file, otherwise default to PNG
+    if input_image_extension in ['jpg', 'jpeg', 'png']:
+        file_extension = input_image_extension
+    else:
+        file_extension = 'png'  # Default to PNG if extension is unavailable or invalid
+
+    # Get the current dimensions of the image
+    width, height = img.size
+
+    # Initialize new dimensions to current size
+    new_width, new_height = width, height
+
+    # Check if the image exceeds the maximum dimensions
+    if width > max_width or height > max_height:
+        # Calculate the new size, maintaining the aspect ratio
+        aspect_ratio = width / height
+
+        if width > max_width:
+            new_width = max_width
+            new_height = int(new_width / aspect_ratio)
+
+        if new_height > max_height:
+            new_height = max_height
+            new_width = int(new_height * aspect_ratio)
+
+    # Generate a unique filename using timestamp
+    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+    filename = f"image_{timestamp}.{file_extension}"
+
+    # Save the image
+    img.save(filename)
+
+    # Get the full path of the saved image
+    full_path = os.path.abspath(filename)
+
+    return full_path, new_width, new_height
+
+
+# Initialize the model and processor globally to optimize performance
+model = Qwen2VLForConditionalGeneration.from_pretrained(
+    "Qwen/Qwen2-VL-7B-Instruct",
+    torch_dtype="auto",
+    device_map="auto"
+)
+
+processor = AutoProcessor.from_pretrained("Qwen/Qwen2-VL-7B-Instruct")
+
+
+@spaces.GPU
+def run_inference(input_imgs, text_input):
+    results = []
+
+    for image in input_imgs:
+        # Convert each image to the required format
+        image_path, width, height = array_to_image_path(image)
+
+        try:
+            # Prepare messages for each image
+            messages = [
+                {
+                    "role": "user",
+                    "content": [
+                        {
+                            "type": "image",
+                            "image": image_path,
+                            "resized_height": height,
+                            "resized_width": width
+                        },
+                        {
+                            "type": "text",
+                            "text": text_input
+                        }
+                    ]
+                }
+            ]
+
+            # Prepare inputs for the model
+            text = processor.apply_chat_template(
+                messages, tokenize=False, add_generation_prompt=True
+            )
+
+            image_inputs, video_inputs = process_vision_info(messages)
+            inputs = processor(
+                text=[text],
+                images=image_inputs,
+                videos=video_inputs,
+                padding=True,
+                return_tensors="pt",
+            )
+            inputs = inputs.to("cuda")
+
+            # Generate inference output
+            generated_ids = model.generate(**inputs, max_new_tokens=4096)
+            generated_ids_trimmed = [
+                out_ids[len(in_ids):] for in_ids, out_ids in zip(inputs.input_ids, generated_ids)
+            ]
+            raw_output = processor.batch_decode(
+                generated_ids_trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=True
+            )
+
+            results.append(raw_output[0])
+            print("Processed: " + image)
+        finally:
+            # Clean up the temporary image file
+            os.remove(image_path)
+
+    return results
+
+
+css = """
+  #output {
+    height: 500px; 
+    overflow: auto; 
+    border: 1px solid #ccc; 
+  }
+"""
+
+with gr.Blocks(css=css) as demo:
+    gr.Markdown(DESCRIPTION)
+    with gr.Tab(label="Qwen2-VL-7B Input"):
+        with gr.Row():
+            with gr.Column():
+                input_imgs = gr.Files(file_types=["image"], label="Upload Document Images")
+                text_input = gr.Textbox(label="Query")
+                submit_btn = gr.Button(value="Submit", variant="primary")
+            with gr.Column():
+                output_text = gr.Textbox(label="Response")
+
+        submit_btn.click(run_inference, [input_imgs, text_input], [output_text])
+
+demo.queue(api_open=True)
+demo.launch(debug=True)

sparrow_parse/vllm/infra/qwen2_vl_7b/requirements.txt

@@ -0,0 +1,9 @@
+numpy==1.24.4
+Pillow==10.3.0
+Requests==2.31.0
+torch
+torchvision
+git+https://github.com/huggingface/transformers.git
+accelerate
+qwen-vl-utils
+gradio==4.44.1

sparrow_parse/vllm/local_gpu_inference.py

@@ -0,0 +1,16 @@
+import torch
+from sparrow_parse.vllm.inference_base import ModelInference
+
+
+class LocalGPUInference(ModelInference):
+    def __init__(self, model, device='cuda'):
+        self.model = model
+        self.device = device
+        self.model.to(self.device)
+
+    def inference(self, input_data, mode=None):
+        self.model.eval()  # Set the model to evaluation mode
+        with torch.no_grad():  # No need to calculate gradients
+            input_tensor = torch.tensor(input_data).to(self.device)
+            output = self.model(input_tensor)
+        return output.cpu().numpy()  # Convert the output back to NumPy if necessary

sparrow_parse/vllm/mlx_inference.py

@@ -0,0 +1,140 @@
+from mlx_vlm import load, generate
+from mlx_vlm.prompt_utils import apply_chat_template
+from mlx_vlm.utils import load_image
+from sparrow_parse.vllm.inference_base import ModelInference
+import os
+import json
+from rich import print
+
+
+class MLXInference(ModelInference):
+    """
+        A class for performing inference using the MLX model.
+        Handles image preprocessing, response formatting, and model interaction.
+        """
+
+    def __init__(self, model_name):
+        """
+        Initialize the inference class with the given model name.
+
+        :param model_name: Name of the model to load.
+        """
+        self.model_name = model_name
+        print(f"MLXInference initialized for model: {model_name}")
+
+
+    @staticmethod
+    def _load_model_and_processor(model_name):
+        """
+        Load the model and processor for inference.
+
+        :param model_name: Name of the model to load.
+        :return: Tuple containing the loaded model and processor.
+        """
+        model, processor = load(model_name)
+        print(f"Loaded model: {model_name}")
+        return model, processor
+
+
+    def process_response(self, output_text):
+        """
+        Process and clean the model's raw output to format as JSON.
+
+        :param output_text: Raw output text from the model.
+        :return: A formatted JSON string or the original text in case of errors.
+        """
+        try:
+            cleaned_text = (
+                output_text.strip("[]'")
+                .replace("```json\n", "")
+                .replace("\n```", "")
+                .replace("'", "")
+            )
+            formatted_json = json.loads(cleaned_text)
+            return json.dumps(formatted_json, indent=2)
+        except json.JSONDecodeError as e:
+            print(f"Failed to parse JSON in MLX inference backend: {e}")
+            return output_text
+
+
+    def load_image_data(self, image_filepath, max_width=1250, max_height=1750):
+        """
+        Load and resize image while maintaining its aspect ratio.
+
+        :param image_filepath: Path to the image file.
+        :param max_width: Maximum allowed width of the image.
+        :param max_height: Maximum allowed height of the image.
+        :return: Tuple containing the image object and its new dimensions.
+        """
+        image = load_image(image_filepath)  # Assuming load_image is defined elsewhere
+        width, height = image.size
+
+        # Calculate new dimensions while maintaining the aspect ratio
+        if width > max_width or height > max_height:
+            aspect_ratio = width / height
+            new_width = min(max_width, int(max_height * aspect_ratio))
+            new_height = min(max_height, int(max_width / aspect_ratio))
+            return image, new_width, new_height
+
+        return image, width, height
+
+
+    def inference(self, input_data, mode=None):
+        """
+        Perform inference on input data using the specified model.
+
+        :param input_data: A list of dictionaries containing image file paths and text inputs.
+        :param mode: Optional mode for inference ("static" for simple JSON output).
+        :return: List of processed model responses.
+        """
+        if mode == "static":
+            return [self.get_simple_json()]
+
+        # Load the model and processor
+        model, processor = self._load_model_and_processor(self.model_name)
+        config = model.config
+
+        # Prepare absolute file paths
+        file_paths = self._extract_file_paths(input_data)
+
+        results = []
+        for file_path in file_paths:
+            image, width, height = self.load_image_data(file_path)
+
+            # Prepare messages for the chat model
+            messages = [
+                {"role": "system", "content": "You are an expert at extracting structured text from image documents."},
+                {"role": "user", "content": input_data[0]["text_input"]},
+            ]
+
+            # Generate and process response
+            prompt = apply_chat_template(processor, config, messages)  # Assuming defined
+            response = generate(
+                model,
+                processor,
+                prompt,
+                image,
+                resize_shape=(width, height),
+                max_tokens=4000,
+                temperature=0.0,
+                verbose=False
+            )
+            results.append(self.process_response(response))
+
+            print("Inference completed successfully for: ", file_path)
+
+        return results
+
+    @staticmethod
+    def _extract_file_paths(input_data):
+        """
+        Extract and resolve absolute file paths from input data.
+
+        :param input_data: List of dictionaries containing image file paths.
+        :return: List of absolute file paths.
+        """
+        return [
+            os.path.abspath(file_path)
+            for data in input_data
+            for file_path in data.get("file_path", [])
+        ]