sparrow-parse 0.3.5__py3-none-any.whl → 0.3.6__py3-none-any.whl

sparrow_parse/__init__.py

@@ -1 +1 @@
-__version__ = '0.3.5'
+__version__ = '0.3.6'

sparrow_parse/extractors/vllm_extractor.py

@@ -1,46 +1,87 @@
 from sparrow_parse.vllm.inference_factory import InferenceFactory
+from sparrow_parse.helpers.pdf_optimizer import PDFOptimizer
 from rich import print
 import os
+import shutil
 
 
 class VLLMExtractor(object):
     def __init__(self):
         pass
 
-    def run_inference(self, model_inference_instance, input_data, generic_query=False, debug=False):
+    def run_inference(self,
+                      model_inference_instance,
+                      input_data,
+                      generic_query=False,
+                      debug_dir=None,
+                      debug=False,
+                      mode=None):
         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 = []
+
+        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)
+
+            # Run inference on each page
+            for page_num, output_file in enumerate(output_files):
+                input_data[0]["file_path"] = output_file
+                if debug:
+                    print(f"Running inference on page {page_num + 1}...")
+
+                # Run inference on the page
+                result = model_inference_instance.inference(input_data, mode)
+                results_array.append(result)
+
+            shutil.rmtree(temp_dir, ignore_errors=True)
+            return results_array, num_pages
+
         result = model_inference_instance.inference(input_data)
+        results_array.append(result)
+
+        return results_array, 1
 
-        return result
+    def is_pdf(self, file_path):
+        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": "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"
-    }
-
-    # Use the factory to get the correct instance
-    factory = InferenceFactory(config)
-    model_inference_instance = factory.get_inference_instance()
-
-    input_data = [
-        {
-            "image": "/Users/andrejb/Documents/work/epik/bankstatement/bonds_table.png",
-            "text_input": "retrieve financial instruments data. return response in JSON format"
-        }
-    ]
-
-    # Now you can run inference without knowing which implementation is used
-    result = extractor.run_inference(model_inference_instance, input_data, generic_query=False, debug=True)
-    print("Inference Result:", result)
+    # # 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"
+    # }
+    #
+    # # Use the factory to get the correct instance
+    # factory = InferenceFactory(config)
+    # model_inference_instance = factory.get_inference_instance()
+    #
+    # 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"
+    #     }
+    # ]
+    #
+    # # 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=True,
+    #                                  mode="static")
+    #
+    # 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/pdf_optimizer.py

@@ -40,17 +40,18 @@ class PDFOptimizer(object):
             return number_of_pages, output_files, temp_dir
         else:
             # Convert the PDF to images
-            images = convert_from_path(file_path, dpi=400)
+            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'page_{i + 1}.jpg')
+                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 output_dir:
                     # Save each image to the debug folder
-                    debug_output_filename = os.path.join(output_dir, f'page_{i + 1}.jpg')
+                    debug_output_filename = os.path.join(output_dir, f'{base_name}_page_{i + 1}.jpg')
                     image.save(debug_output_filename, 'JPEG')
 
             # Return the number of pages, the list of file paths, and the temporary directory
@@ -60,13 +61,17 @@ class PDFOptimizer(object):
 if __name__ == "__main__":
     pdf_optimizer = PDFOptimizer()
 
-    # output_directory = "/Users/andrejb/Documents/work/bankstatement/output_pages"
+    # output_directory = "/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/Documents/work/bankstatement/statement.pdf",
+    # 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",
     #                                                                      output_directory,
-    #                                                                      False)
+    #                                                                      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/vllm/huggingface_inference.py

@@ -24,11 +24,15 @@ class HuggingFaceInference(ModelInference):
             return output_text
 
 
-    def inference(self, input_data):
+    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)
 
         result = client.predict(
-            image=handle_file(input_data[0]["image"]),
+            image=handle_file(input_data[0]["file_path"]),
             text_input=input_data[0]["text_input"],
             api_name="/run_inference"
         )

sparrow_parse/vllm/inference_base.py

@@ -1,7 +1,30 @@
 from abc import ABC, abstractmethod
+import json
+
 
 class ModelInference(ABC):
     @abstractmethod
-    def inference(self, input_data):
+    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/local_gpu_inference.py

@@ -8,7 +8,7 @@ class LocalGPUInference(ModelInference):
         self.device = device
         self.model.to(self.device)
 
-    def inference(self, input_data):
+    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)

