xgen-doc2chunk 0.1.0__py3-none-any.whl → 0.1.2__py3-none-any.whl

xgen_doc2chunk/ocr/ocr_engine/openai_ocr.py

@@ -0,0 +1,100 @@
+# xgen_doc2chunk/ocr/ocr_engine/openai_ocr.py
+# OCR class using OpenAI Vision model
+import logging
+from typing import Any, Optional
+
+from xgen_doc2chunk.ocr.base import BaseOCR
+
+logger = logging.getLogger("ocr-openai")
+
+# Default model
+DEFAULT_OPENAI_MODEL = "gpt-4o"
+
+
+class OpenAIOCR(BaseOCR):
+    """
+    OCR processing class using OpenAI Vision model.
+
+    Supported models: gpt-4-vision-preview, gpt-4o, gpt-4o-mini, etc.
+
+    Example:
+        ```python
+        from xgen_doc2chunk.ocr.ocr_engine import OpenAIOCR
+
+        # Method 1: Initialize with api_key and model
+        ocr = OpenAIOCR(api_key="sk-...", model="gpt-4o")
+
+        # Method 2: Use existing LLM client
+        from langchain_openai import ChatOpenAI
+        llm = ChatOpenAI(model="gpt-4o", temperature=0, api_key="sk-...")
+        ocr = OpenAIOCR(llm_client=llm)
+
+        # Single image conversion
+        result = ocr.convert_image_to_text("/path/to/image.png")
+
+        # Process image tags in text
+        text = "Document content [Image:/path/to/image.png] continues..."
+        processed = ocr.process_text(text)
+        ```
+    """
+
+    def __init__(
+        self,
+        api_key: Optional[str] = None,
+        model: str = DEFAULT_OPENAI_MODEL,
+        llm_client: Optional[Any] = None,
+        prompt: Optional[str] = None,
+        temperature: float = 0.0,
+        max_tokens: Optional[int] = None,
+        base_url: Optional[str] = None,
+    ):
+        """
+        Initialize OpenAI OCR.
+
+        Args:
+            api_key: OpenAI API key (required if llm_client is not provided)
+            model: Model name to use (default: gpt-4o)
+            llm_client: Existing LangChain OpenAI client (if provided, api_key and model are ignored)
+            prompt: Custom prompt (if None, default prompt is used)
+            temperature: Generation temperature (default: 0.0)
+            max_tokens: Maximum number of tokens (if None, model default is used)
+            base_url: OpenAI API base URL (for Azure, etc.)
+        """
+        if llm_client is None:
+            if api_key is None:
+                raise ValueError("Either api_key or llm_client is required.")
+
+            from langchain_openai import ChatOpenAI
+
+            client_kwargs = {
+                "model": model,
+                "api_key": api_key,
+                "temperature": temperature,
+            }
+
+            if max_tokens is not None:
+                client_kwargs["max_tokens"] = max_tokens
+
+            if base_url is not None:
+                client_kwargs["base_url"] = base_url
+
+            llm_client = ChatOpenAI(**client_kwargs)
+            logger.info(f"[OpenAI OCR] Client created: model={model}")
+
+        super().__init__(llm_client=llm_client, prompt=prompt)
+        self.model = model
+        logger.info("[OpenAI OCR] Initialization completed")
+
+    @property
+    def provider(self) -> str:
+        return "openai"
+
+    def build_message_content(self, b64_image: str, mime_type: str) -> list:
+        return [
+            {"type": "text", "text": self.prompt},
+            {
+                "type": "image_url",
+                "image_url": {"url": f"data:{mime_type};base64,{b64_image}"}
+            }
+        ]
+

xgen_doc2chunk/ocr/ocr_engine/vllm_ocr.py

