vlmparse 0.1.0__py3-none-any.whl → 0.1.2__py3-none-any.whl

vlmparse/data_model/document.py

@@ -0,0 +1,148 @@
+import os
+import traceback
+import zipfile
+from pathlib import Path
+from typing import Optional
+
+import orjson
+from PIL import Image
+from PIL import Image as PILImage
+from pydantic import Field
+
+from vlmparse.base_model import VLMParseBaseModel
+from vlmparse.utils import from_base64, to_base64
+
+from .box import BoundingBox
+
+
+class ProcessingError(VLMParseBaseModel):
+    module_class: str
+    traceback: str
+
+    @classmethod
+    def from_class(cls, klass):
+        return cls(
+            module_class=type(klass).__name__,
+            traceback=traceback.format_exc(),
+        )
+
+
+class Item(VLMParseBaseModel):
+    category: str
+    box: BoundingBox
+    text: str
+
+
+class Page(VLMParseBaseModel):
+    text: str | None = None
+    raw_response: str | None = None
+    items: list[Item] | None = None
+    error: ProcessingError | None = None
+    buffer_image: Optional[Image.Image | str | dict] = None
+    latency: Optional[float] = None
+    """Time taken to process the page in seconds."""
+
+    @property
+    def image(self):
+        if isinstance(self.buffer_image, dict):
+            from vlmparse.build_doc import convert_specific_page_to_image, resize_image
+
+            image = convert_specific_page_to_image(
+                self.buffer_image["file_path"],
+                self.buffer_image["page_idx"],
+                self.buffer_image["dpi"],
+            )
+            image = resize_image(image, self.buffer_image["max_image_size"])
+            self.buffer_image = image
+
+        if isinstance(self.buffer_image, str):
+            self.buffer_image = from_base64(self.buffer_image)
+        return self.buffer_image
+
+    def get_image_with_boxes(self, layout=False):
+        from PIL import ImageDraw
+
+        from .box import draw_text_of_box
+
+        image = self.image
+
+        if layout:
+            if self.items is None:
+                return image
+            items = self.items
+            for item in items:
+                box = item.box
+
+                draw = ImageDraw.Draw(image)
+                draw.rectangle(
+                    (box.l, box.t, box.r, box.b),
+                    outline=(255, 0, 0),
+                    width=5,
+                )
+
+                image = draw_text_of_box(
+                    image, box.l, box.t, item.category, font_size=40
+                )
+        return image
+
+
+class Document(VLMParseBaseModel):
+    file_path: str
+    pages: list[Page] = []
+    error: ProcessingError | None = None
+    metadata: dict = Field(default_factory=dict)
+    latency: Optional[float] = None
+    """Time taken to process the document in seconds."""
+
+    @property
+    def text(self):
+        return "\n\n".join([page.text for page in self.pages])
+
+    @property
+    def is_error(self):
+        return self.error is not None or any(
+            page.error is not None for page in self.pages
+        )
+
+    def to_zip(
+        self,
+        file_path,
+        overwrite_file: bool = True,
+        image_extension: str = "webp",
+    ):
+        file_path = Path(file_path)
+        os.makedirs(file_path.parent, exist_ok=True)
+        archive_path = str(file_path).removesuffix(".zip") + ".zip"
+
+        if not overwrite_file:
+            assert not os.path.isfile(archive_path)
+
+        def _custom_encoder(x):
+            if isinstance(x, PILImage.Image):
+                return to_base64(x, image_extension)
+            if isinstance(x, str):
+                return x
+            raise TypeError(
+                f"Object of type {type(x).__name__} is not JSON serializable"
+            )
+
+        json_bytes = orjson.dumps(
+            self.model_dump(),
+            default=_custom_encoder,
+            option=orjson.OPT_INDENT_2,
+        )
+
+        with zipfile.ZipFile(
+            archive_path, "w", compression=zipfile.ZIP_DEFLATED
+        ) as zipf:
+            zipf.writestr("data.json", json_bytes)
+
+    @classmethod
+    def from_zip(cls, file_path):
+        with zipfile.ZipFile(file_path, "r") as zipf:
+            if "data.json" not in zipf.namelist():
+                raise FileNotFoundError("data.json not found in the archive")
+
+            json_bytes = zipf.read("data.json")
+            data = orjson.loads(json_bytes)
+        return cls.model_validate(data)