{sparrow_parse-0.3.5.dist-info → sparrow_parse-0.3.6.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: sparrow-parse
-Version: 0.3.5
+Version: 0.3.6
 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
@@ -42,6 +42,8 @@ pip install sparrow-parse
 ### Sparrow Parse VL (vision-language model) extractor with Hugging Face 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
 
@@ -62,16 +64,24 @@ model_inference_instance = factory.get_inference_instance()
 
 input_data = [
     {
-        "image": "/data/bonds_table.png",
-        "text_input": "retrieve all data. return response in JSON format"
+        "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"
     }
 ]
 
 # Now you can run inference without knowing which implementation is used
-result = extractor.run_inference(model_inference_instance, input_data, generic_query=False, debug=True)
-print("Inference Result:", result)
+results_array, num_pages = extractor.run_inference(model_inference_instance, input_data, generic_query=False,
+                                 debug_dir="/data/",
+                                 debug=True,
+                                 mode="static")
+
+for i, result in enumerate(results_array):
+    print(f"Result for page {i + 1}:", result)
+print(f"Number of pages: {num_pages}")
 ```
 
+Use `mode="static"` if you want to simulate LLM call, without executing LLM backend.
+
 ## PDF pre-processing
 
 ```

sparrow_parse-0.3.6.dist-info/RECORD

@@ -0,0 +1,18 @@
+sparrow_parse/__init__.py,sha256=IbpUPwvtjLOqowcOFsWQ6LKq-FH6cI19IpvfQlxufq0,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=Wo8sOvsQt6YHd7bvB_DB8MUa71FioO9xcQOWA3PQ6eU,3415
+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
+sparrow_parse/processors/table_structure_processor.py,sha256=bG_6jx66n_KNdY_O6hrZD1D4DHX5Qy__RYcKHmrSGnc,23894
+sparrow_parse/vllm/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+sparrow_parse/vllm/huggingface_inference.py,sha256=nalmPJFfrFlRnfd4yTq4HvIwDvIXjhKUlEyZ6gzMqe0,1239
+sparrow_parse/vllm/inference_base.py,sha256=4mwGoAY63MB4cHZpV0czTkJWEzimmiTzqqzKmLNzgjw,820
+sparrow_parse/vllm/inference_factory.py,sha256=r04e95uPWG5l8Q23yeDqKmvFxLyF991aA2m0hfBTNn8,993
+sparrow_parse/vllm/local_gpu_inference.py,sha256=aHoJTejb5xrXjWDIGu5RBQWEyRCOBCB04sMvO2Wyvg8,628
+sparrow_parse-0.3.6.dist-info/METADATA,sha256=ANS8eWCx07bQOOFFnJUKwsiPo-ZT42b8DvMwP9o-jf4,5827
+sparrow_parse-0.3.6.dist-info/WHEEL,sha256=eOLhNAGa2EW3wWl_TU484h7q1UNgy0JXjjoqKoxAAQc,92
+sparrow_parse-0.3.6.dist-info/entry_points.txt,sha256=8CrvTVTTcz1YuZ8aRCYNOH15ZOAaYLlcbYX3t28HwJY,54
+sparrow_parse-0.3.6.dist-info/top_level.txt,sha256=n6b-WtT91zKLyCPZTP7wvne8v_yvIahcsz-4sX8I0rY,14
+sparrow_parse-0.3.6.dist-info/RECORD,,

sparrow_parse-0.3.5.dist-info/RECORD

@@ -1,18 +0,0 @@
-sparrow_parse/__init__.py,sha256=e9arv8KorBrIZFQXAlN4DOQTh91btae1iR36M_3Wafk,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=Qwmf-SW4z_UstiiynX5TkyovlkokVhLuzcbUVZ16TXM,1540
-sparrow_parse/helpers/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-sparrow_parse/helpers/pdf_optimizer.py,sha256=KI_EweGt9Y_rDH1uCpYD5wKCW3rdjSFFhoVtiPBxX8k,3013
-sparrow_parse/processors/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-sparrow_parse/processors/table_structure_processor.py,sha256=bG_6jx66n_KNdY_O6hrZD1D4DHX5Qy__RYcKHmrSGnc,23894
-sparrow_parse/vllm/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-sparrow_parse/vllm/huggingface_inference.py,sha256=Q2Ju65LDzbO-8RWW7cXzrR-pbZ1zKuPVODlKOTWKg_E,1114
-sparrow_parse/vllm/inference_base.py,sha256=W0N2khehGdF1XHzZACG3I1UZaydHMk6BZgWNvaJD4Ck,197
-sparrow_parse/vllm/inference_factory.py,sha256=r04e95uPWG5l8Q23yeDqKmvFxLyF991aA2m0hfBTNn8,993
-sparrow_parse/vllm/local_gpu_inference.py,sha256=I_uWYiFAQhRrykOKbVz69NzftDxuemDKtAye4kWhtnU,617
-sparrow_parse-0.3.5.dist-info/METADATA,sha256=4i_-BJalUQFFUZoo919pfr51ZqvU1Jfq-mEFkHf0gWU,5342
-sparrow_parse-0.3.5.dist-info/WHEEL,sha256=eOLhNAGa2EW3wWl_TU484h7q1UNgy0JXjjoqKoxAAQc,92
-sparrow_parse-0.3.5.dist-info/entry_points.txt,sha256=8CrvTVTTcz1YuZ8aRCYNOH15ZOAaYLlcbYX3t28HwJY,54
-sparrow_parse-0.3.5.dist-info/top_level.txt,sha256=n6b-WtT91zKLyCPZTP7wvne8v_yvIahcsz-4sX8I0rY,14
-sparrow_parse-0.3.5.dist-info/RECORD,,