sparrow-parse 0.3.8__py3-none-any.whl → 0.3.9__py3-none-any.whl

sparrow_parse/__init__.py

@@ -1 +1 @@
-__version__ = '0.3.8'
+__version__ = '0.3.9'

sparrow_parse/extractors/vllm_extractor.py

@@ -9,41 +9,45 @@ class VLLMExtractor(object):
     def __init__(self):
         pass
 
-    def run_inference(self,
-                      model_inference_instance,
-                      input_data,
-                      generic_query=False,
-                      debug_dir=None,
-                      debug=False,
-                      mode=None):
+    def run_inference(self, model_inference_instance, input_data,
+                      generic_query=False, debug_dir=None, debug=False, mode=None):
+        # Modify input for generic queries
         if generic_query:
             input_data[0]["text_input"] = "retrieve document data. return response in JSON format"
 
         if debug:
             print("Input Data:", input_data)
 
-        results_array = []
+        # Check if the input file is a PDF
+        file_path = input_data[0]["file_path"]
+        if self.is_pdf(file_path):
+            return self._process_pdf(model_inference_instance, input_data, debug_dir, mode)
 
-        if self.is_pdf(input_data[0]["file_path"]):
-            pdf_optimizer = PDFOptimizer()
-            num_pages, output_files, temp_dir = pdf_optimizer.split_pdf_to_pages(input_data[0]["file_path"],
-                                                                                 debug_dir,
-                                                                                 True)
+        # Default processing for non-PDF files
+        input_data[0]["file_path"] = [file_path]
+        results_array = model_inference_instance.inference(input_data)
+        return results_array, 1
 
-            input_data[0]["file_path"] = output_files
 
-            # Run inference on the page
-            results_array = model_inference_instance.inference(input_data, mode)
+    def _process_pdf(self, model_inference_instance, input_data, debug_dir, mode):
+        """Handles processing and inference for PDF files."""
+        pdf_optimizer = PDFOptimizer()
+        num_pages, output_files, temp_dir = pdf_optimizer.split_pdf_to_pages(input_data[0]["file_path"],
+                                                                             debug_dir,
+                                                                             True)
+        # Update file paths for PDF pages
+        input_data[0]["file_path"] = output_files
 
-            shutil.rmtree(temp_dir, ignore_errors=True)
-            return results_array, num_pages
+        # Run inference on PDF pages
+        results_array = model_inference_instance.inference(input_data, mode)
 
-        input_data[0]["file_path"] = [input_data[0]["file_path"]]
-        results_array = model_inference_instance.inference(input_data)
-
-        return results_array, 1
+        # Clean up temporary directory
+        shutil.rmtree(temp_dir, ignore_errors=True)
+        return results_array, num_pages
 
-    def is_pdf(self, file_path):
+    @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__":
@@ -53,9 +57,10 @@ if __name__ == "__main__":
 
     # # export HF_TOKEN="hf_"
     # config = {
-    #     "method": "huggingface",  # Could be 'huggingface' or 'local_gpu'
-    #     "hf_space": "katanaml/sparrow-qwen2-vl-7b",
-    #     "hf_token": os.getenv('HF_TOKEN'),
+    #     "method": "mlx",  # Could be 'huggingface', 'mlx' or 'local_gpu'
+    #     "model_name": "mlx-community/Qwen2-VL-72B-Instruct-4bit",
+    #     # "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"
     # }
@@ -66,14 +71,14 @@ if __name__ == "__main__":
     #
     # input_data = [
     #     {
-    #         "file_path": "/Users/andrejb/infra/shared/katana-git/sparrow/sparrow-ml/llm/data/oracle_10k_2014_q1_small.pdf",
-    #         "text_input": "retrieve {\"table\": [{\"description\": \"str\", \"latest_amount\": 0, \"previous_amount\": 0}]}. return response in JSON format"
+    #         "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, generic_query=False,
-    #                                  debug_dir="/Users/andrejb/infra/shared/katana-git/sparrow/sparrow-ml/llm/data/",
+    #                                  debug_dir=None,
     #                                  debug=True,
     #                                  mode=None)
     #

sparrow_parse/vllm/inference_factory.py

@@ -1,5 +1,6 @@
 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:
@@ -12,6 +13,8 @@ class InferenceFactory:
         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']}")
 

sparrow_parse/vllm/mlx_inference.py

@@ -0,0 +1,135 @@
+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
+
+
+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 and load the model once.
+
+        :param model_name: Name of the model to load.
+        """
+        self.model, self.processor = self._load_model_and_processor(model_name)
+        self.config = self.model.config
+
+        print(f"Loaded 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.
+        """
+        return load(model_name)
+
+
+    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: {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()]
+
+        # 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(self.processor, self.config, messages)  # Assuming defined
+            response = generate(
+                self.model,
+                self.processor,
+                image,
+                prompt,
+                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", [])
+        ]

