base script with connection to pixtral to extract pages

Author: christianabbet

config/pixtral_config.yml

@@ -14,10 +14,10 @@ prompt_version: v1.2 # current best is v1.2
 max_document_size_mb: 4.5
 slack_size_mb: 0.2
 # conversation generation
-max_tokens: 5
+max_tokens: 100
 temperature: 0.2
 # throttling
 qps: 3.0
 max_retries: 6
 backoff_base: 0.4
-backoff_cap: 8.0
+backoff_cap: 8.0

prompts/title_extraction_prompts.yml

@@ -0,0 +1,15 @@
+v0:
+  system_prompt: |
+    You are an expert in document layout analysis and content structure recognition.
+    Your task is to identify and extract the main title of the provided document.
+
+    # Guidelines for identifying the title:
+
+    - The title is typically the most prominent text element — largest font, bold, or positioned at the center or top of the document
+    - It represents the primary subject or name of the document, not an author name, or company name
+    - If multiple candidate titles exist, select the one that best represents the document's overall topic
+
+    # Output instructions:
+
+    Return the title text only — no explanations, punctuation wrappers, quotes, or extra formatting
+    Preserve the original capitalization of the title

src/classifiers/pixtral_classifier.py

@@ -7,6 +7,7 @@
 import boto3
 import pymupdf
 from botocore.exceptions import ClientError
+from pydantic import BaseModel, Field
 
 from src.classifiers.classifier_types import Classifier, ClassifierTypes
 from src.classifiers.utils import clean_label, map_string_to_page_class, read_image_bytes
@@ -18,16 +19,61 @@
 logger = logging.getLogger(__name__)
 
 
+class PixtralImageSource(BaseModel):
+    """Raw bytes payload for an image."""
+
+    bytes_: bytes = Field(alias="bytes")
+
+
+class PixtralImage(BaseModel):
+    """Image content block containing its format and raw bytes source."""
+
+    format_: str = Field(alias="format")
+    source: PixtralImageSource
+
+
+class PixtralMessage(BaseModel):
+    """A single content block in a Pixtral conversation, either text or image."""
+
+    text: str | None = None
+    image: PixtralImage | None = None
+
+
+class PixtralMessageStack(BaseModel):
+    """A full conversation turn with a role (e.g. 'user') and a list of content blocks."""
+
+    role: str
+    content: list[PixtralMessage]
+
+
+class PixtralResponseOutput(BaseModel):
+    """The output field of response, wrapping the assistant message."""
+
+    message: PixtralMessageStack
+
+
+class PixtralResponse(BaseModel):
+    """Top-level response, containing the model output."""
+
+    output: PixtralResponseOutput
+
+
 class RateLimiter:
     """Simple token bucket QPS limiter."""
 
     def __init__(self, qps: float):
+        """Initialise the rate limiter with a target queries-per-second rate.
+
+        Args:
+            qps (float): Maximum number of requests allowed per second.
+        """
         self.qps = max(0.1, qps)
         self.lock = threading.Lock()
         self.tokens = 0.0
         self.last = time.monotonic()
 
     def acquire(self):
+        """Block until a token is available, then consume it."""
         while True:
             with self.lock:
                 now = time.monotonic()
@@ -53,23 +99,107 @@ def is_throttle_error(e) -> bool:
         return False
 
 
-class PixtralClassifier(Classifier):
-    """Page Classifier using Pixtral Large."""
+class PixtralConnector:
+    """Low-level client for the Pixtral model.
+
+    Handles authentication, rate limiting, and retries with exponential
+    back-off and full jitter when API throttles requests.
+    """
 
     def __init__(
         self,
         config: dict,
         aws_config: dict,
-        fallback_classifier=None,
     ):
-        self.type = ClassifierTypes.PIXTRAL
+        """Initialise client and rate-limiting settings.
+
+        Args:
+            config (dict): Pixtral configuration dict.
+            aws_config (dict): AWS settings dict.
+        """
         self.config = config
-        self.prompts_dict = read_params(config["prompt_path"])[config["prompt_version"]]
         self.client = boto3.client("bedrock-runtime", region_name=aws_config["region"])
-        self.fallback_classifier = fallback_classifier
         self.model_id = aws_config["model_id"]