vlmparse/servers/docker_server.py

@@ -0,0 +1,199 @@
+import os
+from typing import Callable
+
+from loguru import logger
+from pydantic import BaseModel, Field
+
+from .utils import docker_server
+
+
+class DockerServerConfig(BaseModel):
+    """Base configuration for deploying a Docker server."""
+
+    model_name: str
+    docker_image: str
+    dockerfile_dir: str | None = None
+    command_args: list[str] = Field(default_factory=list)
+    server_ready_indicators: list[str] = Field(
+        default_factory=lambda: [
+            "Application startup complete",
+            "Uvicorn running",
+            "Starting vLLM API server",
+        ]
+    )
+    docker_port: int = 8056
+    gpu_device_ids: list[str] | None = None
+    container_port: int = 8000
+    environment: dict[str, str] = Field(default_factory=dict)
+    volumes: dict[str, dict] | None = None
+    entrypoint: str | None = None
+    aliases: list[str] = Field(default_factory=list)
+
+    class Config:
+        extra = "allow"
+
+    @property
+    def client_config(self):
+        """Override in subclasses to return appropriate client config."""
+        raise NotImplementedError
+
+    def get_client(self, **kwargs):
+        return self.client_config.get_client(**kwargs)
+
+    def get_server(self, auto_stop: bool = True):
+        return ConverterServer(config=self, auto_stop=auto_stop)
+
+    def get_command(self) -> list[str] | None:
+        """Build command for container. Override in subclasses for specific logic."""
+        return self.command_args if self.command_args else None
+
+    def get_volumes(self) -> dict | None:
+        """Setup volumes for container. Override in subclasses for specific logic."""
+        return self.volumes
+
+    def get_environment(self) -> dict | None:
+        """Setup environment variables. Override in subclasses for specific logic."""
+        return self.environment if self.environment else None
+
+    def get_base_url_suffix(self) -> str:
+        """Return URL suffix (e.g., '/v1' for OpenAI-compatible APIs). Override in subclasses."""
+        return ""
+
+
+DEFAULT_MODEL_NAME = "vllm-model"
+
+
+class VLLMDockerServerConfig(DockerServerConfig):
+    """Configuration for deploying a VLLM Docker server."""
+
+    docker_image: str = "vllm/vllm-openai:latest"
+    default_model_name: str = DEFAULT_MODEL_NAME
+    hf_home_folder: str | None = os.getenv("HF_HOME", None)
+    add_model_key_to_server: bool = False
+    container_port: int = 8000
+    aliases: list[str] = Field(default_factory=list)
+
+    @property
+    def llm_params(self):
+        from vlmparse.clients.openai_converter import LLMParams
+
+        return LLMParams(
+            base_url=f"http://localhost:{self.docker_port}/v1",
+            model_name=self.default_model_name,
+        )
+
+    @property
+    def client_config(self):
+        from vlmparse.clients.openai_converter import OpenAIConverterConfig
+
+        return OpenAIConverterConfig(llm_params=self.llm_params)
+
+    def get_command(self) -> list[str]:
+        """Build VLLM-specific command."""
+        model_key = ["--model"] if self.add_model_key_to_server else []
+        command = (
+            model_key
+            + [
+                self.model_name,
+                "--port",
+                str(self.container_port),
+            ]
+            + self.command_args
+            + ["--served-model-name", self.default_model_name]
+        )
+        return command
+
+    def get_volumes(self) -> dict | None:
+        """Setup volumes for HuggingFace model caching."""
+        if self.hf_home_folder is not None:
+            from pathlib import Path
+
+            return {
+                str(Path(self.hf_home_folder).absolute()): {
+                    "bind": "/root/.cache/huggingface",
+                    "mode": "rw",
+                }
+            }
+        return None
+
+    def get_environment(self) -> dict | None:
+        """Setup environment variables for VLLM."""
+        if self.hf_home_folder is not None:
+            return {
+                "HF_HOME": self.hf_home_folder,
+                "TRITON_CACHE_DIR": self.hf_home_folder,
+            }
+        return None
+
+    def get_base_url_suffix(self) -> str:
+        """VLLM uses OpenAI-compatible API with /v1 suffix."""
+        return "/v1"
+
+
+class ConverterServer:
+    """Manages Docker server lifecycle with start/stop methods."""
+
+    def __init__(self, config: DockerServerConfig, auto_stop: bool = True):
+        self.config = config
+        self.auto_stop = auto_stop
+        self._server_context = None
+        self._container = None
+        self.base_url = None
+
+    def start(self):
+        """Start the Docker server."""
+        if self._server_context is not None:
+            logger.warning("Server already started")
+            return self.base_url
+
+        # Use the generic docker_server for all server types
+        self._server_context = docker_server(config=self.config, cleanup=self.auto_stop)
+
+        self.base_url, self._container = self._server_context.__enter__()
+        logger.info(f"Server started at {self.base_url}")
+        logger.info(f"Container ID: {self._container.id}")
+        logger.info(f"Container name: {self._container.name}")
+        return self.base_url, self._container
+
+    def stop(self):
+        """Stop the Docker server."""
+        if self._server_context is not None:
+            self._server_context.__exit__(None, None, None)
+            self._server_context = None
+            self._container = None
+            self.base_url = None
+            logger.info("Server stopped")
+
+    def __del__(self):
+        """Automatically stop server when object is destroyed if auto_stop is True."""
+        if self.auto_stop and self._server_context is not None:
+            self.stop()
+
+
+class DockerConfigRegistry:
+    """Registry for mapping model names to their Docker configurations."""
+
+    def __init__(self):
+        self._registry = dict()
+
+    def register(
+        self, model_name: str, config_factory: Callable[[], DockerServerConfig | None]
+    ):
+        """Register a config factory for a model name."""
+        self._registry[model_name] = config_factory
+
+    def get(self, model_name: str, default=False) -> DockerServerConfig | None:
+        """Get config for a model name. Returns default if not registered."""
+        if model_name not in self._registry:
+            if default:
+                return VLLMDockerServerConfig(model_name=model_name)
+            return None
+        return self._registry[model_name]()
+
+    def list_models(self) -> list[str]:
+        """List all registered model names."""
+        return list(self._registry.keys())
+
+
+# Global registry instance
+docker_config_registry = DockerConfigRegistry()

