multixtract 0.1.1__py3-none-any.whl

multixtract/pipeline.py

@@ -0,0 +1,250 @@
+"""Pipeline orchestrator.
+
+Wires the vendor-neutral core (extraction -> filtering -> chunking) together
+with *injected* providers (vision, embedder, store). The orchestrator never
+imports a vendor SDK — it only talks to the interfaces.
+
+Providers are optional:
+  * no ``vision``   -> images are kept but not described
+  * no ``embedder`` -> chunks are produced without embeddings
+  * no ``store``    -> nothing is persisted (results returned in memory)
+"""
+from __future__ import annotations
+
+import logging
+import os
+from concurrent.futures import ThreadPoolExecutor, as_completed
+from dataclasses import dataclass, field
+from datetime import datetime, timezone
+from typing import Any, Dict, List, Optional
+
+from .chunking import chunk_document
+from .extraction import extract_document
+from .filters import ImageFilterPipeline
+from .interfaces import BlobStore, Embedder, PipelineConfig, VisionModel
+
+log = logging.getLogger("multixtract")
+
+
+@dataclass
+class ExtractionResult:
+    """Everything the pipeline produced for one document."""
+
+    base_name: str
+    document: Dict[str, Any]
+    chunks: List[Dict[str, Any]] = field(default_factory=list)
+    image_index: List[Dict[str, Any]] = field(default_factory=list)
+    filter_stats: Dict[str, int] = field(default_factory=dict)
+
+
+class Pipeline:
+    """End-to-end document extraction pipeline.
+
+    Example::
+
+        pipeline = Pipeline(vision=my_vision, embedder=my_embedder, store=my_store)
+        result = pipeline.process("report.pdf")
+    """
+
+    def __init__(
+        self,
+        vision: Optional[VisionModel] = None,
+        embedder: Optional[Embedder] = None,
+        store: Optional[BlobStore] = None,
+        config: Optional[PipelineConfig] = None,
+    ) -> None:
+        self.vision = vision
+        self.embedder = embedder
+        self.store = store
+        self.config = config or PipelineConfig()
+
+    # ------------------------------------------------------------------
+
+    def process(self, doc_path: str, skip_if_exists: bool = True) -> ExtractionResult:
+        """Run the full pipeline on a single document.
+
+        Raises:
+            ValueError: If no extractor is registered for the file's extension.
+            ImportError: If the required optional extra for this format is not installed.
+        """
+        base_name = os.path.splitext(os.path.basename(doc_path))[0]
+        config = self.config
+
+        # Resume support: skip documents already stored.
+        if skip_if_exists and self.store is not None:
+            doc_key = f"{config.doc_json_subdir}/{base_name}.json"
+            if self.store.exists(doc_key):
+                log.info("Skipping %s (output exists)", base_name)
+                return ExtractionResult(base_name=base_name, document={})
+
+        # Phase 1 — extract + filter
+        _filter = ImageFilterPipeline(
+            min_image_size=config.min_image_size,
+            min_image_size_minor=config.min_image_size_minor,
+            reference_img_dir=config.reference_img_dir,
+        )
+        document, prepared = extract_document(doc_path, image_filter=_filter)
+
+        # Persist raw image bytes before vision frees them.
+        if self.store is not None:
+            self._persist_images(base_name, prepared)
+
+        # Phase 2a — vision (parallel)
+        vision_by_id = self._run_vision(prepared)
+
+        # Phase 2b — embed image descriptions
+        image_embeds = self._embed_images(prepared, vision_by_id)
+
+        # Assemble image metadata onto pages + a flat image index
+        image_index = self._assemble_images(document, prepared, vision_by_id, image_embeds)
+
+        # Phase 3 — chunk + embed chunks (wired to config)
+        chunks = chunk_document(
+            document,
+            base_name,
+            image_embeddings=image_embeds,
+            target_tokens=config.chunk_target_tokens,
+            overlap_tokens=config.chunk_overlap_tokens,
+        )
+        self._embed_chunks(chunks)
+
+        stamp = {
+            "doc_id":     base_name,
+            "file_name":  os.path.basename(doc_path),
+            "file_path":  doc_path,
+            "file_type":  os.path.splitext(doc_path)[1].lstrip(".").lower(),
+            "total_pgs":  len(document.get("pgs", [])),
+            "last_updated": datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ"),
+        }
+        for chunk in chunks:
+            chunk.update(stamp)
+
+        result = ExtractionResult(
+            base_name=base_name,
+            document=document,
+            chunks=chunks,
+            image_index=image_index,
+            filter_stats=_filter.filter_stats,
+        )
+
+        if self.store is not None:
+            self._persist(result)
+        return result
+
+    # ------------------------------------------------------------------
+
+    def _run_vision(self, prepared: List[Dict[str, Any]]) -> Dict[str, Any]:
+        if self.vision is None or not prepared:
+            return {}
+        results: Dict[str, Any] = {}
+        workers = min(self.config.vision_workers, len(prepared))
+        with ThreadPoolExecutor(max_workers=workers) as pool:
+            futures = {
+                pool.submit(
+                    self.vision.analyze,
+                    img["image_bytes"], img["ext"], img["width"], img["height"],
+                ): img["image_id"]
+                for img in prepared
+            }
+            for fut in as_completed(futures):
+                image_id = futures[fut]
+                try:
+                    results[image_id] = fut.result()
+                except Exception as exc:  # provider should not raise, but be safe
+                    log.warning("vision failed for %s: %s", image_id, exc)
+        # Free image bytes once vision is done.
+        for img in prepared:
+            img.pop("image_bytes", None)
+        return results
+
+    def _embed_images(self, prepared, vision_by_id) -> Dict[str, List[float]]:
+        if self.embedder is None:
+            return {}
+        pairs = [
+            (img["image_id"], vision_by_id[img["image_id"]].best_text())
+            for img in prepared
+            if vision_by_id.get(img["image_id"]) and vision_by_id[img["image_id"]].best_text()
+        ]
+        if not pairs:
+            return {}
+        ids, texts = zip(*pairs)
+        all_vectors = self.embedder.embed(list(texts))
+        return {i: v for i, v in zip(ids, all_vectors) if v is not None}
+
+    def _assemble_images(self, document, prepared, vision_by_id, image_embeds) -> List[Dict[str, Any]]:
+        image_index: List[Dict[str, Any]] = []
+        pages_by_num = {page["pg_num"]: page for page in document.get("pgs", [])}
+        for img in prepared:
+            vision_result = vision_by_id.get(img["image_id"])
+            page_meta = {
+                "img_id": img["image_id"],
+                "img_idx": img["img_idx"],
+                "img_path": img["img_path"],
+                "ocr_text": vision_result.ocr_text if vision_result else "",
+                "caption": vision_result.caption if vision_result else "",
+                "description": vision_result.description if vision_result else "",
+            }
+            # Safe page lookup
+            page = pages_by_num.get(img["page_number"])
+            if page is not None:
+                page["imgs"].append(page_meta)
+            image_index.append({
+                **page_meta,
+                "pg_num": img["page_number"],
+                "width": img["width"],
+                "height": img["height"],
+                "format": img["ext"],
+                "embedding": image_embeds.get(img["image_id"]),
+            })
+        return image_index
+
+    def _persist_images(self, base_name: str, prepared: List[Dict[str, Any]]) -> None:
+        config = self.config
+        ext_to_mime = {"png": "image/png", "jpg": "image/jpeg", "jpeg": "image/jpeg",
+                       "gif": "image/gif", "bmp": "image/bmp",
+                       "tiff": "image/tiff", "tif": "image/tiff",
+                       "webp": "image/webp"}
+        for img in prepared:
+            blob_path = f"{config.images_subdir}/{base_name}/{img['img_path']}"
+            mime = ext_to_mime.get(img["ext"].lower(), "application/octet-stream")
+            locator = self.store.put_bytes(blob_path, img["image_bytes"], content_type=mime)
+            img["img_path"] = locator
+
+    def _embed_chunks(self, chunks: List[Dict[str, Any]]) -> None:
+        if self.embedder is None or not chunks:
+            return
+        pending = [(chunk_index, chunk) for chunk_index, chunk in enumerate(chunks) if chunk["embedding"] is None]
+        if not pending:
+            return
+
+        texts = [chunk["content"][: self.config.embed_text_limit] for _, chunk in pending]
+        all_vectors = self.embedder.embed(texts)
+        if len(all_vectors) != len(pending):
+            log.warning("embedder returned %d vectors for %d chunks; some embeddings will be None",
+                        len(all_vectors), len(pending))
+        for (chunk_index, _), vec in zip(pending, all_vectors):
+            chunks[chunk_index]["embedding"] = vec
+
+    def _persist(self, result: ExtractionResult) -> None:
+        config = self.config
+        base_name = result.base_name
+        # Write image and chunk blobs first; doc JSON is written last and acts
+        # as the completion marker checked by skip_if_exists. A partial write
+        # (e.g. network error mid-way) leaves the doc key absent so retries
+        # re-process the document instead of skipping it permanently.
+        self.store.put_json(
+            f"{config.image_json_subdir}/{base_name}_image.json",
+            {"imgs": result.image_index},
+        )
+        self.store.put_json(
+            f"{config.chunks_subdir}/{base_name}_chunks.json",
+            {
+                "_header": {
+                    "file_name": base_name,
+                    "total_pgs": len(result.document.get("pgs", [])),
+                },
+                "chunks": result.chunks,
+            },
+            compact=True,
+        )
+        self.store.put_json(f"{config.doc_json_subdir}/{base_name}.json", result.document)