+        self._stats = {"throttles": 0, "retries": 0}
+        self.qps = config.get("qps", 2.0)
+        self.max_retries = config.get("max_retries", 6)
+        self.backoff_base = config.get("backoff_base", 0.4)
+        self.backoff_cap = config.get("backoff_cap", 8.0)
+        self._rl = RateLimiter(self.qps)
+        self.max_doc_size = self.config["max_document_size_mb"] - self.config["slack_size_mb"]
+
+    def _send_conversation(self, message: PixtralMessageStack, system: PixtralMessage) -> PixtralResponse:
+        """Send a single-turn conversation to the Pixtral model.
+
+        Args:
+            message (PixtralMessageStack): The user message stack to send.
+            system (PixtralMessage): The system prompt message.
+
+        Returns:
+            PixtralResponse: The validated model response.
+        """
+        attempt = 0
+        while True:
+            self._rl.acquire()  # ensure we dont exceed QPS
+            try:
+                answer = self.client.converse(
+                    modelId=self.model_id,
+                    messages=[message.model_dump(by_alias=True, exclude_none=True)],
+                    system=[system.model_dump(by_alias=True, exclude_none=True)],
+                    inferenceConfig={
+                        "maxTokens": self.config.get("max_tokens", 5),
+                        "temperature": self.config.get("temperature", 0.2),
+                    },
+                )
+                return PixtralResponse.model_validate(answer)
+            except ClientError as e:
+                # Retry on throttling
+                if is_throttle_error(e) and attempt < self.max_retries:
+                    delay = min(self.backoff_cap, self.backoff_base * (2**attempt))
+                    # full jitter
+                    delay *= random.uniform(0.5, 1.5)
+                    logger.warning(f"Bedrock throttled (attempt {attempt + 1}/{self.max_retries}); sleep {delay:.2f}s")
+                    time.sleep(delay)
+                    attempt += 1
+
+                    self._stats["retries"] += 1
+                    if "Throttl" in str(e):
+                        self._stats["throttles"] += 1
+                    continue
+                raise  # not retryable or out of retries
+            except Exception:
+                # Non-ClientError; retry a couple of times
+                if attempt < 2:
+                    time.sleep(0.5 * (attempt + 1))
+                    attempt += 1
+                    continue
+                raise
+
+
+class PixtralClassifier(PixtralConnector, Classifier):
+    """Page classifier that uses the Pixtral vision model."""
 