vlmparse/servers/utils.py

@@ -0,0 +1,250 @@
+import getpass
+import time
+from contextlib import contextmanager
+from pathlib import Path
+
+from loguru import logger
+
+import docker
+
+
+def _ensure_image_exists(
+    client: docker.DockerClient,
+    image: str,
+    dockerfile_path: Path,
+):
+    """Check if image exists, build it if not."""
+    try:
+        client.images.get(image)
+        logger.info(f"Docker image {image} found")
+        return
+    except docker.errors.ImageNotFound:
+        logger.info(f"Docker image {image} not found, building...")
+
+        if not dockerfile_path.exists():
+            raise FileNotFoundError(
+                f"Dockerfile directory not found at {dockerfile_path}"
+            ) from None
+
+        logger.info(f"Building image from {dockerfile_path}")
+
+        # Use low-level API for real-time streaming
+        api_client = docker.APIClient(base_url="unix://var/run/docker.sock")
+
+        # Build the image with streaming
+        build_stream = api_client.build(
+            path=str(dockerfile_path),
+            tag=image,
+            rm=True,
+            decode=True,  # Automatically decode JSON responses to dict
+        )
+
+        # Stream build logs in real-time
+        for chunk in build_stream:
+            if "stream" in chunk:
+                for line in chunk["stream"].splitlines():
+                    logger.info(line)
+            elif "error" in chunk:
+                logger.error(chunk["error"])
+                raise docker.errors.BuildError(chunk["error"], build_stream) from None
+            elif "status" in chunk:
+                # Handle status updates (e.g., downloading layers)
+                logger.debug(chunk["status"])
+
+        logger.info(f"Successfully built image {image}")
+
+
+@contextmanager
+def docker_server(
+    config: "DockerServerConfig",  # noqa: F821
+    timeout: int = 1000,
+    cleanup: bool = True,
+):
+    """Generic context manager for Docker server deployment.
+
+    Args:
+        config: DockerServerConfig (can be VLLMDockerServerConfig or GenericDockerServerConfig)
+        timeout: Timeout in seconds to wait for server to be ready
+        cleanup: If True, stop and remove container on exit. If False, leave container running
+
+    Yields:
+        tuple: (base_url, container) - The base URL of the server and the Docker container object
+    """
+
+    client = docker.from_env()
+    container = None
+
+    try:
+        # Ensure image exists
+        logger.info(f"Checking for Docker image {config.docker_image}...")
+
+        if config.dockerfile_dir is not None:
+            _ensure_image_exists(
+                client, config.docker_image, Path(config.dockerfile_dir)
+            )
+        else:
+            # Pull pre-built image
+            try:
+                client.images.get(config.docker_image)
+                logger.info(f"Docker image {config.docker_image} found locally")
+            except docker.errors.ImageNotFound:
+                logger.info(
+                    f"Docker image {config.docker_image} not found locally, pulling..."
+                )
+                client.images.pull(config.docker_image)
+                logger.info(f"Successfully pulled {config.docker_image}")
+
+        logger.info(
+            f"Starting Docker container for {config.model_name} on port {config.docker_port}"
+        )
+
+        # Configure GPU access
+        device_requests = None
+
+        if config.gpu_device_ids is None:
+            # Default: Try to use all GPUs if available
+            device_requests = [
+                docker.types.DeviceRequest(count=-1, capabilities=[["gpu"]])
+            ]
+        elif len(config.gpu_device_ids) > 0 and config.gpu_device_ids[0] != "":
+            # Use specific GPU devices
+            device_requests = [
+                docker.types.DeviceRequest(
+                    device_ids=config.gpu_device_ids, capabilities=[["gpu"]]
+                )
+            ]
+        else:
+            # Empty list means CPU-only, no GPU
+            device_requests = None
+
+        # Use generic methods from config
+        command = config.get_command()
+        volumes = config.get_volumes()
+        environment = config.get_environment()
+        container_port = config.container_port
+        log_prefix = config.model_name
+
+        # Construct URI for label
+        uri = f"http://localhost:{config.docker_port}{config.get_base_url_suffix()}"
+
+        # Determine GPU label
+        if config.gpu_device_ids is None:
+            gpu_label = "all"
+        elif len(config.gpu_device_ids) == 0 or (
+            len(config.gpu_device_ids) == 1 and config.gpu_device_ids[0] == ""
+        ):
+            gpu_label = "cpu"
+        else:
+            gpu_label = ",".join(config.gpu_device_ids)
+
+        # Start container
+        container_kwargs = {
+            "image": config.docker_image,
+            "ports": {f"{container_port}/tcp": config.docker_port},
+            "detach": True,
+            "remove": True,
+            "name": f"vlmparse-{config.model_name.replace('/', '-')}-{getpass.getuser()}",
+            "labels": {
+                "vlmparse_model_name": config.model_name,
+                "vlmparse_uri": uri,
+                "vlmparse_gpus": gpu_label,
+            },
+        }
+
+        if device_requests is not None:
+            container_kwargs["device_requests"] = device_requests
+        if command:
+            container_kwargs["command"] = command
+        if environment:
+            container_kwargs["environment"] = environment
+        if volumes:
+            container_kwargs["volumes"] = volumes
+        if config.entrypoint:
+            container_kwargs["entrypoint"] = config.entrypoint
+
+        container = client.containers.run(**container_kwargs)
+
+        logger.info(
+            f"Container {container.short_id} started, waiting for server to be ready..."
+        )
+
+        # Wait for server to be ready
+        start_time = time.time()
+        server_ready = False
+        last_log_position = 0
+
+        while time.time() - start_time < timeout:
+            try:
+                container.reload()
+            except docker.errors.NotFound as e:
+                logger.error("Container stopped unexpectedly during startup")
+                raise RuntimeError(
+                    "Container crashed during initialization. Check Docker logs for details."
+                ) from e
+
+            if container.status == "running":
+                # Get all logs and display new ones
+                all_logs = container.logs().decode("utf-8")
+
+                # Display new log lines
+                if len(all_logs) > last_log_position:
+                    new_logs = all_logs[last_log_position:]
+                    for line in new_logs.splitlines():
+                        if line.strip():  # Only print non-empty lines
+                            logger.info(f"[{log_prefix}] {line}")
+                    last_log_position = len(all_logs)
+
+                # Check if server is ready
+                for indicator in config.server_ready_indicators:
+                    if indicator in all_logs:
+                        server_ready = True
+                if server_ready:
+                    logger.info(f"Server ready indicator '{indicator}' found in logs")
+                    break
+
+            time.sleep(2)
+
+        if not server_ready:
+            raise TimeoutError(f"Server did not become ready within {timeout} seconds")
+
+        # Build base URL using config's suffix method
+        base_url = (
+            f"http://localhost:{config.docker_port}{config.get_base_url_suffix()}"
+        )
+
+        logger.info(f"{log_prefix} server ready at {base_url}")
+
+        yield base_url, container
+
+    finally:
+        if cleanup and container:
+            logger.info(f"Stopping container {container.short_id}")
+            container.stop(timeout=10)
+            logger.info("Container stopped")
+
+
+def get_model_from_uri(uri: str) -> str:
+    model = None
+    client = docker.from_env()
+    containers = client.containers.list()
+    for container in containers:
+        c_uri = container.labels.get("vlmparse_uri")
+        c_model = container.labels.get("vlmparse_model_name")
+
+        # Check if user URI matches container URI (ignoring /v1 suffix if missing)
+        if c_uri and (
+            c_uri == uri or c_uri.startswith(uri.rstrip("/")) or uri.startswith(c_uri)
+        ):
+            # Update URI to the correct one from container (likely has /v1)
+            if len(c_uri) > len(uri.rstrip("/")):
+                logger.info(f"Updating URI from {uri} to {c_uri}")
+                uri = c_uri
+
+            # Infer model if not provided
+            if model is None and c_model:
+                logger.info(f"Inferred model {c_model} from container")
+                model = c_model
+            break
+    if model is None:
+        raise ValueError(f"No model found for URI {uri}")
+    return model