multixtract/providers/__init__.py

@@ -0,0 +1,38 @@
+"""Concrete provider implementations.
+
+Importing a provider class does not pull in its SDK until you instantiate it.
+Each submodule is imported lazily here so a broken or missing optional
+dependency in one provider never prevents other providers from loading.
+"""
+
+def __getattr__(name: str):
+    _map = {
+        "OpenAIVisionModel":       (".openai",    "OpenAIVisionModel"),
+        "OpenAIEmbedder":          (".openai",    "OpenAIEmbedder"),
+        "AzureOpenAIVisionModel":  (".azure",     "AzureOpenAIVisionModel"),
+        "AzureOpenAIEmbedder":     (".azure",     "AzureOpenAIEmbedder"),
+        "AzureBlobStore":          (".storage",   "AzureBlobStore"),
+        "LocalDiskStore":          (".storage",   "LocalDiskStore"),
+        "Llama32VisionModel":      (".llama",     "Llama32VisionModel"),
+        "Qwen2VLVisionModel":      (".qwen2vl",   "Qwen2VLVisionModel"),
+        "SmolVLMVisionModel":      (".smolvlm",   "SmolVLMVisionModel"),
+    }
+    if name in _map:
+        module_rel, attr = _map[name]
+        import importlib
+        mod = importlib.import_module(module_rel, package=__name__)
+        return getattr(mod, attr)
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
+
+
+__all__ = [
+    "OpenAIVisionModel",
+    "OpenAIEmbedder",
+    "AzureOpenAIVisionModel",
+    "AzureOpenAIEmbedder",
+    "LocalDiskStore",
+    "AzureBlobStore",
+    "Qwen2VLVisionModel",
+    "SmolVLMVisionModel",
+    "Llama32VisionModel",
+]

