vlm4ocr 0.0.1__py3-none-any.whl

vlm4ocr/__init__.py

@@ -0,0 +1,9 @@
+from .ocr_engines import OCREngine
+from .vlm_engines import OllamaVLMEngine, OpenAIVLMEngine, AzureOpenAIVLMEngine
+
+__all__ = [
+    "OCREngine",
+    "OllamaVLMEngine",
+    "OpenAIVLMEngine",
+    "AzureOpenAIVLMEngine"
+]

vlm4ocr/assets/default_prompt_templates/ocr_markdown_system_prompt.txt

@@ -0,0 +1 @@
+You are a helpful assistant that can convert scanned documents into markdown text. Your output is accurate and well-formatted, starting with ```markdown and ending with ```. You will only output the markdown text without any additional explanations or comments. The markdown should include all text, tables, and lists with appropriate headers (e.g., "##"). You will ignore images, icons, or anything that can not be converted into text.

vlm4ocr/assets/default_prompt_templates/ocr_text_system_prompt.txt

@@ -0,0 +1 @@
+You are a helpful assistant that can convert scanned documents into plain text. Your output is accurate and well-formatted. You will only output the plain text without any additional explanations or comments. The plain text should include all text, tables, and lists with appropriate layout (e.g., \n, space, and tab). You will ignore images, icons, or anything that can not be converted into text.

vlm4ocr/assets/default_prompt_templates/ocr_user_prompt.txt

@@ -0,0 +1 @@
+Convert contents in this image into markdown.

vlm4ocr/ocr_engines.py

