""" ocr_confidence.py ----------------- Reads an existing olmOCR .md output file and scores each page for transcription confidence using the DeepInfra API. Does NOT re-run OCR. Requirements: pip install requests Usage: python ocr_confidence.py --api_key "YOUR_KEY" python ocr_confidence.py --api_key "YOUR_KEY" --input "path/to/file.md" --output "path/to/out.json" """ import argparse import json import sys import time import warnings from pathlib import Path import requests # --------------------------------------------------------------------------- # Load .env from project root (if present) — no external dependencies needed # --------------------------------------------------------------------------- def _load_dotenv() -> None: env_path = Path(__file__).resolve().parent.parent / ".env" if not env_path.exists(): return import os with open(env_path, encoding="utf-8") as fh: for line in fh: line = line.strip() if not line or line.startswith("#") or "=" not in line: continue key, _, value = line.partition("=") os.environ.setdefault(key.strip(), value.strip()) _load_dotenv() # --------------------------------------------------------------------------- # Configuration — mirrors ocr_wardiaries.py # --------------------------------------------------------------------------- DEEPINFRA_API_URL = "https://api.deepinfra.com/v1/openai/chat/completions" # olmOCR is a vision model — use a chat LLM for text-based confidence review MODEL = "google/gemma-3-27b-it" MAX_RETRIES = 3 RETRY_DELAY = 5 # seconds between retries _PROJECT_ROOT = Path(__file__).resolve().parent.parent DEFAULT_INPUT = _PROJECT_ROOT / "Inputs" / "ocr-output" / "Calgary-Highlanders_War-Diary_Sep44_olmocr.md" DEFAULT_OUTPUT = _PROJECT_ROOT / "Inputs" / "ocr-output" / "Calgary-Highlanders_War-Diary_Sep44_confidence.json" CONFIDENCE_PROMPT = ( "You are reviewing an OCR transcription of a WWII war diary. Your job is to identify words that may have been misread by the OCR scanner — not to correct the soldier's original spelling or interpret abbreviations. Rate OCR accuracy only. Do not suggest what words "should" be. Respond in JSON only: {"score": 7, "uncertain_words": ["Loon", "Fme"], "notes": "Possible OCR misread in line 3"} ) ERROR_RESULT = {"score": 0, "uncertain_words": [], "notes": "API error"} # --------------------------------------------------------------------------- # Parse olmOCR markdown into {page_num: text} dict # Same logic as parseOCRByPage() in p44-ocr-viewer.html # --------------------------------------------------------------------------- def parse_ocr_by_page(raw_text: str) -> dict[int, str]: """Split olmOCR markdown on '## Page N' headings into a page-keyed dict.""" pages: dict[int, str] = {} lines = raw_text.split("\n") page_num: int | None = None buf: list[str] = [] for line in lines: m_head = _PAGE_HEADING.match(line) if m_head: if page_num is not None: pages[page_num] = "\n".join(buf) page_num = int(m_head.group(1)) buf = [line] # keep the heading at the top elif page_num is not None: buf.append(line) if page_num is not None: pages[page_num] = "\n".join(buf) return pages import re as _re _PAGE_HEADING = _re.compile(r"^## Page (\d+)\s*$") # --------------------------------------------------------------------------- # API call # --------------------------------------------------------------------------- def score_page(page_text: str, page_num: int, api_key: str) -> dict: """ Send one page's OCR text to DeepInfra for confidence scoring. Returns a dict with keys: score, uncertain_words, notes. On failure returns ERROR_RESULT. """ headers = { "Content-Type": "application/json", "Authorization": f"Bearer {api_key}", } payload = { "model": MODEL, "messages": [ { "role": "user", "content": ( f"{CONFIDENCE_PROMPT}\n\n" f"--- BEGIN OCR TEXT (page {page_num}) ---\n" f"{page_text}\n" f"--- END OCR TEXT ---" ), } ], "max_tokens": 256, "temperature": 0.0, } for attempt in range(1, MAX_RETRIES + 1): try: response = requests.post( DEEPINFRA_API_URL, headers=headers, json=payload, timeout=60 ) response.raise_for_status() raw_content = response.json()["choices"][0]["message"]["content"].strip() # Strip markdown code fences if the model wrapped the JSON if raw_content.startswith("```"): raw_content = _re.sub(r"^```[a-z]*\n?", "", raw_content) raw_content = _re.sub(r"\n?```$", "", raw_content) result = json.loads(raw_content) # Validate expected keys; fill missing ones with defaults score = int(result.get("score", 0)) uncertain_words = result.get("uncertain_words", []) notes = result.get("notes", "") if not isinstance(uncertain_words, list): uncertain_words = [] return {"score": score, "uncertain_words": uncertain_words, "notes": notes} except requests.exceptions.HTTPError as exc: warnings.warn(f"Page {page_num}: HTTP error on attempt {attempt}/{MAX_RETRIES}: {exc}") if attempt < MAX_RETRIES: time.sleep(RETRY_DELAY) except json.JSONDecodeError as exc: warnings.warn( f"Page {page_num}: Malformed JSON from API on attempt {attempt}/{MAX_RETRIES}: {exc}" ) if attempt < MAX_RETRIES: time.sleep(RETRY_DELAY) except Exception as exc: # noqa: BLE001 warnings.warn(f"Page {page_num}: Unexpected error on attempt {attempt}/{MAX_RETRIES}: {exc}") if attempt < MAX_RETRIES: time.sleep(RETRY_DELAY) warnings.warn(f"Page {page_num}: All {MAX_RETRIES} attempts failed — storing error result.") return dict(ERROR_RESULT) # return a fresh copy # --------------------------------------------------------------------------- # Summary helpers # --------------------------------------------------------------------------- def print_summary(results: dict) -> None: high = sum(1 for v in results.values() if v["score"] >= 8) med = sum(1 for v in results.values() if 5 <= v["score"] <= 7) low = sum(1 for v in results.values() if 1 <= v["score"] <= 4) err = sum(1 for v in results.values() if v["score"] == 0) total = len(results) print("\n" + "=" * 50) print(f"CONFIDENCE SCORING COMPLETE — {total} pages scored") print("=" * 50) print(f" High (8-10) : {high:>4} ({high/total*100:.1f}%)" if total else " High (8-10) : 0") print(f" Medium (5-7) : {med:>4} ({med/total*100:.1f}%)" if total else " Medium (5-7) : 0") print(f" Low (1-4) : {low:>4} ({low/total*100:.1f}%)" if total else " Low (1-4) : 0") if err: print(f" API errors : {err:>4}") print("=" * 50) # --------------------------------------------------------------------------- # Main # --------------------------------------------------------------------------- def main() -> None: parser = argparse.ArgumentParser( description="Score OCR confidence for each page of an olmOCR .md output file." ) parser.add_argument( "--api_key", default=None, help="DeepInfra API key (defaults to DEEPINFRA_API_KEY env var)." ) parser.add_argument( "--input", default=str(DEFAULT_INPUT), help=f"Path to olmOCR .md file. Default: {DEFAULT_INPUT}" ) parser.add_argument( "--output", default=str(DEFAULT_OUTPUT), help=f"Path to write confidence JSON. Default: {DEFAULT_OUTPUT}" ) parser.add_argument( "--delay", type=float, default=0.5, help="Seconds to pause between API calls (default: 0.5)." ) args = parser.parse_args() import os api_key = args.api_key or os.environ.get("DEEPINFRA_API_KEY", "") if not api_key: print("ERROR: No API key provided. Use --api_key or set DEEPINFRA_API_KEY in .env", file=sys.stderr) sys.exit(1) input_path = Path(args.input) output_path = Path(args.output) # ── Read source file ────────────────────────────────────────────────────── if not input_path.exists(): print(f"ERROR: Input file not found: {input_path}", file=sys.stderr) sys.exit(1) print(f"Reading OCR source: {input_path}") raw_text = input_path.read_text(encoding="utf-8") pages = parse_ocr_by_page(raw_text) if not pages: print("ERROR: No '## Page N' headings found in the input file.", file=sys.stderr) sys.exit(1) total_pages = len(pages) sorted_pages = sorted(pages.keys()) print(f"Found {total_pages} pages (Page {sorted_pages[0]} – {sorted_pages[-1]}).") # ── Load existing results (resume support) ──────────────────────────────── results: dict[str, dict] = {} if output_path.exists(): try: results = json.loads(output_path.read_text(encoding="utf-8")) already_done = len(results) print(f"Resuming — {already_done} page(s) already scored, skipping them.") except (json.JSONDecodeError, OSError) as exc: warnings.warn(f"Could not read existing output ({exc}); starting fresh.") results = {} output_path.parent.mkdir(parents=True, exist_ok=True) # ── Score each page ─────────────────────────────────────────────────────── scored_this_run = 0 for page_num in sorted_pages: page_key = str(page_num) if page_key in results: continue # already scored — skip page_text = pages[page_num] result = score_page(page_text, page_num, api_key) results[page_key] = result scored_this_run += 1 # Progress line score_str = str(result["score"]) if result["score"] > 0 else "ERR" print(f" Page {page_num}/{sorted_pages[-1]} — score: {score_str}") # Write after every page so a crash loses minimal work output_path.write_text( json.dumps(results, indent=2, ensure_ascii=False), encoding="utf-8" ) if scored_this_run > 0 and page_num != sorted_pages[-1]: time.sleep(args.delay) # ── Final save & summary ────────────────────────────────────────────────── output_path.write_text( json.dumps(results, indent=2, ensure_ascii=False), encoding="utf-8" ) print(f"\nResults written to: {output_path}") print_summary(results) if __name__ == "__main__": main()