multixtract/providers/_utils.py

@@ -0,0 +1,32 @@
+"""Shared helpers for provider implementations."""
+from __future__ import annotations
+
+import contextlib
+import io
+
+
+def _infer_device() -> str:
+    """Return 'cuda' if a CUDA GPU is available, else 'cpu'."""
+    try:
+        import torch
+        return "cuda" if torch.cuda.is_available() else "cpu"
+    except Exception:
+        return "cpu"
+
+
+@contextlib.contextmanager
+def _open_image_rgb(image_bytes: bytes):
+    """Open raw image bytes and yield a converted RGB PIL Image.
+
+    Closes both the raw and the converted image on exit, even on error.
+    """
+    from PIL import Image
+    raw = Image.open(io.BytesIO(image_bytes))
+    try:
+        image = raw.convert("RGB")
+    finally:
+        raw.close()
+    try:
+        yield image
+    finally:
+        image.close()

multixtract/providers/azure.py

@@ -0,0 +1,91 @@
+"""Azure OpenAI providers (optional extra: ``pip install multixtract[azure]``).
+
+These reuse the OpenAI provider logic but construct an ``AzureOpenAI`` client.
+Pass credentials explicitly — never hard-code secrets. Source them from
+environment variables, a secrets manager (Azure Key Vault, AWS Secrets Manager,
+HashiCorp Vault, etc.), or your platform's secret store.
+
+Both classes support passwordless / managed-identity auth via
+``azure_ad_token_provider`` — a callable that returns a fresh bearer token.
+When provided, ``api_key`` should be omitted (or ``None``).
+
+Example (managed identity)::
+
+    from azure.identity import DefaultAzureCredential, get_bearer_token_provider
+    credential = DefaultAzureCredential()
+    token_provider = get_bearer_token_provider(
+        credential, "https://cognitiveservices.azure.com/.default"
+    )
+    vision = AzureOpenAIVisionModel(
+        endpoint="https://<resource>.openai.azure.com/",
+        azure_ad_token_provider=token_provider,
+    )
+"""
+from __future__ import annotations
+
+from typing import Callable, Optional
+
+from .openai import OpenAIEmbedder, OpenAIVisionModel
+
+
+def _azure_client(
+    endpoint: str,
+    api_key: Optional[str],
+    api_version: str,
+    azure_ad_token_provider: Optional[Callable[[], str]] = None,
+):
+    from openai import AzureOpenAI  # keeps openai optional for core users
+    return AzureOpenAI(
+        azure_endpoint=endpoint,
+        api_key=api_key,
+        api_version=api_version,
+        azure_ad_token_provider=azure_ad_token_provider,
+    )
+
+
+class AzureOpenAIVisionModel(OpenAIVisionModel):
+    """Vision provider backed by an Azure OpenAI deployment (e.g. gpt-4o).
+
+    Supports both API-key and passwordless (Azure AD / managed identity) auth.
+    Pass ``azure_ad_token_provider`` for passwordless; omit ``api_key`` in that case.
+    """
+
+    def __init__(
+        self,
+        endpoint: str,
+        api_key: Optional[str] = None,
+        deployment: str = "gpt-4o",
+        api_version: str = "2024-10-21",
+        azure_ad_token_provider: Optional[Callable[[], str]] = None,
+        **kwargs,
+    ) -> None:
+        super().__init__(
+            model=deployment,
+            client=_azure_client(endpoint, api_key, api_version, azure_ad_token_provider),
+            **kwargs,
+        )
+
+
+class AzureOpenAIEmbedder(OpenAIEmbedder):
+    """Embedding provider backed by an Azure OpenAI deployment.
+
+    Supports both API-key and passwordless (Azure AD / managed identity) auth.
+    Pass ``azure_ad_token_provider`` for passwordless; omit ``api_key`` in that case.
+    """
+
+    def __init__(
+        self,
+        endpoint: str,
+        api_key: Optional[str] = None,
+        deployment: str = "text-embedding-3-large",
+        api_version: str = "2024-10-21",
+        dim: int = 1024,
+        azure_ad_token_provider: Optional[Callable[[], str]] = None,
+        **kwargs,
+    ) -> None:
+        super().__init__(
+            model=deployment,
+            dim=dim,
+            client=_azure_client(endpoint, api_key, api_version, azure_ad_token_provider),
+            **kwargs,
+        )

