vlm4ocr 0.4.0__tar.gz → 0.4.1__tar.gz

{vlm4ocr-0.4.0 → vlm4ocr-0.4.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: vlm4ocr
-Version: 0.4.0
+Version: 0.4.1
 Summary: Python package and Web App for OCR with vision language models.
 License: MIT
 Author: Enshuo (David) Hsu
@@ -12,6 +12,7 @@ Classifier: Programming Language :: Python :: 3.12
 Provides-Extra: tesseract
 Requires-Dist: colorama (>=0.4.4)
 Requires-Dist: json-repair (>=0.30.0)
+Requires-Dist: llm-inference-engine (>=0.1.1,<0.2.0)
 Requires-Dist: pdf2image (>=1.16.0)
 Requires-Dist: pillow (>=10.0.0)
 Requires-Dist: pytesseract (>=0.3.13) ; extra == "tesseract"

{vlm4ocr-0.4.0 → vlm4ocr-0.4.1}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "vlm4ocr"
-version = "0.4.0"
+version = "0.4.1"
 description = "Python package and Web App for OCR with vision language models."
 authors = ["Enshuo (David) Hsu"]
 license = "MIT"
@@ -18,6 +18,7 @@ pdf2image = ">=1.16.0"
 colorama = ">=0.4.4"
 pillow = ">=10.0.0"
 json-repair = ">=0.30.0"
+llm-inference-engine = "^0.1.1"
 pytesseract = { version = ">=0.3.13", optional = true }
 
 [tool.poetry.scripts]

vlm4ocr-0.4.1/vlm4ocr/vlm_engines.py

@@ -0,0 +1,276 @@
+import abc
+from typing import List, Dict
+from PIL import Image
+from vlm4ocr.utils import image_to_base64
+from vlm4ocr.data_types import FewShotExample
+from llm_inference_engine.llm_configs import (
+    LLMConfig as VLMConfig, 
+    BasicLLMConfig as BasicVLMConfig, 
+    ReasoningLLMConfig as ReasoningVLMConfig, 
+    OpenAIReasoningLLMConfig as OpenAIReasoningVLMConfig
+)
+from llm_inference_engine.utils import MessagesLogger
+from llm_inference_engine.engines import (
+    InferenceEngine,
+    OllamaInferenceEngine,
+    OpenAICompatibleInferenceEngine,
+    VLLMInferenceEngine,
+    OpenRouterInferenceEngine,
+    OpenAIInferenceEngine,
+    AzureOpenAIInferenceEngine,
+)
+
+
+class VLMEngine(InferenceEngine):
+    @abc.abstractmethod
+    def get_ocr_messages(self, system_prompt:str, user_prompt:str, image:Image.Image, few_shot_examples:List[FewShotExample]=None) -> List[Dict[str,str]]:
+        """
+        This method inputs an image and returns the correesponding chat messages for the inference engine.
+
+        Parameters:
+        ----------
+        system_prompt : str
+            the system prompt.
+        user_prompt : str
+            the user prompt.
+        image : Image.Image
+            the image for OCR.
+        few_shot_examples : List[FewShotExample], Optional
+            list of few-shot examples. 
+        """
+        return NotImplemented
+
+
+class OllamaVLMEngine(OllamaInferenceEngine, VLMEngine):
+    def get_ocr_messages(self, system_prompt:str, user_prompt:str, image:Image.Image, few_shot_examples:List[FewShotExample]=None) -> List[Dict[str,str]]:
+        """
+        This method inputs an image and returns the correesponding chat messages for the inference engine.
+
+        Parameters:
+        ----------
+        system_prompt : str
+            the system prompt.
+        user_prompt : str
+            the user prompt.
+        image : Image.Image
+            the image for OCR.
+        few_shot_examples : List[FewShotExample], Optional
+            list of few-shot examples. 
+        """
+        base64_str = image_to_base64(image)
+        output_messages = []
+        # system message
+        system_message = {"role": "system", "content": system_prompt}
+        output_messages.append(system_message)
+
+        # few-shot examples
+        if few_shot_examples is not None:
+            for example in few_shot_examples:
+                if not isinstance(example, FewShotExample):
+                    raise ValueError("Few-shot example must be a FewShotExample object.")
+                
+                example_image_b64 = image_to_base64(example.image)
+                example_user_message = {"role": "user", "content": user_prompt, "images": [example_image_b64]}
+                example_agent_message = {"role": "assistant", "content": example.text}
+                output_messages.append(example_user_message)
+                output_messages.append(example_agent_message)
+
+        # user message
+        user_message = {"role": "user", "content": user_prompt, "images": [base64_str]}
+        output_messages.append(user_message)
+
+        return output_messages
+
+
+class OpenAICompatibleVLMEngine(OpenAICompatibleInferenceEngine, VLMEngine):
+    def get_ocr_messages(self, system_prompt:str, user_prompt:str, image:Image.Image, format:str='png', 
+                         detail:str="high", few_shot_examples:List[FewShotExample]=None) -> List[Dict[str,str]]:
+        """
+        This method inputs an image and returns the correesponding chat messages for the inference engine.
+
+        Parameters:
+        ----------
+        system_prompt : str
+            the system prompt.
+        user_prompt : str
+            the user prompt.
+        image : Image.Image
+            the image for OCR.
+        format : str, Optional
+            the image format. 
+        detail : str, Optional
+            the detail level of the image. Default is "high". 
+        few_shot_examples : List[FewShotExample], Optional
+            list of few-shot examples.
+        """
+        base64_str = image_to_base64(image)
+        output_messages = []
+        # system message
+        system_message = {"role": "system", "content": system_prompt}
+        output_messages.append(system_message)
+
+        # few-shot examples
+        if few_shot_examples is not None:
+            for example in few_shot_examples:
+                if not isinstance(example, FewShotExample):
+                    raise ValueError("Few-shot example must be a FewShotExample object.")
+                
+                example_image_b64 = image_to_base64(example.image)
+                example_user_message = {
+                    "role": "user",
+                    "content": [
+                        {
+                            "type": "image_url",
+                            "image_url": {
+                                "url": f"data:image/{format};base64,{example_image_b64}",
+                                "detail": detail
+                            },
+                        },
+                        {"type": "text", "text": user_prompt},
+                    ],
+                }
+                example_agent_message = {"role": "assistant", "content": example.text}
+                output_messages.append(example_user_message)
+                output_messages.append(example_agent_message)
+
+        # user message
+        user_message = {
+            "role": "user",
+            "content": [
+                {
+                    "type": "image_url",
+                    "image_url": {
+                        "url": f"data:image/{format};base64,{base64_str}",
+                        "detail": detail
+                    },
+                },
+                {"type": "text", "text": user_prompt},
+            ],
+        }
+        output_messages.append(user_message)
+        return output_messages
+
+
+class VLLMVLMEngine(VLLMInferenceEngine, OpenAICompatibleVLMEngine):
+    """
+    vLLM OpenAI compatible server inference engine.
+    https://docs.vllm.ai/en/latest/serving/openai_compatible_server.html
+
+    For parameters and documentation, refer to https://platform.openai.com/docs/api-reference/introduction
+
+    Parameters:
+    ----------
+    model_name : str
+        model name as shown in the vLLM server
+    api_key : str, Optional
+        the API key for the vLLM server.
+    base_url : str, Optional
+        the base url for the vLLM server. 
+    config : LLMConfig
+        the LLM configuration.
+    """
+    pass
+
+class OpenRouterVLMEngine(OpenRouterInferenceEngine, OpenAICompatibleVLMEngine):
+    """
+    OpenRouter OpenAI-compatible server inference engine.
+
+    Parameters:
+    ----------
+    model_name : str
+        model name as shown in the vLLM server
+    api_key : str, Optional
+        the API key for the vLLM server. If None, will use the key in os.environ['OPENROUTER_API_KEY'].
+    base_url : str, Optional
+        the base url for the vLLM server. 
+    config : LLMConfig
+        the LLM configuration.
+    """
+    pass
+
+class OpenAIVLMEngine(OpenAIInferenceEngine, VLMEngine):
+    def get_ocr_messages(self, system_prompt:str, user_prompt:str, image:Image.Image, format:str='png', 
+                         detail:str="high", few_shot_examples:List[FewShotExample]=None) -> List[Dict[str,str]]:
+        """
+        This method inputs an image and returns the correesponding chat messages for the inference engine.
+
+        Parameters:
+        ----------
+        system_prompt : str
+            the system prompt.
+        user_prompt : str
+            the user prompt.
+        image : Image.Image
+            the image for OCR.
+        format : str, Optional
+            the image format. 
+        detail : str, Optional
+            the detail level of the image. Default is "high". 
+        few_shot_examples : List[FewShotExample], Optional
+            list of few-shot examples. Each example is a dict with keys "image" (PIL.Image.Image) and "text" (str).
+        """
+        base64_str = image_to_base64(image)
+        output_messages = []
+        # system message
+        system_message = {"role": "system", "content": system_prompt}
+        output_messages.append(system_message)
+
+        # few-shot examples
+        if few_shot_examples is not None:
+            for example in few_shot_examples:
+                if not isinstance(example, FewShotExample):
+                    raise ValueError("Few-shot example must be a FewShotExample object.")
+                
+                example_image_b64 = image_to_base64(example.image)
+                example_user_message = {
+                    "role": "user",
+                    "content": [
+                        {
+                            "type": "image_url",
+                            "image_url": {
+                                "url": f"data:image/{format};base64,{example_image_b64}",
+                                "detail": detail
+                            },
+                        },
+                        {"type": "text", "text": user_prompt},
+                    ],
+                }
+                example_agent_message = {"role": "assistant", "content": example.text}
+                output_messages.append(example_user_message)
+                output_messages.append(example_agent_message)
+
+        # user message
+        user_message = {
+            "role": "user",
+            "content": [
+                {
+                    "type": "image_url",
+                    "image_url": {
+                        "url": f"data:image/{format};base64,{base64_str}",
+                        "detail": detail
+                    },
+                },
+                {"type": "text", "text": user_prompt},
+            ],
+        }
+        output_messages.append(user_message)
+        return output_messages
+
+
+class AzureOpenAIVLMEngine(AzureOpenAIInferenceEngine, OpenAIVLMEngine):
+    """
+    The Azure OpenAI API inference engine.
+    For parameters and documentation, refer to 
+    - https://azure.microsoft.com/en-us/products/ai-services/openai-service
+    - https://learn.microsoft.com/en-us/azure/ai-services/openai/quickstart
+    
+    Parameters:
+    ----------
+    model : str
+        model name as described in https://platform.openai.com/docs/models
+    api_version : str
+        the Azure OpenAI API version
+    config : LLMConfig
+        the LLM configuration.
+    """
+    pass