@@ -0,0 +1,311 @@
+import os
+from typing import List, Dict, Union, Generator, Iterable
+import importlib
+import asyncio
+from vlm4ocr.utils import get_images_from_pdf, get_image_from_file, clean_markdown
+from vlm4ocr.vlm_engines import VLMEngine
+
+SUPPORTED_IMAGE_EXTS = ['.png', '.jpg', '.jpeg', '.bmp', '.gif', '.webp']
+
+class OCREngine:
+    def __init__(self, vlm_engine:VLMEngine, output_mode:str="markdown", system_prompt:str=None, user_prompt:str=None, page_delimiter:str="\n\n---\n\n"):
+        """
+        This class inputs a image or PDF file path and processes them using a VLM inference engine. Outputs plain text or markdown.
+
+        Parameters:
+        -----------
+        inference_engine : InferenceEngine
+            The inference engine to use for OCR.
+        output_mode : str, Optional
+            The output format. Can be 'markdown' or 'text'.
+        system_prompt : str, Optional
+            Custom system prompt. We recommend use a default system prompt by leaving this blank. 
+        user_prompt : str, Optional
+            Custom user prompt. It is good to include some information regarding the document. If not specified, a default will be used.
+        page_delimiter : str, Optional
+            The delimiter to use between PDF pages. 
+        """
+        # Check inference engine
+        if not isinstance(vlm_engine, VLMEngine):
+            raise TypeError("vlm_engine must be an instance of VLMEngine")
+        self.vlm_engine = vlm_engine
+
+        # Check output mode
+        if output_mode not in ["markdown", "text"]:
+            raise ValueError("output_mode must be 'markdown' or 'text'")
+        self.output_mode = output_mode
+
+        # System prompt
+        if isinstance(system_prompt, str) and system_prompt:
+            self.system_prompt = system_prompt
+        else:
+            file_path = importlib.resources.files('vlm4ocr.assets.default_prompt_templates').joinpath(f'ocr_{self.output_mode}_system_prompt.txt')
+            with open(file_path, 'r', encoding='utf-8') as f:
+                self.system_prompt =  f.read()
+
+        # User prompt
+        if isinstance(user_prompt, str) and user_prompt:
+            self.user_prompt = user_prompt
+        else:
+            file_path = importlib.resources.files('vlm4ocr.assets.default_prompt_templates').joinpath('ocr_user_prompt.txt')
+            with open(file_path, 'r', encoding='utf-8') as f:
+                self.user_prompt =  f.read()
+
+        # Page delimiter
+        if isinstance(page_delimiter, str):
+            self.page_delimiter = page_delimiter
+        else:
+            raise ValueError("page_delimiter must be a string")
+        
+
+    def stream_ocr(self, file_path: str, max_new_tokens:int=4096, temperature:float=0.0, **kwrs) -> Generator[str, None, None]:
+        """
+        This method inputs a file path (image or PDF) and stream OCR results in real-time. This is useful for frontend applications.
+
+        Parameters:
+        -----------
+        file_path : str
+            The path to the image or PDF file. Must be one of '.pdf', '.png', '.jpg', '.jpeg', '.bmp', '.gif', '.webp'
+        max_new_tokens : int, Optional
+            The maximum number of tokens to generate.
+        temperature : float, Optional
+            The temperature to use for sampling.
+        
+        Returns:
+        --------
+        Generator[str, None, None]
+            A generator that yields the output.
+        """
+        # Check file path
+        if not isinstance(file_path, str):
+            raise TypeError("file_path must be a string")
+        
+        # Check file extension
+        file_ext = os.path.splitext(file_path)[1].lower()
+        if file_ext not in SUPPORTED_IMAGE_EXTS and file_ext != '.pdf':
+            raise ValueError(f"Unsupported file type: {file_ext}. Supported types are: {SUPPORTED_IMAGE_EXTS + ['.pdf']}")
+
+        # PDF
+        if file_ext == '.pdf':
+            images = get_images_from_pdf(file_path)
+            if not images:
+                raise ValueError(f"No images extracted from PDF: {file_path}")
+            for i, image in enumerate(images):
+                messages = self.vlm_engine.get_ocr_messages(self.system_prompt, self.user_prompt, image)
+                response_stream = self.vlm_engine.chat(
+                    messages,
+                    max_new_tokens=max_new_tokens,
+                    temperature=temperature,
+                    stream=True,
+                    **kwrs
+                )
+                for chunk in response_stream:
+                    yield chunk
+
+                if i < len(images) - 1:
+                    yield self.page_delimiter
+
+        # Image
+        else:
+            image = get_image_from_file(file_path)
+            messages = self.vlm_engine.get_ocr_messages(self.system_prompt, self.user_prompt, image)
+            response_stream = self.vlm_engine.chat(
+                    messages,
+                    max_new_tokens=max_new_tokens,
+                    temperature=temperature,
+                    stream=True,
+                    **kwrs
+                )
+            for chunk in response_stream:
+                yield chunk
+            
+
+    def run_ocr(self, file_paths: Union[str, Iterable[str]], max_new_tokens:int=4096, temperature:float=0.0, 
+                verbose:bool=False, concurrent:bool=False, concurrent_batch_size:int=32, **kwrs) -> Union[str, Generator[str, None, None]]:
+        """
+        This method takes a list of file paths (image or PDF) and perform OCR using the VLM inference engine.
+
+        Parameters:
+        -----------
+        file_paths : Union[str, Iterable[str]]
+            A file path or a list of file paths to process. Must be one of '.pdf', '.png', '.jpg', '.jpeg', '.bmp', '.gif', '.webp'
+        max_new_tokens : int, Optional
+            The maximum number of tokens to generate. 
+        temperature : float, Optional
+            The temperature to use for sampling. 
+        verbose : bool, Optional
+            If True, the function will print the output in terminal. 
+        concurrent : bool, Optional
+            If True, the function will process the files concurrently.
+        concurrent_batch_size : int, Optional
+            The number of images/pages to process concurrently. 
+        """
+        # if file_paths is a string, convert it to a list
+        if isinstance(file_paths, str):
+            file_paths = [file_paths]
+        
+        if not isinstance(file_paths, Iterable):
+            raise TypeError("file_paths must be a string or an iterable of strings")
+        
+        # check if all file paths are valid
+        for file_path in file_paths:
+            if not isinstance(file_path, str):
+                raise TypeError("file_paths must be a string or an iterable of strings")
+            file_ext = os.path.splitext(file_path)[1].lower()
+            if file_ext not in SUPPORTED_IMAGE_EXTS and file_ext != '.pdf':
+                raise ValueError(f"Unsupported file type: {file_ext}. Supported types are: {SUPPORTED_IMAGE_EXTS + ['.pdf']}")
+
+
+        # Concurrent processing
+        if concurrent:
+            # Check concurrent_batch_size
+            if concurrent_batch_size <= 0:
+                raise ValueError("concurrent_batch_size must be greater than 0")
+
+            if verbose:
+                Warning("verbose is not supported for concurrent processing.", UserWarning)
+
+            return asyncio.run(self._run_ocr_async(file_paths, 
+                                                   max_new_tokens=max_new_tokens, 
+                                                   temperature=temperature, 
+                                                   concurrent_batch_size=concurrent_batch_size, 
+                                                   **kwrs))
+        
+        # Sync processing
+        return self._run_ocr(file_paths, max_new_tokens=max_new_tokens, temperature=temperature, verbose=verbose, **kwrs)
+    
+
+    def _run_ocr(self, file_paths: Union[str, Iterable[str]], max_new_tokens:int=4096, 
+                 temperature:float=0.0, verbose:bool=False, **kwrs) -> Iterable[str]:
+        """
+        This method inputs a file path or a list of file paths (image or PDF) and performs OCR using the VLM inference engine.
+
+        Parameters:
+        -----------
+        file_paths : Union[str, Iterable[str]]
+            A file path or a list of file paths to process. Must be one of '.pdf', '.png', '.jpg', '.jpeg', '.bmp', '.gif', '.webp'
+        max_new_tokens : int, Optional
+            The maximum number of tokens to generate.
+        temperature : float, Optional
+            The temperature to use for sampling.
+        verbose : bool, Optional
+            If True, the function will print the output in terminal.
+        
+        Returns:
+        --------
+        Iterable[str]
+            A list of strings containing the OCR results.
+        """
+        ocr_results = []
+        for file_path in file_paths:
+            file_ext = os.path.splitext(file_path)[1].lower()
+            # PDF
+            if file_ext == '.pdf':
+                images = get_images_from_pdf(file_path)
+                if not images:
+                    raise ValueError(f"No images extracted from PDF: {file_path}")
+                pdf_results = []
+                for image in images:
+                    messages = self.vlm_engine.get_ocr_messages(self.system_prompt, self.user_prompt, image)
+                    response = self.vlm_engine.chat(
+                        messages,
+                        max_new_tokens=max_new_tokens,
+                        temperature=temperature,
+                        verbose=verbose,
+                        stream=False,
+                        **kwrs
+                    )
+                    pdf_results.append(response)
+
+                ocr_text = self.page_delimiter.join(pdf_results)
+            # Image
+            else:
+                image = get_image_from_file(file_path)
+                messages = self.vlm_engine.get_ocr_messages(self.system_prompt, self.user_prompt, image)
+                ocr_text = self.vlm_engine.chat(
+                    messages,
+                    max_new_tokens=max_new_tokens,
+                    temperature=temperature,
+                    verbose=verbose,
+                    stream=False,
+                    **kwrs
+                )
+            
+            # Clean markdown
+            if self.output_mode == "markdown":
+                ocr_text = clean_markdown(ocr_text)
+            ocr_results.append(ocr_text)
+
+        return ocr_results
+
+
+    async def _run_ocr_async(self, file_paths: Union[str, Iterable[str]], max_new_tokens:int=4096, 
+                       temperature:float=0.0, concurrent_batch_size:int=32, **kwrs) -> List[str]:
+        """
+        This is the async version of the _run_ocr method. 
+        """
+        # flatten pages/images in file_paths
+        flat_page_list = []
+        for file_path in file_paths:
+            file_ext = os.path.splitext(file_path)[1].lower()
+            # PDF
+            if file_ext == '.pdf':
+                images = get_images_from_pdf(file_path)
+                if not images:
+                    flat_page_list.append({'file_path': file_path, 'file_type': "PDF", "image": image, "page_num": 0, "total_page_count": 0})
+                for page_num, image in enumerate(images):
+                    flat_page_list.append({'file_path': file_path, 'file_type': "PDF", "image": image, "page_num": page_num, "total_page_count": len(images)})
+            # Image
+            else:
+                image = get_image_from_file(file_path)
+                flat_page_list.append({'file_path': file_path, 'file_type': "image", "image": image})
+
+        # Process images with asyncio.Semaphore
+        semaphore = asyncio.Semaphore(concurrent_batch_size)
+        async def semaphore_helper(page:List[Dict[str,str]], max_new_tokens:int, temperature:float, **kwrs):
+            try:
+                messages = self.vlm_engine.get_ocr_messages(self.system_prompt, self.user_prompt, page["image"])
+                async with semaphore:
+                    async_task = self.vlm_engine.chat_async(
+                        messages,
+                        max_new_tokens=max_new_tokens,
+                        temperature=temperature,
+                        **kwrs
+                    )
+                return await async_task
+            except Exception as e:
+                print(f"Error processing image: {e}")
+                return f"[Error: {e}]"
+
+        tasks = []
+        for page in flat_page_list:
+            async_task = semaphore_helper(
+                page,
+                max_new_tokens=max_new_tokens,
+                temperature=temperature,
+                **kwrs
+            )
+            tasks.append(asyncio.create_task(async_task))
+
+        responses = await asyncio.gather(*tasks)
+
+        # Restructure the results
+        ocr_results = []
+        pdf_page_text_buffer = ""
+        for page, ocr_text in zip(flat_page_list, responses):
+            # PDF
+            if page['file_type'] == "PDF":                
+                pdf_page_text_buffer += ocr_text + self.page_delimiter
+                if page['page_num'] == page['total_page_count'] - 1:
+                    if self.output_mode == "markdown":
+                        pdf_page_text_buffer = clean_markdown(pdf_page_text_buffer)
+                    ocr_results.append(pdf_page_text_buffer)
+                    pdf_page_text_buffer = ""
+            # Image
+            if page['file_type'] == "image":
+                if self.output_mode == "markdown":
+                    ocr_text = clean_markdown(ocr_text)
+                ocr_results.append(ocr_text)
+            
+        return ocr_results
+