multixtract/providers/llama.py

@@ -0,0 +1,151 @@
+"""Llama 3.2 Vision provider (optional extra: ``pip install multixtract[llama]``).
+
+Runs Llama 3.2 Vision **fully offline** via HuggingFace transformers —
+no cloud, no API key. It produces the same GPT-4o-style
+caption / OCR / description output as the cloud providers by reusing the shared
+:data:`DEFAULT_SYSTEM_PROMPT` and :func:`parse_vision_response`, so it is a
+drop-in ``VisionModel`` for the pipeline or the standalone recipes.
+
+A CUDA GPU is strongly recommended — the 11B model requires ≥16 GB VRAM. The
+first call downloads the model weights from the HuggingFace hub; set
+``HF_HOME`` to a persistent path to cache them.
+
+Example::
+
+    from multixtract import extract_document
+    from multixtract.providers import Llama32VisionModel
+
+    vision = Llama32VisionModel()                                    # 11B default
+    vision = Llama32VisionModel("meta-llama/Llama-3.2-11B-Vision-Instruct")
+    _, images = extract_document("report.pdf")
+    for img in images:
+        r = vision.analyze(img["image_bytes"], img["ext"])
+        print(r.caption, "|", r.description)
+"""
+from __future__ import annotations
+
+import logging
+import threading
+from typing import Optional
+
+from ..interfaces import VisionModel, VisionResult
+from ..vision import DEFAULT_SYSTEM_PROMPT, parse_vision_response
+from ._utils import _infer_device, _open_image_rgb
+
+log = logging.getLogger("multixtract")
+
+
+class Llama32VisionModel(VisionModel):
+    """Offline ``VisionModel`` backed by Llama 3.2 Vision (transformers).
+
+    The heavy deps (``torch`` / ``transformers`` / ``accelerate``) are imported
+    lazily, so simply importing this module never pulls them in. The model is
+    loaded once, on construction, and reused across :meth:`analyze` calls.
+    """
+
+    def __init__(
+        self,
+        model_id: str = "meta-llama/Llama-3.2-11B-Vision-Instruct",
+        system_prompt: str = DEFAULT_SYSTEM_PROMPT,
+        max_new_tokens: int = 512,
+        device: Optional[str] = None,
+        torch_dtype: str = "bfloat16",
+        load_in_4bit: bool = False,
+        model=None,
+        processor=None,
+    ) -> None:
+        """Load a Llama 3.2 Vision model.
+
+        Args:
+            model_id: HuggingFace model id.
+                ``meta-llama/Llama-3.2-11B-Vision-Instruct`` (default, best accuracy),
+                ``meta-llama/Llama-3.2-90B-Vision-Instruct`` (highest accuracy, needs 80+ GB VRAM).
+            system_prompt: Instruction sent with every image. Defaults to the
+                shared, parseable CAPTION/OCR_TEXT/DESCRIPTION prompt.
+            max_new_tokens: Generation budget per image.
+            device: ``"cuda"`` or ``"cpu"``. Auto-detected when omitted.
+            torch_dtype: dtype name used on GPU (``"bfloat16"`` recommended for Llama 3.2).
+            load_in_4bit: 4-bit quantised load (needs ``bitsandbytes``); halves
+                VRAM for the 11B model.
+            model / processor: Pre-built instances to reuse (skips loading).
+        """
+        self.model_id = model_id
+        self.system_prompt = system_prompt
+        self.max_new_tokens = max_new_tokens
+
+        self._lock = threading.Lock()
+
+        if model is not None and processor is not None:
+            self._model, self._processor = model, processor
+            self.device = device or _infer_device()
+            return
+
+        try:
+            import torch
+            from transformers import AutoProcessor, MllamaForConditionalGeneration
+        except ImportError as e:
+            raise ImportError(
+                "Llama 3.2 Vision support requires transformers + torch + accelerate: "
+                "pip install 'multixtract[llama]'"
+            ) from e
+
+        self.device = device or _infer_device()
+        load_kwargs = {"low_cpu_mem_usage": True}
+        if self.device == "cuda":
+            load_kwargs["torch_dtype"] = getattr(torch, torch_dtype, torch.bfloat16)
+        else:
+            load_kwargs["torch_dtype"] = torch.float32
+        if load_in_4bit:
+            load_kwargs["load_in_4bit"] = True
+            load_kwargs["device_map"] = "auto"  # required by bitsandbytes 4-bit
+        else:
+            load_kwargs["device_map"] = self.device
+
+        self._processor = AutoProcessor.from_pretrained(model_id)
+        self._model = MllamaForConditionalGeneration.from_pretrained(model_id, **load_kwargs)
+        self._model.eval()
+
+    def analyze(
+        self,
+        image_bytes: bytes,
+        ext: str = "png",
+        width: int = 0,
+        height: int = 0,
+    ) -> VisionResult:
+        """Describe one image. Never raises — returns an empty result on failure."""
+        try:
+            import torch
+
+            with _open_image_rgb(image_bytes) as image:
+                messages = [
+                    {
+                        "role": "user",
+                        "content": [
+                            {"type": "image"},
+                            {"type": "text", "text": self.system_prompt},
+                        ],
+                    }
+                ]
+                with self._lock:
+                    input_text = self._processor.apply_chat_template(
+                        messages, add_generation_prompt=True
+                    )
+                    inputs = self._processor(
+                        image, input_text, return_tensors="pt"
+                    ).to(self._model.device)
+                    with torch.inference_mode():
+                        output_ids = self._model.generate(
+                            **inputs,
+                            max_new_tokens=self.max_new_tokens,
+                            do_sample=False,
+                        )
+                    # Decode only the newly generated tokens (skip the prompt).
+                    prompt_len = inputs["input_ids"].shape[-1]
+                    text = self._processor.decode(
+                        output_ids[0][prompt_len:], skip_special_tokens=True
+                    )
+
+            return parse_vision_response(text)
+        except Exception:  # noqa: BLE001 — never break the caller/pipeline
+            log.warning("Llama32Vision analyze failed", exc_info=True)
+            return VisionResult()

