sparrow-parse 0.3.5__tar.gz → 0.3.6__tar.gz

{sparrow-parse-0.3.5 → sparrow-parse-0.3.6}/PKG-INFO

@@ -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
@@ -34,6 +34,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
 
@@ -54,16 +56,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.5 → sparrow-parse-0.3.6}/README.md

@@ -15,6 +15,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
 
@@ -35,16 +37,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.5 → sparrow-parse-0.3.6}/setup.py

@@ -8,7 +8,7 @@ with open("requirements.txt", "r", encoding="utf-8") as fh:
 
 setup(
     name="sparrow-parse",
-    version="0.3.5",
+    version="0.3.6",
     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.",

sparrow-parse-0.3.6/sparrow_parse/__init__.py

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

sparrow-parse-0.3.6/sparrow_parse/extractors/vllm_extractor.py

@@ -0,0 +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_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
+
+    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 = [
+    #     {
+    #         "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-0.3.5 → sparrow-parse-0.3.6}/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-0.3.5 → sparrow-parse-0.3.6}/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-0.3.6/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-0.3.5 → sparrow-parse-0.3.6}/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 → sparrow-parse-0.3.6}/sparrow_parse.egg-info/PKG-INFO

@@ -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
@@ -34,6 +34,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
 
@@ -54,16 +56,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.5/sparrow_parse/__init__.py

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

sparrow-parse-0.3.5/sparrow_parse/extractors/vllm_extractor.py

@@ -1,46 +0,0 @@
-from sparrow_parse.vllm.inference_factory import InferenceFactory
-from rich import print
-import os
-
-
-class VLLMExtractor(object):
-    def __init__(self):
-        pass
-
-    def run_inference(self, model_inference_instance, input_data, generic_query=False, debug=False):
-        if generic_query:
-            input_data[0]["text_input"] = "retrieve document data. return response in JSON format"
-
-        if debug:
-            print("Input Data:", input_data)
-
-        result = model_inference_instance.inference(input_data)
-
-        return result
-
-if __name__ == "__main__":
-    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)

sparrow-parse-0.3.5/sparrow_parse/vllm/inference_base.py

@@ -1,7 +0,0 @@
-from abc import ABC, abstractmethod
-
-class ModelInference(ABC):
-    @abstractmethod
-    def inference(self, input_data):
-        """This method should be implemented by subclasses."""
-        pass