@@ -0,0 +1,116 @@
+# xgen_doc2chunk/ocr/ocr_engine/vllm_ocr.py
+# OCR class using vLLM-based Vision model
+import logging
+from typing import Any, Optional
+
+from xgen_doc2chunk.ocr.base import BaseOCR
+
+logger = logging.getLogger("ocr-vllm")
+
+# Default model (varies by user environment)
+DEFAULT_VLLM_MODEL = "Qwen/Qwen2-VL-7B-Instruct"
+
+
+class VllmOCR(BaseOCR):
+    """
+    OCR processing class using vLLM-based Vision model.
+
+    Uses OpenAI-compatible API provided by vLLM server.
+
+    Example:
+        ```python
+        from xgen_doc2chunk.ocr.ocr_engine import VllmOCR
+
+        # Method 1: Initialize with base_url and model
+        ocr = VllmOCR(
+            base_url="http://localhost:8000/v1",
+            model="Qwen/Qwen2-VL-7B-Instruct"
+        )
+
+        # Method 2: When api_key is required
+        ocr = VllmOCR(
+            base_url="http://your-vllm-server:8000/v1",
+            api_key="your-api-key",
+            model="Qwen/Qwen2-VL-7B-Instruct"
+        )
+
+        # Method 3: Use existing LLM client
+        from langchain_openai import ChatOpenAI
+        llm = ChatOpenAI(
+            model="Qwen/Qwen2-VL-7B-Instruct",
+            base_url="http://localhost:8000/v1",
+            api_key="EMPTY"
+        )
+        ocr = VllmOCR(llm_client=llm)
+
+        # Single image conversion
+        result = ocr.convert_image_to_text("/path/to/image.png")
+        ```
+    """
+
+    # vLLM uses simple prompt
+    DEFAULT_PROMPT = "Describe the contents of this image."
+
+    def __init__(
+        self,
+        base_url: Optional[str] = None,
+        api_key: Optional[str] = "EMPTY",
+        model: str = DEFAULT_VLLM_MODEL,
+        llm_client: Optional[Any] = None,
+        prompt: Optional[str] = None,
+        temperature: float = 0.0,
+        max_tokens: Optional[int] = None,
+    ):
+        """
+        Initialize vLLM OCR.
+
+        Args:
+            base_url: vLLM server URL (e.g., "http://localhost:8000/v1")
+            api_key: API key (default: "EMPTY", vLLM default setting)
+            model: Model name to use (default: Qwen/Qwen2-VL-7B-Instruct)
+            llm_client: Existing LangChain client (if provided, base_url, api_key, and model are ignored)
+            prompt: Custom prompt (if None, SIMPLE_PROMPT is used)
+            temperature: Generation temperature (default: 0.0)
+            max_tokens: Maximum number of tokens (if None, model default is used)
+        """
+        # vLLM uses simple prompt by default
+        if prompt is None:
+            prompt = self.DEFAULT_PROMPT
+
+        if llm_client is None:
+            if base_url is None:
+                raise ValueError("Either base_url or llm_client is required.")
+
+            from langchain_openai import ChatOpenAI
+
+            client_kwargs = {
+                "model": model,
+                "base_url": base_url,
+                "api_key": api_key,
+                "temperature": temperature,
+            }
+
+            if max_tokens is not None:
+                client_kwargs["max_tokens"] = max_tokens
+
+            llm_client = ChatOpenAI(**client_kwargs)
+            logger.info(f"[vLLM OCR] Client created: base_url={base_url}, model={model}")
+
+        super().__init__(llm_client=llm_client, prompt=prompt)
+        self.model = model
+        self.base_url = base_url
+        logger.info("[vLLM OCR] Initialization completed")
+
+    @property
+    def provider(self) -> str:
+        return "vllm"
+
+    def build_message_content(self, b64_image: str, mime_type: str) -> list:
+        return [
+            {"type": "text", "text": self.prompt},
+            {
+                "type": "image_url",
+                "image_url": {"url": f"data:{mime_type};base64,{b64_image}"}
+            }
+        ]
+

xgen_doc2chunk/ocr/ocr_processor.py