vlm4ocr/utils.py

@@ -0,0 +1,47 @@
+import os
+import io
+import base64
+from typing import List
+from pdf2image import convert_from_path
+from PIL import Image
+
+def get_images_from_pdf(file_path: str) -> List[Image.Image]:
+    """ Extracts images from a PDF file. """
+    try:
+        images = convert_from_path(file_path)
+        if not images:
+            print(f"Warning: No images extracted from PDF: {file_path}")
+        return images
+    except Exception as e:
+        print(f"Error converting PDF to images: {e}")
+        raise ValueError(f"Failed to process PDF file '{os.path.basename(file_path)}'. Ensure poppler is installed and the file is valid.") from e
+
+
+def get_image_from_file(file_path: str) -> Image.Image:
+        """ Loads a single image file. """
+        try:
+            image = Image.open(file_path)
+            image.load()
+            return image
+        except FileNotFoundError:
+            raise FileNotFoundError(f"Image file not found: {file_path}")
+        except Exception as e:
+            raise ValueError(f"Failed to load image file '{os.path.basename(file_path)}': {e}") from e
+
+
+def image_to_base64(image:Image.Image, format:str="png") -> str:
+    """ Converts an image to a base64 string. """
+    try:
+        buffered = io.BytesIO()
+        image.save(buffered, format=format)
+        img_bytes = buffered.getvalue()
+        encoded_bytes = base64.b64encode(img_bytes)
+        base64_encoded_string = encoded_bytes.decode('utf-8')
+        return base64_encoded_string
+    except Exception as e:
+        print(f"Error converting image to base64: {e}")
+        raise ValueError(f"Failed to convert image to base64: {e}") from e
+    
+def clean_markdown(text:str) -> str:
+    cleaned_text = text.replace("```markdown", "").replace("```", "")
+    return cleaned_text