-        self.system_content = [{"text": self.prompts_dict["system_prompt"]}]
+    def __init__(
+        self,
+        config: dict,
+        aws_config: dict,
+        fallback_classifier: Callable = None,
+    ):
+        """Initialise the classifier, loading prompts and example images.
+
+        Args:
+            config (dict): Pixtral configuration dict.
+            aws_config (dict): AWS settings dict.
+            fallback_classifier (Callable): Optional classifier to use when Pixtral
+                returns an unrecognised label or errors out.
+        """
+        # Create connection to remote model
+        PixtralConnector.__init__(self, config=config, aws_config=aws_config)
+
+        self.type = ClassifierTypes.PIXTRAL
+        self.prompts_dict = read_params(config["prompt_path"])[config["prompt_version"]]
+        self.fallback_classifier = fallback_classifier
+        self.system_content = PixtralMessage(text=self.prompts_dict["system_prompt"])
         self.examples_bytes = {
             "borehole": read_image_bytes(config["borehole_img_path"]),
             "text": read_image_bytes(config["text_img_path"]),
@@ -79,12 +209,6 @@ def __init__(
             "diagram": read_image_bytes(config["diagram_img_path"]),
             "table": read_image_bytes(config["table_img_path"]),
         }
-        self._stats = {"throttles": 0, "retries": 0}
-        self.qps = config.get("qps", 2.0)
-        self.max_retries = config.get("max_retries", 6)
-        self.backoff_base = config.get("backoff_base", 0.4)
-        self.backoff_cap = config.get("backoff_cap", 8.0)
-        self._rl = RateLimiter(self.qps)
 
     def determine_class(
         self, page: pymupdf.Page, page_number: int, context_builder: Callable[[], PageContext] = None, **kwargs
@@ -102,14 +226,12 @@ def determine_class(
         Returns:
             PageClasses: The predicted page class.
         """
-        max_doc_size = self.config["max_document_size_mb"] - self.config["slack_size_mb"]
-        image_bytes = get_page_image_bytes(page, page_number, max_mb=max_doc_size)
-
-        conversation = self._build_conversation(image_bytes=image_bytes)
+        image_bytes = get_page_image_bytes(page, max_mb=self.max_doc_size)
+        message = self._build_conversation(image_bytes=image_bytes)
 
         try:
-            response = self._send_conversation(conversation)
-            raw_label = response["output"]["message"]["content"][0]["text"]
+            response = self._send_conversation(message=message, system=self.system_content)
+            raw_label = response.output.message.content[0].text
 
             label = clean_label(raw_label)
             category = map_string_to_page_class(label)
@@ -138,52 +260,90 @@ def determine_class(
                 )
             return PageClasses.UNKNOWN
 
-    def _build_conversation(self, image_bytes: bytes) -> list[dict]:
-        content = [
-            {"image": {"format": "jpeg", "source": {"bytes": self.examples_bytes[text.strip("@")]}}}
-            if text.startswith("@")  # @category encodes the image of the category and adds it to the content
-            else {"text": text}
+    def _build_conversation(self, image_bytes: bytes) -> PixtralMessageStack:
+        """Build the user message containing few-shot examples and the target image.
+
+        Args:
+            image_bytes: Eencoded bytes of the page to classify.
+
+        Returns:
+            PixtralMessageStack: A user turn ready to send.
+        """
+        # List of examples for pixtral model
+        content_examples = [
+            PixtralMessage(
+                image=PixtralImage(
+                    format="jpeg",
+                    source=PixtralImageSource(bytes=self.examples_bytes[text.strip("@")]),
+                )
+            )
+            if text.startswith("@")
+            else PixtralMessage(text=text)
             for text in self.prompts_dict.get("examples_prompt", [])
         ]
-        content.append({"text": self.prompts_dict["user_prompt"]})
-        content.append({"image": {"format": "jpeg", "source": {"bytes": image_bytes}}})
 
-        return [{"role": "user", "content": content}]
+        # User prompt with content to classify
+        content_user_text = PixtralMessage(text=self.prompts_dict["user_prompt"])
+        content_user_img = PixtralMessage(
+            image=PixtralImage(
+                format="jpeg",
+                source=PixtralImageSource(bytes=image_bytes),
+            ),
+        )
 
-    def _send_conversation(self, conversation: list) -> dict:
-        """Sends the conversation to Bedrock with retry-on-throttle."""
-        attempt = 0
-        while True:
-            self._rl.acquire()  # ensure we dont exceed QPS
-            try:
-                return self.client.converse(
-                    modelId=self.model_id,
-                    messages=conversation,
-                    system=self.system_content,
-                    inferenceConfig={
-                        "maxTokens": self.config.get("max_tokens", 5),
-                        "temperature": self.config.get("temperature", 0.2),
-                    },
+        return PixtralMessageStack(role="user", content=content_examples + [content_user_text, content_user_img])
+
+
+class PixtralFeatureExtraction(PixtralConnector):
+    """Uses the Pixtral vision model to extract features from PDF pages."""
+
+    def __init__(self, config: dict, aws_config: dict, system_prompt: str):
+        """Initialise the extractor with a custom system prompt.
+
+        Args:
+            config (dict): Pixtral configuration dict.
+            aws_config (dict): AWS settings dict.
+            system_prompt (Callable): Instruction text sent as the system message for
+                every extraction request.
+        """
+        # Create connection to remote model
+        PixtralConnector.__init__(self, config=config, aws_config=aws_config)
+        self.system_prompt = PixtralMessage(text=system_prompt)
+
+    def _build_conversation(self, image_bytes: bytes) -> PixtralMessageStack:
+        """Build a minimal user message containing only the target page image.
+
+        Args:
+            image_bytes (bytes): Encoded bytes of the page to process.
+
+        Returns:
+            PixtralMessageStack: A 'user' turn with a single image content block.
+        """
+        # List of examples for pixtral model
+        return PixtralMessageStack(
+            role="user",
+            content=[
+                PixtralMessage(
+                    image=PixtralImage(
+                        format="jpeg",
+                        source=PixtralImageSource(bytes=image_bytes),
+                    )
                 )
-            except ClientError as e:
-                # Retry on throttling
-                if is_throttle_error(e) and attempt < self.max_retries:
-                    delay = min(self.backoff_cap, self.backoff_base * (2**attempt))
-                    # full jitter
-                    delay *= random.uniform(0.5, 1.5)
-                    logger.warning(f"Bedrock throttled (attempt {attempt + 1}/{self.max_retries}); sleep {delay:.2f}s")
-                    time.sleep(delay)
-                    attempt += 1
+            ],
+        )
 
-                    self._stats["retries"] += 1
-                    if "Throttl" in str(e):
-                        self._stats["throttles"] += 1
-                    continue
-                raise  # not retryable or out of retries
-            except Exception:
-                # Non-ClientError; retry a couple of times
-                if attempt < 2:
-                    time.sleep(0.5 * (attempt + 1))
-                    attempt += 1
-                    continue
-                raise
+    def find(self, page: pymupdf.Page) -> str:
+        """Extract a feature from a single PDF page using the Pixtral model.
+
+        Args:
+            page (pymupdf.Page): The PyMuPDF page object to process.
+
+        Returns:
+            str: The raw text returned by the model (e.g. an extracted title).
+        """
+        # User prompt with content to classify
+        image_bytes = get_page_image_bytes(page, max_mb=self.max_doc_size)
+        content_user = self._build_conversation(image_bytes=image_bytes)
+
+        response = self._send_conversation(message=content_user, system=self.system_prompt)
+        return response.output.message.content[0].text

src/page_graphics.py

@@ -51,15 +51,15 @@ def get_images_from_page(page: pymupdf.Page) -> list[ImageRect]:
     return extracted_images
 
 
-def get_page_image_bytes(page: pymupdf.Page, page_number: int, max_mb: float = 4.5) -> bytes:
+def get_page_image_bytes(page: pymupdf.Page, max_mb: float = 4.5) -> bytes:
     """Returns JPEG image bytes of a single PDF page. Downscales if image exceeds allowed size."""
     max_bytes = int(max_mb * 1024 * 1024)
     scale = 1.0
 
     for attempt in range(10):
         # Render and convert to JPEG
         with pymupdf.open() as pdf_doc:
-            pdf_doc.insert_pdf(page.parent, from_page=page_number, to_page=page_number)
+            pdf_doc.insert_pdf(page.parent, from_page=0, to_page=0)
             page_bytes = pdf_doc.tobytes(deflate=True, garbage=3, use_objstms=1)
 
         image_bytes = convert_pdf_to_jpeg(page_bytes, scale=scale)