@@ -0,0 +1,387 @@
+# xgen_doc2chunk/ocr/ocr_processor.py
+# Module for loading image files and processing OCR.
+import re
+import base64
+import logging
+import os
+from typing import Any, Callable, Dict, List, Optional, Pattern
+
+logger = logging.getLogger("ocr-processor")
+
+# Default image tag pattern: [Image:{path}] or [image:{path}] (case-insensitive)
+DEFAULT_IMAGE_TAG_PATTERN = re.compile(r'\[[Ii]mage:([^\]]+)\]')
+
+# Keep for backward compatibility
+IMAGE_TAG_PATTERN = DEFAULT_IMAGE_TAG_PATTERN
+
+
+def _b64_from_file(path: str) -> str:
+    """Encode file to Base64"""
+    with open(path, "rb") as f:
+        return base64.b64encode(f.read()).decode("utf-8")
+
+
+def _get_mime_type(file_path: str) -> str:
+    """Return MIME type based on file extension"""
+    ext = os.path.splitext(file_path)[1].lower()
+    mime_map = {
+        ".jpg": "image/jpeg",
+        ".jpeg": "image/jpeg",
+        ".png": "image/png",
+        ".gif": "image/gif",
+        ".bmp": "image/bmp",
+        ".webp": "image/webp",
+        ".tiff": "image/tiff",
+        ".svg": "image/svg+xml",
+    }
+    return mime_map.get(ext, "image/jpeg")
+
+
+def extract_image_tags(
+    text: str,
+    pattern: Optional[Pattern[str]] = None
+) -> List[str]:
+    """
+    Extract image tags from text.
+
+    Args:
+        text: Text containing image tags
+        pattern: Custom regex pattern with capture group for path.
+                 If None, uses default [Image:{path}] pattern.
+
+    Returns:
+        List of extracted image_path values
+    """
+    if pattern is None:
+        pattern = DEFAULT_IMAGE_TAG_PATTERN
+    matches = pattern.findall(text)
+    return matches
+
+
+def load_image_from_path(image_path: str) -> Optional[str]:
+    """
+    Validate and return local image file path.
+
+    Args:
+        image_path: Image file path
+
+    Returns:
+        Valid local file path or None
+    """
+    try:
+        # Convert to absolute path
+        if not os.path.isabs(image_path):
+            image_path = os.path.abspath(image_path)
+
+        # Check file existence
+        if os.path.exists(image_path) and os.path.getsize(image_path) > 0:
+            logger.info(f"[OCR] Image loaded: {image_path}")
+            return image_path
+
+        logger.warning(f"[OCR] Image file not found: {image_path}")
+        return None
+
+    except Exception as e:
+        logger.error(f"[OCR] Image load failed: {image_path}, error: {e}")
+        return None
+
+
+def convert_image_to_text_with_llm(
+    image_path: str,
+    llm_client: Any,
+    provider: str
+) -> str:
+    """
+    Convert image to text using VL model.
+
+    Args:
+        image_path: Local image file path
+        llm_client: LangChain LLM client
+        provider: LLM provider (openai, anthropic, gemini, vllm)
+
+    Returns:
+        Text extracted from image
+    """
+    try:
+        from langchain_core.messages import HumanMessage
+
+        b64_image = _b64_from_file(image_path)
+        mime_type = _get_mime_type(image_path)
+
+        # vllm uses simple prompt
+        if provider == "vllm":
+            prompt = "Describe the contents of this image."
+        else:
+            prompt = (
+                "Extract meaningful information from this image.\n\n"
+                "**If the image contains a TABLE:**\n"
+                "- Convert to HTML table format (<table>, <tr>, <td>, <th>)\n"
+                "- Use 'rowspan' and 'colspan' attributes for merged cells\n"
+                "- Preserve all cell content exactly as shown\n"
+                "- Example:\n"
+                "  <table>\n"
+                "    <tr><th colspan=\"2\">Header</th></tr>\n"
+                "    <tr><td rowspan=\"2\">Merged</td><td>A</td></tr>\n"
+                "    <tr><td>B</td></tr>\n"
+                "  </table>\n\n"
+                "**If the image contains TEXT (non-table):**\n"
+                "- Extract all text exactly as shown\n"
+                "- Keep layout, hierarchy, and structure\n\n"
+                "**If the image contains DATA (charts, graphs, diagrams):**\n"
+                "- Extract the data and its meaning\n"
+                "- Describe trends, relationships, or key insights\n\n"
+                "**If the image is decorative or has no semantic meaning:**\n"
+                "- Simply state what it is in one short sentence\n"
+                "- Example: 'A decorative geometric shape' or 'Company logo'\n"
+                "- Do NOT over-analyze decorative elements\n\n"
+                "**Rules:**\n"
+                "- Output in Korean (except HTML tags)\n"
+                "- Tables MUST use HTML format with proper rowspan/colspan\n"
+                "- Be concise - only include what is semantically meaningful\n"
+                "- No filler words or unnecessary descriptions"
+            )
+
+        # Build message by provider
+        if provider in ("openai", "vllm"):
+            content = [
+                {"type": "text", "text": prompt},
+                {
+                    "type": "image_url",
+                    "image_url": {"url": f"data:{mime_type};base64,{b64_image}"}
+                }
+            ]
+            message = HumanMessage(content=content)
+
+        elif provider == "anthropic":
+            content = [
+                {
+                    "type": "image",
+                    "source": {
+                        "type": "base64",
+                        "media_type": mime_type,
+                        "data": b64_image
+                    }
+                },
+                {"type": "text", "text": prompt}
+            ]
+            message = HumanMessage(content=content)
+
+        elif provider == "gemini":
+            content = [
+                {"type": "text", "text": prompt},
+                {
+                    "type": "image_url",
+                    "image_url": {"url": f"data:{mime_type};base64,{b64_image}"}
+                }
+            ]
+            message = HumanMessage(content=content)
+
+        elif provider == "aws_bedrock":
+            # AWS Bedrock (Claude via Bedrock) - Uses same format as Anthropic
+            content = [
+                {
+                    "type": "image",
+                    "source": {
+                        "type": "base64",
+                        "media_type": mime_type,
+                        "data": b64_image
+                    }
+                },
+                {"type": "text", "text": prompt}
+            ]
+            message = HumanMessage(content=content)
+
+        else:
+            return None  # Unsupported provider
+
+        response = llm_client.invoke([message])
+        result = response.content.strip()
+
+        # Wrap result in [Figure:...] format
+        result = f"[Figure:{result}]"
+
+        logger.info(f"[OCR] Image to text conversion completed: {os.path.basename(image_path)}")
+        return result
+
+    except Exception as e:
+        logger.error(f"[OCR] Image to text conversion failed: {image_path}, error: {e}")
+        return f"[Image conversion error: {str(e)}]"
+
+
+def process_text_with_ocr(
+    text: str,
+    llm_client: Any,
+    provider: str
+) -> str:
+    """
+    Detect image tags in text and replace with OCR results.
+
+    Args:
+        text: Text containing [Image:{path}] tags
+        llm_client: LangChain LLM client
+        provider: LLM provider
+
+    Returns:
+        Text with image tags replaced by OCR results
+    """
+    if not llm_client:
+        logger.warning("[OCR] Skipping OCR processing: no LLM client")
+        return text
+
+    # Extract image tags
+    image_paths = extract_image_tags(text)
+
+    if not image_paths:
+        logger.debug("[OCR] No image tags found in text")
+        return text
+
+    logger.info(f"[OCR] Detected {len(image_paths)} image tags")
+
+    result_text = text
+
+    for img_path in image_paths:
+        # Case-insensitive tag matching
+        tag_pattern = re.compile(r'\[[Ii]mage:' + re.escape(img_path) + r'\]')
+
+        # Load image from local path
+        local_path = load_image_from_path(img_path)
+
+        if local_path is None:
+            # Keep original tag on load failure
+            logger.warning(f"[OCR] Image load failed, keeping original tag: {img_path}")
+            continue
+
+        # Convert image to text using VL model
+        ocr_result = convert_image_to_text_with_llm(
+            image_path=local_path,
+            llm_client=llm_client,
+            provider=provider
+        )
+
+        # Keep original tag on OCR failure (None or error message)
+        if ocr_result is None or ocr_result.startswith("[Image conversion error:"):
+            logger.warning(f"[OCR] Image conversion failed, keeping original tag: {img_path}")
+            continue
+
+        # Replace tag with OCR result
+        result_text = tag_pattern.sub(ocr_result, result_text)
+        logger.info(f"[OCR] Tag replacement completed: {img_path[:50]}...")
+
+    return result_text
+
+
+def process_text_with_ocr_progress(
+    text: str,
+    llm_client: Any,
+    provider: str,
+    progress_callback: Optional[Callable[[Dict[str, Any]], Any]] = None
+) -> str:
+    """
+    Detect image tags in text and replace with OCR results (with progress callback support).
+
+    Args:
+        text: Text containing [Image:{path}] tags
+        llm_client: LangChain LLM client
+        provider: LLM provider
+        progress_callback: Progress callback function
+
+    Returns:
+        Text with image tags replaced by OCR results
+    """
+    if not llm_client:
+        logger.warning("[OCR] Skipping OCR processing: no LLM client")
+        return text
+
+    # Extract image tags
+    image_paths = extract_image_tags(text)
+
+    if not image_paths:
+        logger.debug("[OCR] No image tags found in text")
+        return text
+
+    total_chunks = len(image_paths)
+    logger.info(f"[OCR] Detected {total_chunks} image tags")
+
+    result_text = text
+    success_count = 0
+    failed_count = 0
+
+    for idx, img_path in enumerate(image_paths):
+        # Progress callback - processing started
+        if progress_callback:
+            progress_callback({
+                'event': 'ocr_tag_processing',
+                'chunk_index': idx,
+                'total_chunks': total_chunks,
+                'image_path': img_path
+            })
+
+        # Case-insensitive tag matching
+        tag_pattern = re.compile(r'\[[Ii]mage:' + re.escape(img_path) + r'\]')
+
+        # Load image from local path
+        local_path = load_image_from_path(img_path)
+
+        if local_path is None:
+            # Keep original tag on load failure
+            logger.warning(f"[OCR] Image load failed, keeping original tag: {img_path}")
+            failed_count += 1
+            if progress_callback:
+                progress_callback({
+                    'event': 'ocr_chunk_processed',
+                    'chunk_index': idx,
+                    'total_chunks': total_chunks,
+                    'status': 'failed',
+                    'error': f'Load failed: {img_path}'
+                })
+            continue
+
+        try:
+            # Convert image to text using VL model
+            ocr_result = convert_image_to_text_with_llm(
+                image_path=local_path,
+                llm_client=llm_client,
+                provider=provider
+            )
+
+            # Keep original tag on OCR failure (None or error message)
+            if ocr_result is None or ocr_result.startswith("[Image conversion error:"):
+                logger.warning(f"[OCR] Image conversion failed, keeping original tag: {img_path}")
+                failed_count += 1
+                if progress_callback:
+                    progress_callback({
+                        'event': 'ocr_chunk_processed',
+                        'chunk_index': idx,
+                        'total_chunks': total_chunks,
+                        'status': 'failed',
+                        'error': ocr_result or 'OCR returned None'
+                    })
+                continue
+
+            # Replace tag with OCR result
+            result_text = tag_pattern.sub(ocr_result, result_text)
+            success_count += 1
+            logger.info(f"[OCR] Tag replacement completed: {img_path[:50]}...")
+
+            if progress_callback:
+                progress_callback({
+                    'event': 'ocr_chunk_processed',
+                    'chunk_index': idx,
+                    'total_chunks': total_chunks,
+                    'status': 'success'
+                })
+
+        except Exception as e:
+            logger.error(f"[OCR] Image processing error: {img_path}, error: {e}")
+            failed_count += 1
+            if progress_callback:
+                progress_callback({
+                    'event': 'ocr_chunk_processed',
+                    'chunk_index': idx,
+                    'total_chunks': total_chunks,
+                    'status': 'failed',
+                    'error': str(e)
+                })
+
+    return result_text
+

{xgen_doc2chunk-0.1.0.dist-info → xgen_doc2chunk-0.1.2.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: xgen-doc2chunk
-Version: 0.1.0
+Version: 0.1.2
 Summary: Convert raw documents into AI-understandable context with intelligent text extraction, table detection, and semantic chunking
 Project-URL: Homepage, https://github.com/master0419/doc2chunk
 Project-URL: Documentation, https://github.com/master0419/doc2chunk#readme