ocr/transcribe.py

"""Blind handwriting transcription via an OpenAI-compatible VLM (Qwen3-VL on
Modal, see modal/ocr/modal_qwen3vl.py).

The one rule (specs/ocr.md §3): this module NEVER sees the reference text. It
asks the model for a faithful, uncorrected transcription of the page; grading
happens afterwards in ocr.grading. Note there is deliberately no ``reference``
parameter anywhere here."""

import base64
import json
import mimetypes
from pathlib import Path

# specs/ocr.md §10, hardened. The "word by word" + explicit spelling-error +
# umlaut-dot framing beat both the original prompt and a "letter by letter"
# variant on a real handwriting sample: on Qwen3-VL-8B it fixed gross misreads
# without adding noise. The one residual failure is the language-model prior
# auto-correcting missing umlauts on very common words (e.g. a deliberately
# written "Fruhstück" still came back "Frühstück") — see specs/ocr.md §9.
SYSTEM_PROMPT = (
    "You are a forensic handwriting transcriber. Transcribe the page WORD BY "
    "WORD, copying each written word exactly as it appears on the paper.\n"
    "\n"
    "This text was written in a SPELLING TEST and DELIBERATELY CONTAINS "
    "SPELLING ERRORS. Reproduce every error faithfully. NEVER replace a word "
    "with its dictionary-correct spelling. If a word is written wrong, "
    "transcribe it exactly as wrongly written.\n"
    "\n"
    "Umlaut rule (critical): write a/o/u with two dots (a->ä, o->ö, u->ü) ONLY "
    "when you can clearly see two dots above that exact letter. If a u, a or o "
    "has NO dots above it, keep it plain (u, a, o) even if correct German would "
    "use an umlaut. Never add dots that are not on the page; never drop dots "
    "that are.\n"
    "Do NOT change ß to ss or ss to ß. Do NOT translate. Do NOT fix spacing, "
    "grammar, or capitalization.\n"
    "Output ONLY the transcribed text, one line per written line."
)

USER_INSTRUCTION = (
    "Transcribe this page word by word, reproducing every spelling error "
    "exactly as written."
)

def guess_mime(path: str) -> str:
    """MIME type for an image path, defaulting to PNG."""
    mime, _ = mimetypes.guess_type(path)
    return mime or "image/png"


def data_uri(image_bytes: bytes, mime: str) -> str:
    """Encode raw image bytes as an OpenAI ``image_url`` data URI."""
    b64 = base64.b64encode(image_bytes).decode("ascii")
    return f"data:{mime};base64,{b64}"


def encode_image(path: str) -> str:
    """Read an image file and return its data URI."""
    return data_uri(Path(path).read_bytes(), guess_mime(path))


def build_payload(image_uri: str, model: str, temperature: float = 0.0) -> dict:
    """Build the chat-completion request for a blind transcription. No
    reference text appears anywhere (specs/ocr.md §3)."""
    return {
        "model": model,
        "temperature": temperature,
        "messages": [
            {"role": "system", "content": SYSTEM_PROMPT},
            {
                "role": "user",
                "content": [
                    {"type": "text", "text": USER_INSTRUCTION},
                    {"type": "image_url", "image_url": {"url": image_uri}},
                ],
            },
        ],
    }


def parse_response(data: dict) -> str:
    """Extract the transcription text from an OpenAI-style response dict."""
    try:
        return data["choices"][0]["message"]["content"].strip()
    except (KeyError, IndexError, TypeError, AttributeError) as e:
        raise ValueError(f"Unexpected VLM response: {json.dumps(data)[:500]}") from e


def transcribe_image(image_path: str, client, model: str = "qwen3-vl") -> str:
    """Transcribe a handwritten page blind.

    ``client`` is an OpenAI-compatible client pointed at the Modal VLM server
    (build one with ``openai_client.make_client``). Cold-start retries live in
    the client, not here."""
    payload = build_payload(encode_image(image_path), model)
    completion = client.chat.completions.create(**payload)
    return parse_response(completion.model_dump())