{sparrow_parse-0.3.8.dist-info → sparrow_parse-0.3.9.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: sparrow-parse
-Version: 0.3.8
+Version: 0.3.9
 Summary: Sparrow Parse is a Python package (part of Sparrow) for parsing and extracting information from documents.
 Home-page: https://github.com/katanaml/sparrow/tree/main/sparrow-data/parse
 Author: Andrej Baranovskij
@@ -24,7 +24,7 @@ Requires-Dist: pypdf==4.3.0
 Requires-Dist: easyocr==1.7.1
 Requires-Dist: gradio-client
 Requires-Dist: pdf2image
-Requires-Dist: mlx-vlm==0.1.1; sys_platform == "darwin" and platform_machine == "arm64"
+Requires-Dist: mlx-vlm==0.1.3; sys_platform == "darwin" and platform_machine == "arm64"
 
 # Sparrow Parse
 
@@ -40,7 +40,7 @@ pip install sparrow-parse
 
 ## Parsing and extraction
 
-### Sparrow Parse VL (vision-language model) extractor with Hugging Face GPU infra
+### 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
@@ -50,13 +50,9 @@ from sparrow_parse.extractors.vllm_extractor import VLLMExtractor
 
 extractor = VLLMExtractor()
 
-# export HF_TOKEN="hf_"
 config = {
-    "method": "huggingface",  # Could be 'huggingface' or 'local_gpu'
-    "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"
+    "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
@@ -65,14 +61,14 @@ model_inference_instance = factory.get_inference_instance()
 
 input_data = [
     {
-        "file_path": "/data/oracle_10k_2014_q1_small.pdf",
-        "text_input": "retrieve {"table": [{"description": "str", "latest_amount": 0, "previous_amount": 0}]}. return response in JSON format"
+        "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, generic_query=False,
-                                 debug_dir="/data/",
+                                 debug_dir=None,
                                  debug=True,
                                  mode=None)
 
@@ -85,6 +81,16 @@ Use `mode="static"` if you want to simulate LLM call, without executing LLM back
 
 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

{sparrow_parse-0.3.8.dist-info → sparrow_parse-0.3.9.dist-info}/RECORD

@@ -1,7 +1,7 @@
-sparrow_parse/__init__.py,sha256=nDN5c_hLef3c3ZFmKN7QOhgRDGgxKZlDHKhzZ-D9x28,21
+sparrow_parse/__init__.py,sha256=OJRl30XMXR02jOFfAqFJB3-IohqKdn-uwu-tvsSqqhc,21
 sparrow_parse/__main__.py,sha256=Xs1bpJV0n08KWOoQE34FBYn6EBXZA9HIYJKrE4ZdG78,153
 sparrow_parse/extractors/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-sparrow_parse/extractors/vllm_extractor.py,sha256=mBPgeyMuHUa6jN_OZLVE-426tD4zYnFT61oxebk7XJc,3191
+sparrow_parse/extractors/vllm_extractor.py,sha256=FI1y5WGG6Z-msH0XHkvjyA8Oh2x2nBc1_kswNlsw27Y,3464
 sparrow_parse/helpers/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 sparrow_parse/helpers/pdf_optimizer.py,sha256=GIqQYWtixFeZGCRFXL0lQfQByapCDuQzzRHAkzcPwLE,3302
 sparrow_parse/processors/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -9,10 +9,11 @@ sparrow_parse/processors/table_structure_processor.py,sha256=bG_6jx66n_KNdY_O6hr
 sparrow_parse/vllm/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 sparrow_parse/vllm/huggingface_inference.py,sha256=EJnG6PesGKMc_0qGPN8ufE6pSnhAgFu0XjCbaLCNVyM,1980
 sparrow_parse/vllm/inference_base.py,sha256=4mwGoAY63MB4cHZpV0czTkJWEzimmiTzqqzKmLNzgjw,820
-sparrow_parse/vllm/inference_factory.py,sha256=r04e95uPWG5l8Q23yeDqKmvFxLyF991aA2m0hfBTNn8,993
+sparrow_parse/vllm/inference_factory.py,sha256=FTM65O-dW2WZchHOrNN7_Q3-FlVoAc65iSptuuUuClM,1166
 sparrow_parse/vllm/local_gpu_inference.py,sha256=aHoJTejb5xrXjWDIGu5RBQWEyRCOBCB04sMvO2Wyvg8,628
-sparrow_parse-0.3.8.dist-info/METADATA,sha256=eszXQ77bD2K7HAouFLzidLJ5ev2YR7rGDeWcR1tvma4,6275
-sparrow_parse-0.3.8.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
-sparrow_parse-0.3.8.dist-info/entry_points.txt,sha256=8CrvTVTTcz1YuZ8aRCYNOH15ZOAaYLlcbYX3t28HwJY,54
-sparrow_parse-0.3.8.dist-info/top_level.txt,sha256=n6b-WtT91zKLyCPZTP7wvne8v_yvIahcsz-4sX8I0rY,14
-sparrow_parse-0.3.8.dist-info/RECORD,,
+sparrow_parse/vllm/mlx_inference.py,sha256=3srKhlhVRoYqeVy_PziNZRHcHl6D3ksibvcOAug4irA,4723
+sparrow_parse-0.3.9.dist-info/METADATA,sha256=kbHkfNEfmoqd2yt3aM1IqwVwf6HYBWR-FTnc8XOIWSQ,6351
+sparrow_parse-0.3.9.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
+sparrow_parse-0.3.9.dist-info/entry_points.txt,sha256=8CrvTVTTcz1YuZ8aRCYNOH15ZOAaYLlcbYX3t28HwJY,54
+sparrow_parse-0.3.9.dist-info/top_level.txt,sha256=n6b-WtT91zKLyCPZTP7wvne8v_yvIahcsz-4sX8I0rY,14
+sparrow_parse-0.3.9.dist-info/RECORD,,