multixtract/providers/openai.py

@@ -0,0 +1,125 @@
+"""OpenAI providers (optional extra: ``pip install multixtract[openai]``).
+
+Implements VisionModel and Embedder against the OpenAI Python SDK. The core
+pipeline depends only on the interfaces — importing this module is the only
+place ``openai`` is required.
+"""
+from __future__ import annotations
+
+import logging
+import time
+from typing import List, Optional
+
+from ..interfaces import Embedder, VisionModel, VisionResult
+from ..vision import DEFAULT_SYSTEM_PROMPT, parse_vision_response, to_data_url
+
+log = logging.getLogger("multixtract")
+
+
+def _is_permanent(exc: Exception) -> bool:
+    """Return True for errors that will never succeed on retry."""
+    try:
+        from openai import AuthenticationError, BadRequestError, PermissionDeniedError
+        if isinstance(exc, (AuthenticationError, BadRequestError, PermissionDeniedError)):
+            return True
+    except ImportError:
+        pass
+    return False
+
+
+def _retry(func, *, label: str = "API", max_retries: int = 3):
+    """Call *func()* with exponential backoff and rate-limit awareness."""
+    for attempt in range(max_retries):
+        try:
+            return func()
+        except Exception as exc:  # noqa: BLE001
+            if attempt >= max_retries - 1 or _is_permanent(exc):
+                log.warning("%s failed after %d attempts: %s", label, attempt + 1, exc)  # noqa: BLE001
+                raise
+            wait = 2 ** attempt
+            time.sleep(wait)
+
+
+class OpenAIVisionModel(VisionModel):
+    """Vision provider backed by an OpenAI chat-completions vision model."""
+
+    def __init__(
+        self,
+        api_key: Optional[str] = None,
+        model: str = "gpt-4o",
+        system_prompt: str = DEFAULT_SYSTEM_PROMPT,
+        max_tokens: int = 800,
+        temperature: float = 0.1,
+        client=None,
+    ) -> None:
+        if client is None:
+            from openai import OpenAI  # local import keeps openai optional
+            client = OpenAI(api_key=api_key)
+        self._client = client
+        self.model = model
+        self.system_prompt = system_prompt
+        self.max_tokens = max_tokens
+        self.temperature = temperature
+
+    def analyze(self, image_bytes: bytes, ext: str = "png", width: int = 0, height: int = 0) -> VisionResult:
+        try:
+            data_url = to_data_url(image_bytes, ext, width, height)
+            resp = _retry(
+                lambda: self._client.chat.completions.create(
+                    model=self.model,
+                    messages=[
+                        {"role": "system", "content": self.system_prompt},
+                        {"role": "user", "content": [
+                            {"type": "text", "text": "Analyze this image."},
+                            {"type": "image_url", "image_url": {"url": data_url, "detail": "high"}},
+                        ]},
+                    ],
+                    max_tokens=self.max_tokens,
+                    temperature=self.temperature,
+                ),
+                label="vision",
+            )
+            return parse_vision_response(resp.choices[0].message.content or "")
+        except Exception:  # noqa: BLE001 — never break the pipeline
+            return VisionResult()
+
+
+class OpenAIEmbedder(Embedder):
+    """Embedding provider backed by an OpenAI embeddings model."""
+
+    def __init__(
+        self,
+        api_key: Optional[str] = None,
+        model: str = "text-embedding-3-large",
+        dim: int = 1024,
+        batch_size: int = 16,
+        text_limit: int = 8000,
+        client=None,
+    ) -> None:
+        if client is None:
+            from openai import OpenAI
+            client = OpenAI(api_key=api_key)
+        self._client = client
+        self.model = model
+        self.dim = dim
+        self.batch_size = batch_size
+        self.text_limit = text_limit
+
+    def embed(self, texts: List[str]) -> List[Optional[List[float]]]:
+        results: List[Optional[List[float]]] = [None] * len(texts)
+        work = [(original_index, text[: self.text_limit]) for original_index, text in enumerate(texts) if text]
+        for start in range(0, len(work), self.batch_size):
+            batch = work[start : start + self.batch_size]
+            inputs = [text for _, text in batch]
+            try:
+                resp = _retry(
+                    lambda inp=inputs: self._client.embeddings.create(
+                        model=self.model, input=inp, dimensions=self.dim,
+                    ),
+                    label="embed",
+                )
+                for (original_index, _), item in zip(batch, resp.data):
+                    results[original_index] = item.embedding
+            except Exception:  # noqa: BLE001
+                pass  # leave those entries as None
+        return results