refactor: enable pass model adapter

.vscode/settings.json

@@ -4,7 +4,8 @@
         "editor.codeActionsOnSave": {
             "source.organizeImports": "explicit"
         },
-        "files.trimTrailingWhitespace": true
+        "python.linting.enabled": true,
+        "files.trimTrailingWhitespace": true,
     },
     "cSpell.words": [
         "deepseek",

doc_page_extractor/__init__.py

@@ -1,21 +1,28 @@
 from .extraction_context import (
     AbortError,
     ExtractionAbortedError,
-    ExtractionContext,
     TokenLimitError,
 )
-from .extractor import Layout, PageExtractor
-from .model import DeepSeekOCRSize
+from .extractor import create_page_extractor
 from .plot import plot
+from .types import (
+    Layout,
+    PageExtractor,
+    DeepSeekOCRModel,
+    ExtractionContext,
+    DeepSeekOCRSize,
+)
 
 __version__ = "1.0.0"
 __all__ = [
+    "plot",
+    "create_page_extractor",
+    "PageExtractor",
     "DeepSeekOCRSize",
+    "DeepSeekOCRModel",
     "ExtractionContext",
     "AbortError",
     "ExtractionAbortedError",
     "TokenLimitError",
     "Layout",
-    "PageExtractor",
-    "plot",
 ]

doc_page_extractor/extraction_context.py

@@ -1,17 +1,11 @@
-from dataclasses import dataclass
 from typing import Any, Callable, cast
 
 import torch
 from transformers import StoppingCriteria
 
+from .types import ExtractionContext
+
 
-@dataclass
-class ExtractionContext:
-    check_aborted: Callable[[], bool]
-    max_tokens: int | None = None
-    max_output_tokens: int | None = None
-    input_tokens: int = 0
-    output_tokens: int = 0
 
 
 class ExtractionAbortedError(Exception):

doc_page_extractor/extractor.py

@@ -1,35 +1,37 @@
 import tempfile
-from dataclasses import dataclass
 from os import PathLike
 from pathlib import Path
 from typing import Generator, cast
 
 from PIL import Image
 
 from .check_env import check_env
-from .extraction_context import ExtractionContext
-from .model import DeepSeekOCRModel, DeepSeekOCRSize
+from .model import DeepSeekOCRHugginfaceModel
 from .parser import ParsedItemKind, parse_ocr_response
 from .redacter import background_color, redact
+from .types import Layout, PageExtractor, ExtractionContext, DeepSeekOCRModel, DeepSeekOCRSize
 
 
-@dataclass
-class Layout:
-    ref: str
-    det: tuple[int, int, int, int]
-    text: str | None
 
+def create_page_extractor(
+    model_path: PathLike | None = None,
+    local_only: bool = False,
+) -> PageExtractor:
+    model: DeepSeekOCRHugginfaceModel = DeepSeekOCRHugginfaceModel(
+        model_path=Path(model_path) if model_path else None,
+        local_only=local_only,
+    )
+    return _PageExtractorImpls(model)
 
-class PageExtractor:
-    def __init__(
-        self,
-        model_path: PathLike | None = None,
-        local_only: bool = False,
-    ) -> None:
-        self._model: DeepSeekOCRModel = DeepSeekOCRModel(
-            model_path=Path(model_path) if model_path else None,
-            local_only=local_only,
-        )
+def create_page_extractor_with_model(model: DeepSeekOCRModel) -> PageExtractor:
+    if not isinstance(model, DeepSeekOCRModel):
+        raise TypeError("model must implement DeepSeekOCRModel protocol")
+    return _PageExtractorImpls(model)
+
+
+class _PageExtractorImpls:
+    def __init__(self, model: DeepSeekOCRModel) -> None:
+        self._model: DeepSeekOCRModel = model
 
     def download_models(self) -> None:
         self._model.download()

doc_page_extractor/injection.py

@@ -59,7 +59,8 @@
 
 from transformers import StoppingCriteria
 
-from .extraction_context import AbortStoppingCriteria, ExtractionContext
+from .types import ExtractionContext
+from .extraction_context import AbortStoppingCriteria
 
 _LOCAL = threading.local()
 _LOCAL_KEY = "value"

doc_page_extractor/model.py

@@ -2,7 +2,7 @@
 from dataclasses import dataclass
 from importlib.util import find_spec
 from pathlib import Path
-from typing import Any, Literal
+from typing import Any
 
 import torch
 from huggingface_hub import snapshot_download
@@ -12,8 +12,8 @@
 
 from .extraction_context import ExtractionContext
 from .injection import InferWithInterruption, preprocess_model
+from .types import DeepSeekOCRSize
 
-DeepSeekOCRSize = Literal["tiny", "small", "base", "large", "gundam"]
 
 
 @dataclass
@@ -40,7 +40,7 @@ class _SizeConfig:
 _Models = tuple[Any, Any]
 
 
-class DeepSeekOCRModel:
+class DeepSeekOCRHugginfaceModel:
     def __init__(self, model_path: Path | None, local_only: bool) -> None:
         if local_only and model_path is None:
             raise ValueError("model_path must be provided when local_only is True")

doc_page_extractor/plot.py

@@ -4,7 +4,7 @@
 from PIL.Image import Image
 from PIL.ImageFont import FreeTypeFont, load_default
 
-from .extractor import Layout
+from .types import Layout
 
 _FRAGMENT_COLOR = (0x49, 0xCF, 0xCB)  # Light Green
 _Color = tuple[int, int, int]

doc_page_extractor/types.py

@@ -0,0 +1,62 @@
+from dataclasses import dataclass
+from typing import Protocol, runtime_checkable
+from typing import Generator, Literal, Callable
+
+from PIL import Image
+
+
+DeepSeekOCRSize = Literal["tiny", "small", "base", "large", "gundam"]
+
+@dataclass
+class Layout:
+    ref: str
+    det: tuple[int, int, int, int]
+    text: str | None
+
+@dataclass
+class ExtractionContext:
+    check_aborted: Callable[[], bool]
+    max_tokens: int | None = None
+    max_output_tokens: int | None = None
+    input_tokens: int = 0
+    output_tokens: int = 0
+
+
+@runtime_checkable
+class PageExtractor(Protocol):
+    def download_models(self) -> None:
+        ...
+
+    def load_models(self) -> None:
+        ...
+
+    def extract(
+        self,
+        image: Image.Image,
+        size: DeepSeekOCRSize,
+        stages: int = 1,
+        context: ExtractionContext | None = None,
+    ) -> Generator[tuple[Image.Image, list[Layout]], None, None]:
+        ...
+
+
+@runtime_checkable
+class DeepSeekOCRModel(Protocol):
+    def download(self) -> None:
+        ...
+
+    def load(self) -> None:
+        ...
+
+    def unload(self) -> None:
+        ...
+
+    def generate(
+        self,
+        image: Image.Image,
+        prompt: str,
+        temp_path: str,
+        size: DeepSeekOCRSize,
+        context: ExtractionContext | None,
+    ) -> str:
+        ...

main.py

@@ -3,7 +3,7 @@
 
 from PIL import Image
 
-from doc_page_extractor import ExtractionContext, PageExtractor, plot
+from doc_page_extractor import plot, create_page_extractor, ExtractionContext
 
 _ABORT_TIMEOUT = 9999.0  # seconds
 
@@ -12,7 +12,7 @@ def main() -> None:
     project_root = Path(__file__).parent
     image_dir_path = project_root / "tests" / "images"
     image_name = "double_column.png"
-    extractor = PageExtractor(
+    extractor = create_page_extractor(
         model_path=project_root / "models-cache",
         local_only=False,
     )

pyproject.toml

@@ -48,11 +48,6 @@ name = "pytorch-cu121"
 url = "https://download.pytorch.org/whl/cu121"
 priority = "explicit"
 
-[[tool.poetry.source]]
-name = "pytorch-cu118"
-url = "https://download.pytorch.org/whl/cu118"
-priority = "explicit"
-
 [tool.poetry.group.dev.dependencies]
 pylint = "^3.3.7"
 # Development environment includes CUDA 12.1 PyTorch