sparrow-parse 0.4.2__py3-none-any.whl → 0.4.4__py3-none-any.whl

sparrow_parse/__init__.py

@@ -1 +1 @@
-__version__ = '0.4.2'
+__version__ = '0.4.4'

sparrow_parse/extractors/vllm_extractor.py

@@ -7,7 +7,6 @@ from rich import print
 import os
 import tempfile
 import shutil
-from typing import Any, Dict, List, Union
 
 
 class VLLMExtractor(object):
@@ -156,7 +155,7 @@ if __name__ == "__main__":
     # # export HF_TOKEN="hf_"
     # config = {
     #     "method": "mlx",  # Could be 'huggingface', 'mlx' or 'local_gpu'
-    #     "model_name": "mlx-community/Qwen2-VL-72B-Instruct-4bit",
+    #     "model_name": "mlx-community/Qwen2-VL-7B-Instruct-8bit",
     #     # "hf_space": "katanaml/sparrow-qwen2-vl-7b",
     #     # "hf_token": os.getenv('HF_TOKEN'),
     #     # Additional fields for local GPU inference

sparrow_parse/vllm/mlx_inference.py

@@ -14,25 +14,48 @@ class MLXInference(ModelInference):
 
     def __init__(self, model_name):
         """
-        Initialize the inference class with the given model name and load the model once.
+        Initialize the inference class with the given model name.
 
         :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
+        self.model_name = model_name
+        self.model = None
+        self.processor = None
+        print(f"MLXInference initialized with model: {model_name}")
 
-        print(f"Loaded model: {model_name}")
 
+    def __del__(self):
+        """
+        Destructor to clean up resources when the object is deleted.
+        """
+        self.unload_model()
 
-    @staticmethod
-    def _load_model_and_processor(model_name):
+
+    def unload_model(self):
         """
-        Load the model and processor for inference.
+        Unload the model and release resources.
+        """
+        if self.model:
+            del self.model
+            self.model = None
+        if self.processor:
+            del self.processor
+            self.processor = None
+        print(f"Model {self.model_name} and its resources have been unloaded.")
+
 
+    def _load_model_and_processor(self, 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)
+        print(f"Loading model and processor for: {model_name}...")
+        model, processor = load(model_name)
+        self.model = model  # Store model instance
+        self.processor = processor  # Store processor instance
+        print(f"Model and processor for '{model_name}' loaded successfully.")
+        return model, processor
 
 
     def process_response(self, output_text):
@@ -81,44 +104,54 @@ class MLXInference(ModelInference):
     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)
+        try:
+            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)
+                response = generate(
+                    model,
+                    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
+
+        finally:
+            # Always unload the model after inference
+            self.unload_model()
 
-        return results
 
     @staticmethod
     def _extract_file_paths(input_data):

{sparrow_parse-0.4.2.dist-info → sparrow_parse-0.4.4.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: sparrow-parse
-Version: 0.4.2
+Version: 0.4.4
 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

{sparrow_parse-0.4.2.dist-info → sparrow_parse-0.4.4.dist-info}/RECORD

@@ -1,7 +1,7 @@
-sparrow_parse/__init__.py,sha256=ZQn_AcWYegaUtOl4-txMtrEFR2pE4wBpoPlmrITggnY,21
+sparrow_parse/__init__.py,sha256=jP9l7AhBCN2A-6tezbTIihxoMTDna4SLTYhvVxwbdNM,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=ybWpRpDH0YHoYpHkjIJtm7DQoHJBKNsirK2YIAlMvGo,7863
+sparrow_parse/extractors/vllm_extractor.py,sha256=PDLgLlKiq3Bv-UOQTzX3AgxNOLcEU2EniGAXLjMC30U,7820
 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
@@ -11,9 +11,9 @@ sparrow_parse/vllm/huggingface_inference.py,sha256=EJnG6PesGKMc_0qGPN8ufE6pSnhAg
 sparrow_parse/vllm/inference_base.py,sha256=4mwGoAY63MB4cHZpV0czTkJWEzimmiTzqqzKmLNzgjw,820
 sparrow_parse/vllm/inference_factory.py,sha256=FTM65O-dW2WZchHOrNN7_Q3-FlVoAc65iSptuuUuClM,1166
 sparrow_parse/vllm/local_gpu_inference.py,sha256=aHoJTejb5xrXjWDIGu5RBQWEyRCOBCB04sMvO2Wyvg8,628
-sparrow_parse/vllm/mlx_inference.py,sha256=xR40qwjIR0HvrN8x58oOq6F4r1hEANRB-9kcokUQHHU,4748
-sparrow_parse-0.4.2.dist-info/METADATA,sha256=JTSkKdB2X5o3tXovyjnLSEANrPnkvJLtTfeYH7PZRDw,6432
-sparrow_parse-0.4.2.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
-sparrow_parse-0.4.2.dist-info/entry_points.txt,sha256=8CrvTVTTcz1YuZ8aRCYNOH15ZOAaYLlcbYX3t28HwJY,54
-sparrow_parse-0.4.2.dist-info/top_level.txt,sha256=n6b-WtT91zKLyCPZTP7wvne8v_yvIahcsz-4sX8I0rY,14
-sparrow_parse-0.4.2.dist-info/RECORD,,
+sparrow_parse/vllm/mlx_inference.py,sha256=c6-s493jLXE3DfYnwsybiqgk3GU9GEaWt3CrfqLSWKQ,5872
+sparrow_parse-0.4.4.dist-info/METADATA,sha256=x_jaR76FUv5-kR9R5YR9So3OZbWj_rG8hjFZ07lZuto,6432
+sparrow_parse-0.4.4.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
+sparrow_parse-0.4.4.dist-info/entry_points.txt,sha256=8CrvTVTTcz1YuZ8aRCYNOH15ZOAaYLlcbYX3t28HwJY,54
+sparrow_parse-0.4.4.dist-info/top_level.txt,sha256=n6b-WtT91zKLyCPZTP7wvne8v_yvIahcsz-4sX8I0rY,14
+sparrow_parse-0.4.4.dist-info/RECORD,,