kreuzberg 3.11.4__py3-none-any.whl → 3.13.1__py3-none-any.whl

kreuzberg/_gmft.py

@@ -7,16 +7,17 @@ import queue
 import signal
 import time
 import traceback
-from dataclasses import dataclass, field
 from io import StringIO
 from pathlib import Path
-from typing import TYPE_CHECKING, Any, Literal
+from typing import TYPE_CHECKING, Any
 
 import anyio
 import msgspec
+import pandas as pd
 from PIL import Image
 
-from kreuzberg._types import TableData
+from kreuzberg._types import GMFTConfig, TableData
+from kreuzberg._utils._cache import get_table_cache
 from kreuzberg._utils._sync import run_sync
 from kreuzberg.exceptions import MissingDependencyError, ParsingError
 
@@ -27,139 +28,9 @@ if TYPE_CHECKING:
     from pandas import DataFrame
 
 
-@dataclass(unsafe_hash=True, slots=True)
-class GMFTConfig:
-    """Configuration options for GMFT.
-
-    This class encapsulates the configuration options for GMFT, providing a way to customize its behavior.
-    """
-
-    verbosity: int = 0
-    """
-    Verbosity level for logging.
-
-    0: errors only
-    1: print warnings
-    2: print warnings and info
-    3: print warnings, info, and debug
-    """
-    formatter_base_threshold: float = 0.3
-    """
-    Base threshold for the confidence demanded of a table feature (row/column).
-
-    Note that a low threshold is actually better, because overzealous rows means that generally, numbers are still aligned and there are just many empty rows (having fewer rows than expected merges cells, which is bad).
-    """
-    cell_required_confidence: dict[Literal[0, 1, 2, 3, 4, 5, 6], float] = field(
-        default_factory=lambda: {
-            0: 0.3,
-            1: 0.3,
-            2: 0.3,
-            3: 0.3,
-            4: 0.5,
-            5: 0.5,
-            6: 99,
-        },
-        hash=False,
-    )
-    """
-    Confidences required (>=) for a row/column feature to be considered good. See TATRFormattedTable.id2label
-
-    But low confidences may be better than too high confidence (see formatter_base_threshold)
-    """
-    detector_base_threshold: float = 0.9
-    """Minimum confidence score required for a table"""
-    remove_null_rows: bool = True
-    """
-    Flag to remove rows with no text.
-    """
-    enable_multi_header: bool = False
-    """
-    Enable multi-indices in the dataframe.
-
-    If false, then multiple headers will be merged column-wise.
-    """
-    semantic_spanning_cells: bool = False
-    """
-    [Experimental] Enable semantic spanning cells, which often encode hierarchical multi-level indices.
-    """
-    semantic_hierarchical_left_fill: Literal["algorithm", "deep"] | None = "algorithm"
-    """
-    [Experimental] When semantic spanning cells is enabled, when a left header is detected which might represent a group of rows, that same value is reduplicated for each row.
-
-    Possible values: 'algorithm', 'deep', None.
-
-    'algorithm': assumes that the higher-level header is always the first row followed by several empty rows.
-    'deep': merges headers according to the spanning cells detected by the Table Transformer.
-    None: headers are not duplicated.
-    """
-    large_table_if_n_rows_removed: int = 8
-    """
-    If >= n rows are removed due to non-maxima suppression (NMS), then this table is classified as a large table.
-    """
-    large_table_threshold: int = 10
-    """
-    With large tables, table transformer struggles with placing too many overlapping rows. Luckily, with more rows, we have more info on the usual size of text, which we can use to make a guess on the height such that no rows are merged or overlapping.
-
-    Large table assumption is only applied when (# of rows > large_table_threshold) AND (total overlap > large_table_row_overlap_threshold). Set 9999 to disable; set 0 to force large table assumption to run every time.
-    """
-    large_table_row_overlap_threshold: float = 0.2
-    """
-    With large tables, table transformer struggles with placing too many overlapping rows. Luckily, with more rows, we have more info on the usual size of text, which we can use to make a guess on the height such that no rows are merged or overlapping.
-
-    Large table assumption is only applied when (# of rows > large_table_threshold) AND (total overlap > large_table_row_overlap_threshold).
-    """
-    large_table_maximum_rows: int = 1000
-    """
-    Maximum number of rows allowed for a large table.
-    """
-    force_large_table_assumption: bool | None = None
-    """
-    Force the large table assumption to be applied, regardless of the number of rows and overlap.
-    """
-    total_overlap_reject_threshold: float = 0.9
-    """
-    Reject if total overlap is > 90% of table area.
-    """
-    total_overlap_warn_threshold: float = 0.1
-    """
-    Warn if total overlap is > 10% of table area.
-    """
-    nms_warn_threshold: int = 5
-    """
-    Warn if non maxima suppression removes > 5 rows.
-    """
-    iob_reject_threshold: float = 0.05
-    """
-    Reject if iob between textbox and cell is < 5%.
-    """
-    iob_warn_threshold: float = 0.5
-    """
-    Warn if iob between textbox and cell is < 50%.
-    """
-
-
 async def extract_tables(
     file_path: str | PathLike[str], config: GMFTConfig | None = None, use_isolated_process: bool | None = None
 ) -> list[TableData]:
-    """Extracts tables from a PDF file.
-
-    This function takes a file path to a PDF file, and an optional configuration object.
-    It returns a list of strings, where each string is a markdown-formatted table.
-
-    Args:
-        file_path: The path to the PDF file.
-        config: An optional configuration object.
-        use_isolated_process: Whether to use an isolated process for extraction.
-            If None, uses environment variable KREUZBERG_GMFT_ISOLATED (default: True).
-
-    Raises:
-        MissingDependencyError: Raised when the required dependencies are not installed.
-
-    Returns:
-        A list of table data dictionaries.
-    """
-    from kreuzberg._utils._cache import get_table_cache  # noqa: PLC0415
-
     # Determine if we should use isolated process  # ~keep
     if use_isolated_process is None:
         use_isolated_process = os.environ.get("KREUZBERG_GMFT_ISOLATED", "true").lower() in ("true", "1", "yes")
@@ -211,15 +82,15 @@ async def extract_tables(
             return result
 
         try:
-            from gmft.auto import (  # type: ignore[attr-defined]  # noqa: PLC0415  # noqa: PLC0415
+            from gmft.auto import (  # type: ignore[attr-defined]  # noqa: PLC0415
                 AutoTableDetector,
                 AutoTableFormatter,
             )
             from gmft.detectors.tatr import TATRDetectorConfig  # type: ignore[attr-defined]  # noqa: PLC0415
-            from gmft.formatters.tatr import TATRFormatConfig  # noqa: PLC0415  # noqa: PLC0415
-            from gmft.pdf_bindings.pdfium import PyPDFium2Document  # noqa: PLC0415  # noqa: PLC0415
+            from gmft.formatters.tatr import TATRFormatConfig  # noqa: PLC0415
+            from gmft.pdf_bindings.pdfium import PyPDFium2Document  # noqa: PLC0415
 
-            formatter: Any = AutoTableFormatter(  # type: ignore[no-untyped-call]  # type: ignore[no-untyped-call]
+            formatter: Any = AutoTableFormatter(  # type: ignore[no-untyped-call]
                 config=TATRFormatConfig(
                     verbosity=config.verbosity,
                     formatter_base_threshold=config.formatter_base_threshold,
@@ -235,7 +106,7 @@ async def extract_tables(
                     force_large_table_assumption=config.force_large_table_assumption,
                 )
             )
-            detector: Any = AutoTableDetector(  # type: ignore[no-untyped-call]  # type: ignore[no-untyped-call]
+            detector: Any = AutoTableDetector(  # type: ignore[no-untyped-call]
                 config=TATRDetectorConfig(detector_base_threshold=config.detector_base_threshold)
             )
             doc = await run_sync(PyPDFium2Document, str(file_path))
@@ -276,19 +147,6 @@ async def extract_tables(
 def extract_tables_sync(
     file_path: str | PathLike[str], config: GMFTConfig | None = None, use_isolated_process: bool | None = None
 ) -> list[TableData]:
-    """Synchronous wrapper for extract_tables.
-
-    Args:
-        file_path: The path to the PDF file.
-        config: An optional configuration object.
-        use_isolated_process: Whether to use an isolated process for extraction.
-            If None, uses environment variable KREUZBERG_GMFT_ISOLATED (default: True).
-
-    Returns:
-        A list of table data dictionaries.
-    """
-    from kreuzberg._utils._cache import get_table_cache  # noqa: PLC0415
-
     # Determine if we should use isolated process  # ~keep
     if use_isolated_process is None:
         use_isolated_process = os.environ.get("KREUZBERG_GMFT_ISOLATED", "true").lower() in ("true", "1", "yes")
@@ -390,13 +248,6 @@ def _extract_tables_in_process(
     config_dict: dict[str, Any],
     result_queue: queue.Queue[tuple[bool, Any]],
 ) -> None:
-    """Extract tables in an isolated process to handle potential segfaults.
-
-    Args:
-        file_path: Path to the PDF file
-        config_dict: Serialized GMFTConfig as a dict
-        result_queue: Queue to put results or errors
-    """
     signal.signal(signal.SIGINT, signal.SIG_IGN)
 
     try:
@@ -480,19 +331,6 @@ def _extract_tables_isolated(
     config: GMFTConfig | None = None,
     timeout: float = 300.0,
 ) -> list[TableData]:
-    """Extract tables using an isolated process to handle segfaults.
-
-    Args:
-        file_path: Path to the PDF file
-        config: GMFT configuration
-        timeout: Maximum time to wait for extraction
-
-    Returns:
-        List of extracted tables
-
-    Raises:
-        RuntimeError: If extraction fails or times out
-    """
     config = config or GMFTConfig()
     config_dict = msgspec.to_builtins(config)
 
@@ -542,7 +380,6 @@ def _extract_tables_isolated(
             tables = []
             for table_dict in result:
                 img = Image.open(io.BytesIO(table_dict["cropped_image_bytes"]))
-                import pandas as pd  # noqa: PLC0415
 
                 if table_dict["df_csv"] is None:
                     df = pd.DataFrame(columns=table_dict["df_columns"])
@@ -592,19 +429,6 @@ async def _extract_tables_isolated_async(
     config: GMFTConfig | None = None,
     timeout: float = 300.0,  # noqa: ASYNC109
 ) -> list[TableData]:
-    """Async version of extract_tables_isolated using asyncio.
-
-    Args:
-        file_path: Path to the PDF file
-        config: GMFT configuration
-        timeout: Maximum time to wait for extraction
-
-    Returns:
-        List of extracted tables
-
-    Raises:
-        RuntimeError: If extraction fails or times out
-    """
     config = config or GMFTConfig()
     config_dict = msgspec.to_builtins(config)
 
@@ -620,38 +444,29 @@ async def _extract_tables_isolated_async(
 
     try:
 
-        async def wait_for_result() -> tuple[bool, Any]:
+        def get_result_sync() -> tuple[bool, Any]:
             while True:
                 try:
-                    return result_queue.get_nowait()  # type: ignore[no-any-return]
+                    return result_queue.get(timeout=0.1)  # type: ignore[no-any-return]
                 except queue.Empty:  # noqa: PERF203
-                    await anyio.sleep(0.1)
                     if not process.is_alive():
-                        # Process died without putting result  # ~keep
                         if process.exitcode == -signal.SIGSEGV:
                             raise ParsingError(
                                 "GMFT process crashed with segmentation fault",
-                                context={
-                                    "file_path": str(file_path),
-                                    "exit_code": process.exitcode,
-                                },
+                                context={"file_path": str(file_path), "exit_code": process.exitcode},
                             ) from None
                         raise ParsingError(
                             f"GMFT process died unexpectedly with exit code {process.exitcode}",
-                            context={
-                                "file_path": str(file_path),
-                                "exit_code": process.exitcode,
-                            },
+                            context={"file_path": str(file_path), "exit_code": process.exitcode},
                         ) from None
 
         with anyio.fail_after(timeout):
-            success, result = await wait_for_result()
+            success, result = await anyio.to_thread.run_sync(get_result_sync)
 
         if success:
             tables = []
             for table_dict in result:
                 img = Image.open(io.BytesIO(table_dict["cropped_image_bytes"]))
-                import pandas as pd  # noqa: PLC0415
 
                 if table_dict["df_csv"] is None:
                     df = pd.DataFrame(columns=table_dict["df_columns"])

kreuzberg/_language_detection.py

@@ -1,9 +1,9 @@
 from __future__ import annotations
 
-from dataclasses import dataclass
 from functools import lru_cache
 from typing import TYPE_CHECKING, Any
 
+from kreuzberg._types import LanguageDetectionConfig
 from kreuzberg.exceptions import MissingDependencyError
 
 if TYPE_CHECKING:
@@ -23,29 +23,7 @@ except ImportError:  # pragma: no cover
 _CACHE_SIZE = 128
 
 
-@dataclass(frozen=True, slots=True)
-class LanguageDetectionConfig:
-    """Configuration for language detection.
-
-    Attributes:
-        low_memory: If True, uses a smaller model (~200MB). If False, uses a larger, more accurate model.
-            Defaults to True for better memory efficiency.
-        top_k: Maximum number of languages to return for multilingual detection. Defaults to 3.
-        multilingual: If True, uses multilingual detection to handle mixed-language text.
-            If False, uses single language detection. Defaults to False.
-        cache_dir: Custom directory for model cache. If None, uses system default.
-        allow_fallback: If True, falls back to small model if large model fails. Defaults to True.
-    """
-
-    low_memory: bool = True
-    top_k: int = 3
-    multilingual: bool = False
-    cache_dir: str | None = None
-    allow_fallback: bool = True
-
-
 def _create_fast_langdetect_config(config: LanguageDetectionConfig) -> FastLangDetectConfig | None:
-    """Create FastLangDetectConfig from our config."""
     if not HAS_FAST_LANGDETECT or FastLangDetectConfig is None:
         return None
 
@@ -60,19 +38,6 @@ def _create_fast_langdetect_config(config: LanguageDetectionConfig) -> FastLangD
 
 @lru_cache(maxsize=_CACHE_SIZE)
 def detect_languages(text: str, config: LanguageDetectionConfig | None = None) -> list[str] | None:
-    """Detect the most probable languages in the given text using fast-langdetect.
-
-    Args:
-        text: The text to analyze.
-        config: Configuration for language detection. If None, uses defaults.
-
-    Returns:
-        A list of detected language codes in lowercase (e.g., ['en', 'de', 'fr']),
-        or None if detection fails.
-
-    Raises:
-        MissingDependencyError: If fast-langdetect is not installed.
-    """
     if not HAS_FAST_LANGDETECT or detect is None or detect_multilingual is None:
         raise MissingDependencyError.create_for_package(
             dependency_group="langdetect", functionality="language detection", package_name="fast-langdetect"

kreuzberg/_mcp/__init__.py

@@ -1,5 +1,3 @@
-"""MCP server for Kreuzberg text extraction."""
-
 from .server import mcp
 
 __all__ = ["mcp"]

kreuzberg/_mcp/server.py

@@ -1,5 +1,3 @@
-"""Kreuzberg MCP server implementation."""
-
 from __future__ import annotations
 
 import base64
@@ -10,11 +8,10 @@ import msgspec
 from mcp.server import FastMCP
 from mcp.types import TextContent
 
-from kreuzberg._config import try_discover_config
+from kreuzberg._config import discover_config
 from kreuzberg._types import ExtractionConfig, OcrBackendType
 from kreuzberg.extraction import extract_bytes_sync, extract_file_sync
 
-# Create the MCP server
 mcp = FastMCP("Kreuzberg Text Extraction")
 
 
@@ -27,14 +24,11 @@ def _create_config_with_overrides(**kwargs: Any) -> ExtractionConfig:
     Returns:
         ExtractionConfig instance.
     """
-    # Try to discover configuration from files
-    base_config = try_discover_config()
+    base_config = discover_config()
 
     if base_config is None:
-        # No config file found, use defaults
         return ExtractionConfig(**kwargs)
 
-    # Merge discovered config with tool parameters (tool params take precedence)
     config_dict: dict[str, Any] = {
         "force_ocr": base_config.force_ocr,
         "chunk_content": base_config.chunk_content,
@@ -50,7 +44,6 @@ def _create_config_with_overrides(**kwargs: Any) -> ExtractionConfig:
         "gmft_config": base_config.gmft_config,
     }
 
-    # Override with provided parameters
     config_dict = config_dict | kwargs
 
     return ExtractionConfig(**config_dict)
@@ -189,7 +182,7 @@ def get_default_config() -> str:
 @mcp.resource("config://discovered")
 def get_discovered_config() -> str:
     """Get the discovered configuration from config files."""
-    config = try_discover_config()
+    config = discover_config()
     if config is None:
         return "No configuration file found"
     return json.dumps(msgspec.to_builtins(config, order="deterministic"), indent=2)

kreuzberg/_mime_types.py

@@ -4,6 +4,7 @@ from mimetypes import guess_type
 from pathlib import Path
 from typing import TYPE_CHECKING, Final
 
+from kreuzberg._utils._cache import get_mime_cache
 from kreuzberg.exceptions import ValidationError
 
 if TYPE_CHECKING:  # pragma: no cover
@@ -172,27 +173,10 @@ SUPPORTED_MIME_TYPES: Final[set[str]] = (
 def validate_mime_type(
     *, file_path: PathLike[str] | str | None = None, mime_type: str | None = None, check_file_exists: bool = True
 ) -> str:
-    """Validate and detect the MIME type for a given file.
-
-    Args:
-        file_path: The path to the file.
-        mime_type: Optional explicit MIME type. If provided, this will be validated.
-            If not provided, the function will attempt to detect the MIME type.
-        check_file_exists: Whether to check if the file exists. Default is True.
-            Set to False in tests where you want to validate a mime type without an actual file.
-
-    Raises:
-        ValidationError: If the MIME type is not supported or cannot be determined.
-
-    Returns:
-        The validated MIME type.
-    """
     if mime_type:
         return _validate_explicit_mime_type(mime_type)
 
     if file_path:
-        from kreuzberg._utils._cache import get_mime_cache  # noqa: PLC0415
-
         path = Path(file_path)
 
         try:
@@ -228,7 +212,6 @@ def validate_mime_type(
 
 
 def _validate_explicit_mime_type(mime_type: str) -> str:
-    """Validate an explicitly provided MIME type."""
     if mime_type in SUPPORTED_MIME_TYPES:
         return mime_type
 
@@ -243,7 +226,6 @@ def _validate_explicit_mime_type(mime_type: str) -> str:
 
 
 def _detect_mime_type_uncached(file_path: PathLike[str] | str | None = None, check_file_exists: bool = True) -> str:
-    """Detect MIME type without caching (internal function)."""
     if file_path and check_file_exists:
         path = Path(file_path)
         if not path.exists():

kreuzberg/_ocr/_base.py

@@ -16,98 +16,26 @@ T = TypeVar("T")
 
 
 class OCRBackend(ABC, Generic[T]):
-    """Abstract base class for Optical Character Recognition (OCR) backend implementations.
-
-    This class provides the blueprint for OCR backend implementations,
-    offering both synchronous and asynchronous methods to process images
-    and files for text extraction.
-    """
-
     @abstractmethod
-    async def process_image(self, image: Image, **kwargs: Unpack[T]) -> ExtractionResult:
-        """Asynchronously process an image and extract its text and metadata.
-
-        Args:
-            image: An instance of PIL.Image representing the input image.
-            **kwargs: Any kwargs related to the given backend
-
-        Returns:
-            The extraction result object
-        """
-        ...
+    async def process_image(self, image: Image, **kwargs: Unpack[T]) -> ExtractionResult: ...
 
     @abstractmethod
-    async def process_file(self, path: Path, **kwargs: Unpack[T]) -> ExtractionResult:
-        """Asynchronously process a file and extract its text and metadata.
-
-        Args:
-            path: A Path object representing the file to be processed.
-            **kwargs: Any kwargs related to the given backend
-
-        Returns:
-            The extraction result object
-        """
-        ...
+    async def process_file(self, path: Path, **kwargs: Unpack[T]) -> ExtractionResult: ...
 
     @abstractmethod
-    def process_image_sync(self, image: Image, **kwargs: Unpack[T]) -> ExtractionResult:
-        """Synchronously process an image and extract its text and metadata.
-
-        Args:
-            image: An instance of PIL.Image representing the input image.
-            **kwargs: Any kwargs related to the given backend
-
-        Returns:
-            The extraction result object
-        """
-        ...
+    def process_image_sync(self, image: Image, **kwargs: Unpack[T]) -> ExtractionResult: ...
 
     @abstractmethod
-    def process_file_sync(self, path: Path, **kwargs: Unpack[T]) -> ExtractionResult:
-        """Synchronously process a file and extract its text and metadata.
-
-        Args:
-            path: A Path object representing the file to be processed.
-            **kwargs: Any kwargs related to the given backend
-
-        Returns:
-            The extraction result object
-        """
-        ...
+    def process_file_sync(self, path: Path, **kwargs: Unpack[T]) -> ExtractionResult: ...
 
     def process_batch_sync(self, paths: list[Path], **kwargs: Unpack[T]) -> list[ExtractionResult]:
-        """Synchronously process a batch of files and extract their text and metadata.
-
-        Default implementation processes files sequentially. Backends can override
-        for more efficient batch processing.
-
-        Args:
-            paths: List of Path objects representing files to be processed.
-            **kwargs: Any kwargs related to the given backend
-
-        Returns:
-            List of extraction result objects in the same order as input paths
-        """
         return [self.process_file_sync(path, **kwargs) for path in paths]  # pragma: no cover
 
     async def process_batch(self, paths: list[Path], **kwargs: Unpack[T]) -> list[ExtractionResult]:
-        """Asynchronously process a batch of files and extract their text and metadata.
-
-        Default implementation processes files concurrently. Backends can override
-        for more efficient batch processing.
-
-        Args:
-            paths: List of Path objects representing files to be processed.
-            **kwargs: Any kwargs related to the given backend
-
-        Returns:
-            List of extraction result objects in the same order as input paths
-        """
         from kreuzberg._utils._sync import run_taskgroup  # noqa: PLC0415
 
         tasks = [self.process_file(path, **kwargs) for path in paths]
         return await run_taskgroup(*tasks)  # pragma: no cover
 
     def __hash__(self) -> int:
-        """Hash function for allowing caching."""
         return hash(type(self).__name__)  # pragma: no cover