vlmparse/st_viewer/fs_nav.py

@@ -0,0 +1,53 @@
+from __future__ import annotations
+
+import glob
+import os
+from typing import List, Optional
+
+import streamlit as st
+
+
+def get_gz_files_count(folder_path: str) -> int:
+    return len(glob.glob(os.path.join(folder_path, "*.json*")))
+
+
+def get_subdirectories(path: str) -> List[str]:
+    return sorted([d for d in os.listdir(path) if os.path.isdir(os.path.join(path, d))])
+
+
+def file_selector(root_folder: str) -> Optional[str]:
+    st.title("Folder Navigation")
+    if not root_folder or not os.path.exists(root_folder):
+        return None
+
+    current_path = selected_path = root_folder
+    level = 0
+    while True:
+        subdirs = get_subdirectories(current_path)
+        if not subdirs:
+            break
+        dir_options = [
+            f"{d} ({get_gz_files_count(os.path.join(current_path, d))} .zip files)"
+            for d in subdirs
+        ]
+        selected = st.selectbox(
+            f"Level {level} Selection",
+            ["--Select--"] + dir_options,
+            key=f"level_{level}",
+        )
+        if selected == "--Select--" or not selected:
+            break
+        selected_dir = selected.split(" (", 1)[0]
+        current_path = os.path.join(current_path, selected_dir)
+        selected_path = current_path
+        level += 1
+
+    gz_files = sorted(glob.glob(os.path.join(selected_path, "*.zip")))
+    if gz_files:
+        selected_file = st.selectbox(
+            "Select .zip file",
+            ["--Select--"] + [os.path.basename(f) for f in gz_files],
+        )
+        if selected_file != "--Select--":
+            return os.path.join(selected_path, selected_file)
+    return None