vlm4ocr/vlm_engines.py

@@ -0,0 +1,444 @@
+import abc
+import importlib
+from typing import List, Dict, Union, Generator
+import warnings
+from PIL import Image
+from vlm4ocr.utils import image_to_base64
+
+
+class VLMEngine:
+    @abc.abstractmethod
+    def __init__(self):
+        """
+        This is an abstract class to provide interfaces for VLM inference engines. 
+        Children classes that inherts this class can be used in extrators. Must implement chat() method.
+        """
+        return NotImplemented
+
+    @abc.abstractmethod
+    def chat(self, messages:List[Dict[str,str]], max_new_tokens:int=4096, temperature:float=0.0, 
+             verbose:bool=False, stream:bool=False, **kwrs) -> Union[str, Generator[str, None, None]]:
+        """
+        This method inputs chat messages and outputs VLM generated text.
+
+        Parameters:
+        ----------
+        messages : List[Dict[str,str]]
+            a list of dict with role and content. role must be one of {"system", "user", "assistant"}
+        max_new_tokens : str, Optional
+            the max number of new tokens VLM can generate. 
+        temperature : float, Optional
+            the temperature for token sampling. 
+        verbose : bool, Optional
+            if True, VLM generated text will be printed in terminal in real-time.
+        stream : bool, Optional
+            if True, returns a generator that yields the output in real-time.
+        """
+        return NotImplemented
+    
+    @abc.abstractmethod
+    def chat_async(self, messages:List[Dict[str,str]], max_new_tokens:int=4096, temperature:float=0.0, **kwrs) -> str:
+        """
+        The async version of chat method. Streaming is not supported.
+        """
+        return NotImplemented
+
+    @abc.abstractmethod
+    def get_ocr_messages(self, system_prompt:str, user_prompt:str, image_path:str) -> 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_path : str
+            the image path for OCR.
+        """
+        return NotImplemented
+
+
+class OllamaVLMEngine(VLMEngine):
+    def __init__(self, model_name:str, num_ctx:int=4096, keep_alive:int=300, **kwrs):
+        """
+        The Ollama inference engine.
+
+        Parameters:
+        ----------
+        model_name : str
+            the model name exactly as shown in >> ollama ls
+        num_ctx : int, Optional
+            context length that VLM will evaluate.
+        keep_alive : int, Optional
+            seconds to hold the VLM after the last API call.
+        """
+        if importlib.util.find_spec("ollama") is None:
+            raise ImportError("ollama-python not found. Please install ollama-python (```pip install ollama```).")
+        
+        from ollama import Client, AsyncClient
+        self.client = Client(**kwrs)
+        self.async_client = AsyncClient(**kwrs)
+        self.model_name = model_name
+        self.num_ctx = num_ctx
+        self.keep_alive = keep_alive
+
+    def chat(self, messages:List[Dict[str,str]], max_new_tokens:int=4096, temperature:float=0.0, 
+             verbose:bool=False, stream:bool=False, **kwrs) -> Union[str, Generator[str, None, None]]:
+        """
+        This method inputs chat messages and outputs VLM generated text.
+
+        Parameters:
+        ----------
+        messages : List[Dict[str,str]]
+            a list of dict with role and content. role must be one of {"system", "user", "assistant"}
+        max_new_tokens : str, Optional
+            the max number of new tokens VLM can generate. 
+        temperature : float, Optional
+            the temperature for token sampling. 
+        verbose : bool, Optional
+            if True, VLM generated text will be printed in terminal in real-time.
+        stream : bool, Optional
+            if True, returns a generator that yields the output in real-time.
+        """
+        options={'temperature':temperature, 'num_ctx': self.num_ctx, 'num_predict': max_new_tokens, **kwrs}
+        if stream:
+            def _stream_generator():
+                response_stream = self.client.chat(
+                    model=self.model_name, 
+                    messages=messages, 
+                    options=options,
+                    stream=True, 
+                    keep_alive=self.keep_alive
+                )
+                for chunk in response_stream:
+                    content_chunk = chunk.get('message', {}).get('content')
+                    if content_chunk:
+                        yield content_chunk
+
+            return _stream_generator()
+
+        elif verbose:
+            response = self.client.chat(
+                            model=self.model_name, 
+                            messages=messages, 
+                            options=options,
+                            stream=True,
+                            keep_alive=self.keep_alive
+                        )
+            
+            res = ''
+            for chunk in response:
+                content_chunk = chunk.get('message', {}).get('content')
+                print(content_chunk, end='', flush=True)
+                res += content_chunk
+            print('\n')
+            return res
+        
+        return response.get('message', {}).get('content', '')
+    
+
+    async def chat_async(self, messages:List[Dict[str,str]], max_new_tokens:int=4096, temperature:float=0.0, **kwrs) -> str:
+        """
+        Async version of chat method. Streaming is not supported.
+        """
+        response = await self.async_client.chat(
+                            model=self.model_name, 
+                            messages=messages, 
+                            options={'temperature':temperature, 'num_ctx': self.num_ctx, 'num_predict': max_new_tokens, **kwrs},
+                            stream=False,
+                            keep_alive=self.keep_alive
+                        )
+        
+        return response.get('message', {}).get('content', '')
+    
+    def get_ocr_messages(self, system_prompt:str, user_prompt:str, image:Image.Image) -> 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_path : str
+            the image path for OCR.
+        """
+        base64_str = image_to_base64(image)
+        return [
+            {"role": "system", "content": system_prompt},
+            {
+                "role": "user",
+                "content": user_prompt,
+                "images": [base64_str]
+            }
+        ]
+
+
+class OpenAIVLMEngine(VLMEngine):
+    def __init__(self, model:str, reasoning_model:bool=False, **kwrs):
+        """
+        The OpenAI API inference engine. Supports OpenAI models and OpenAI compatible servers:
+        - vLLM OpenAI compatible server (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 described in https://platform.openai.com/docs/models
+        reasoning_model : bool, Optional
+            indicator for OpenAI reasoning models ("o" series).
+        """
+        if importlib.util.find_spec("openai") is None:
+            raise ImportError("OpenAI Python API library not found. Please install OpanAI (```pip install openai```).")
+        
+        from openai import OpenAI, AsyncOpenAI
+        self.client = OpenAI(**kwrs)
+        self.async_client = AsyncOpenAI(**kwrs)
+        self.model = model
+        self.reasoning_model = reasoning_model
+
+    def chat(self, messages:List[Dict[str,str]], max_new_tokens:int=4096, temperature:float=0.0, 
+             verbose:bool=False, stream:bool=False, **kwrs) -> Union[str, Generator[str, None, None]]:
+        """
+        This method inputs chat messages and outputs VLM generated text.
+
+        Parameters:
+        ----------
+        messages : List[Dict[str,str]]
+            a list of dict with role and content. role must be one of {"system", "user", "assistant"}
+        max_new_tokens : str, Optional
+            the max number of new tokens VLM can generate. 
+        temperature : float, Optional
+            the temperature for token sampling. 
+        verbose : bool, Optional
+            if True, VLM generated text will be printed in terminal in real-time. 
+        stream : bool, Optional
+            if True, returns a generator that yields the output in real-time.
+        """
+        # For reasoning models
+        if self.reasoning_model:
+            # Reasoning models do not support temperature parameter
+            if temperature != 0.0:
+                warnings.warn("Reasoning models do not support temperature parameter. Will be ignored.", UserWarning)
+
+            # Reasoning models do not support system prompts
+            if any(msg['role'] == 'system' for msg in messages):
+                warnings.warn("Reasoning models do not support system prompts. Will be ignored.", UserWarning)
+                messages = [msg for msg in messages if msg['role'] != 'system']
+            
+
+            if stream:
+                def _stream_generator():
+                    response_stream = self.client.chat.completions.create(
+                                            model=self.model,
+                                            messages=messages,
+                                            max_completion_tokens=max_new_tokens,
+                                            stream=True,
+                                            **kwrs
+                                        )
+                    for chunk in response_stream:
+                        if len(chunk.choices) > 0:
+                            if chunk.choices[0].delta.content is not None:
+                                yield chunk.choices[0].delta.content
+                            if chunk.choices[0].finish_reason == "length":
+                                warnings.warn("Model stopped generating due to context length limit.", RuntimeWarning)
+                                if self.reasoning_model:
+                                    warnings.warn("max_new_tokens includes reasoning tokens and output tokens.", UserWarning)
+                return _stream_generator()
+
+            elif verbose:
+                response = self.client.chat.completions.create(
+                    model=self.model,
+                    messages=messages,
+                    max_completion_tokens=max_new_tokens,
+                    stream=True,
+                    **kwrs
+                )
+                res = ''
+                for chunk in response:
+                    if len(chunk.choices) > 0:
+                        if chunk.choices[0].delta.content is not None:
+                            res += chunk.choices[0].delta.content
+                            print(chunk.choices[0].delta.content, end="", flush=True)
+                        if chunk.choices[0].finish_reason == "length":
+                            warnings.warn("Model stopped generating due to context length limit.", RuntimeWarning)
+                            if self.reasoning_model:
+                                warnings.warn("max_new_tokens includes reasoning tokens and output tokens.", UserWarning)
+
+                print('\n')
+                return res
+            else:
+                response = self.client.chat.completions.create(
+                    model=self.model,
+                    messages=messages,
+                    max_completion_tokens=max_new_tokens,
+                    stream=False,
+                    **kwrs
+                )
+                return response.choices[0].message.content
+
+        # For non-reasoning models
+        else:
+            if stream:
+                def _stream_generator():
+                    response_stream = self.client.chat.completions.create(
+                        model=self.model,
+                        messages=messages,
+                        max_tokens=max_new_tokens,
+                        temperature=temperature,
+                        stream=True,
+                        **kwrs
+                    )
+                    for chunk in response_stream:
+                        if len(chunk.choices) > 0:
+                            if chunk.choices[0].delta.content is not None:
+                                yield chunk.choices[0].delta.content
+                            if chunk.choices[0].finish_reason == "length":
+                                warnings.warn("Model stopped generating due to context length limit.", RuntimeWarning)
+                                if self.reasoning_model:
+                                    warnings.warn("max_new_tokens includes reasoning tokens and output tokens.", UserWarning)
+                return _stream_generator()
+            
+            elif verbose:
+                response = self.client.chat.completions.create(
+                    model=self.model,
+                    messages=messages,
+                    max_tokens=max_new_tokens,
+                    temperature=temperature,
+                    stream=True,
+                    **kwrs
+                )
+                res = ''
+                for chunk in response:
+                    if len(chunk.choices) > 0:
+                        if chunk.choices[0].delta.content is not None:
+                            res += chunk.choices[0].delta.content
+                            print(chunk.choices[0].delta.content, end="", flush=True)
+                        if chunk.choices[0].finish_reason == "length":
+                            warnings.warn("Model stopped generating due to context length limit.", RuntimeWarning)
+                            if self.reasoning_model:
+                                warnings.warn("max_new_tokens includes reasoning tokens and output tokens.", UserWarning)
+
+                print('\n')
+                return res
+
+            else:
+                response = self.client.chat.completions.create(
+                    model=self.model,
+                    messages=messages,
+                    max_tokens=max_new_tokens,
+                    temperature=temperature,
+                    stream=False,
+                    **kwrs
+                )
+
+            return response.choices[0].message.content
+    
+
+    async def chat_async(self, messages:List[Dict[str,str]], max_new_tokens:int=4096, temperature:float=0.0, **kwrs) -> str:
+        """
+        Async version of chat method. Streaming is not supported.
+        """
+        if self.reasoning_model:
+            # Reasoning models do not support temperature parameter
+            if temperature != 0.0:
+                warnings.warn("Reasoning models do not support temperature parameter. Will be ignored.", UserWarning)
+
+            # Reasoning models do not support system prompts
+            if any(msg['role'] == 'system' for msg in messages):
+                warnings.warn("Reasoning models do not support system prompts. Will be ignored.", UserWarning)
+                messages = [msg for msg in messages if msg['role'] != 'system']
+
+            response = await self.async_client.chat.completions.create(
+                model=self.model,
+                messages=messages,
+                max_completion_tokens=max_new_tokens,
+                stream=False,
+                **kwrs
+            )
+
+        else:
+            response = await self.async_client.chat.completions.create(
+                model=self.model,
+                messages=messages,
+                max_tokens=max_new_tokens,
+                temperature=temperature,
+                stream=False,
+                **kwrs
+            )
+        
+        if response.choices[0].finish_reason == "length":
+            warnings.warn("Model stopped generating due to context length limit.", RuntimeWarning)
+            if self.reasoning_model:
+                warnings.warn("max_new_tokens includes reasoning tokens and output tokens.", UserWarning)
+
+        return response.choices[0].message.content
+    
+    def get_ocr_messages(self, system_prompt:str, user_prompt:str, image:Image.Image, format:str='png', detail:str="high") -> 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". 
+        """
+        base64_str = image_to_base64(image)
+        return [
+            {"role": "system", "content": system_prompt},
+            {
+                "role": "user",
+                "content": [
+                    {
+                        "type": "image_url",
+                        "image_url": {
+                            "url": f"data:image/{format};base64,{base64_str}",
+                            "detail": detail
+                        },
+                    },
+                    {"type": "text", "text": user_prompt},
+                ],
+            },
+        ]
+
+
+class AzureOpenAIVLMEngine(OpenAIVLMEngine):
+    def __init__(self, model:str, api_version:str, reasoning_model:bool=False, **kwrs):
+        """
+        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
+        reasoning_model : bool, Optional
+            indicator for OpenAI reasoning models ("o" series).
+        """
+        if importlib.util.find_spec("openai") is None:
+            raise ImportError("OpenAI Python API library not found. Please install OpanAI (```pip install openai```).")
+        
+        from openai import AzureOpenAI, AsyncAzureOpenAI
+        self.model = model
+        self.api_version = api_version
+        self.client = AzureOpenAI(api_version=self.api_version, 
+                                  **kwrs)
+        self.async_client = AsyncAzureOpenAI(api_version=self.api_version, 
+                                             **kwrs)
+        self.reasoning_model = reasoning_model

vlm4ocr-0.0.1.dist-info/METADATA

@@ -0,0 +1,16 @@
+Metadata-Version: 2.1
+Name: vlm4ocr
+Version: 0.0.1
+Summary: OCR with vision language models.
+License: MIT
+Author: Enshuo (David) Hsu
+Requires-Python: >=3.11,<4.0
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Requires-Dist: pdf2image (>=1.16.0)
+Requires-Dist: pillow (>=10.0.0)
+Description-Content-Type: text/markdown
+
+This is the readme for vlm4ocr Python package.

vlm4ocr-0.0.1.dist-info/RECORD

@@ -0,0 +1,10 @@
+vlm4ocr/__init__.py,sha256=gD7WuuxH9_iUu8dgbfonRQD89xiveVP3e9YP4RAHeIQ,219
+vlm4ocr/assets/default_prompt_templates/ocr_markdown_system_prompt.txt,sha256=pIsYO2G3jkZ5EWg7MJixre3Itz1oPqJSduUZT34_RNY,436
+vlm4ocr/assets/default_prompt_templates/ocr_text_system_prompt.txt,sha256=WbLSOerqFjlYGaGWJ-w2enhky1WhnPl011s0fgRPgnQ,398
+vlm4ocr/assets/default_prompt_templates/ocr_user_prompt.txt,sha256=61EJv8POsQGIIUVwCjDU73lMXJE7F3qhPIYl6zSbl1Q,45
+vlm4ocr/ocr_engines.py,sha256=gNwN_itHqIQLF3cfMoGV51lAuY8ZJI2A78LWJivqVi4,13687
+vlm4ocr/utils.py,sha256=3o6TLY8YIEorYHluHrNoJesEM8th89uRO_KNEzfdDA8,1763
+vlm4ocr/vlm_engines.py,sha256=3Aymjo3xa-CTHbktOIHSeLSLAdbo7fk9qRw8d3uRmOM,18836
+vlm4ocr-0.0.1.dist-info/METADATA,sha256=_8E0__RQ6bJy_WLlkirJRvf3G8LqmG5_WtOF1ddBaZw,520
+vlm4ocr-0.0.1.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
+vlm4ocr-0.0.1.dist-info/RECORD,,

vlm4ocr-0.0.1.dist-info/WHEEL

@@ -0,0 +1,4 @@
+Wheel-Version: 1.0
+Generator: poetry-core 1.9.0
+Root-Is-Purelib: true
+Tag: py3-none-any