#!/usr/bin/env python3 """ Smoke Signal — Stage 7: LLM Normalisation Layer ================================================= Takes low-confidence or visually complex regions from Stage 5 and sends them to an LLM (via HF Inference API) for cleanup. Rules (non-negotiable): - LLM receives: page image crop + raw OCR candidates - LLM must return strict JSON only — no free text, no preamble - Uncertain words must be marked uncertain, NEVER silently guessed - LLM may NOT invent text that isn't visually present - Only low-confidence / flagged pages are sent (cost + drift control) - Every call logs: prompt version, model, input hash, output hash, cost Output per region: { "text_raw": "original OCR text", "text_clean": "LLM corrected text", "uncertain_words": ["word1", "word2"], "confidence_notes": "why confidence is low", "changed_from_ocr": true/false, "visible_context": "brief note on what the LLM can see" } Usage: python scripts/05_llm_normalise.py python scripts/05_llm_normalise.py --book-id SS-BOOK-0001 python scripts/05_llm_normalise.py --dry-run python scripts/05_llm_normalise.py --confidence-below 0.80 """ import argparse import base64 import hashlib import json import sys import time from datetime import datetime from pathlib import Path from typing import Optional # ── Paths ────────────────────────────────────────────────────────────────────── ROOT = Path(__file__).resolve().parents[1] MANIFEST_CSV = ROOT / "manifest" / "source_manifest.csv" REGIONS_DIR = ROOT / "regions" OCR_RAW_DIR = ROOT / "ocr_raw" RENDERS_DIR = ROOT / "renders" LOGS_DIR = ROOT / "logs" REVIEW_DIR = ROOT / "review" CLEANED_DIR = ROOT / "regions" / "cleaned" CLEANED_DIR.mkdir(parents=True, exist_ok=True) LOGS_DIR.mkdir(parents=True, exist_ok=True) # ── Config (freeze before running — do not change mid-batch) ────────────────── CONFIG = { "config_version": "ss_llm_v0.1", "prompt_version": "ss_prompt_v0.1", "model": "meta-llama/Llama-3.2-11B-Vision-Instruct", # HF Inference API "confidence_threshold": 0.75, # only send pages below this "max_pages_per_run": 50, # cost control — cap per batch "temperature": 0.1, # low = deterministic, no hallucination "max_new_tokens": 512, "eligible_statuses": ["ocred"], "story_classes": ["narration", "dialogue-speech-bubble", "caption"], } # ── Prompt (versioned — never change without bumping prompt_version) ─────────── SYSTEM_PROMPT = """You are a precise OCR correction assistant for children's picture books. RULES — follow exactly: 1. Return ONLY valid JSON. No preamble, no explanation, no markdown fences. 2. Correct OCR errors you can see in the image. Do NOT invent text. 3. If a word is unclear or unreadable, add it to uncertain_words — do NOT guess. 4. Keep the author's exact words, punctuation, and line breaks. 5. Do not add, remove, or reorder words unless fixing a clear OCR error. 6. changed_from_ocr must be true only if you changed something. Return this exact JSON structure: { "text_clean": "corrected text here", "uncertain_words": ["list", "of", "unclear", "words"], "confidence_notes": "brief note on what made this hard to read", "changed_from_ocr": false, "visible_context_notes": "brief note on what you can see in the image" }""" USER_PROMPT_TEMPLATE = """Here is the raw OCR output for a picture book page region: RAW OCR: {raw_text} REGION CLASS: {region_class} OCR CONFIDENCE: {confidence} Please examine the image crop and return the corrected JSON.""" # ── HF Inference API ─────────────────────────────────────────────────────────── def _get_hf_token() -> Optional[str]: """Get HF token from environment or huggingface_hub cache.""" import os token = os.environ.get("HF_TOKEN") or os.environ.get("HUGGINGFACE_HUB_TOKEN") if token: return token try: from huggingface_hub import get_token return get_token() except Exception: return None def call_llm_with_image( image_path: Path, raw_text: str, region_class: str, confidence: float, ) -> dict: """ Call HF Inference API with image + OCR text. Returns parsed JSON result or error dict. """ import urllib.request import urllib.error token = _get_hf_token() if not token: return { "error": "no_hf_token", "text_clean": raw_text, "uncertain_words": [], "confidence_notes": "HF token not found — set HF_TOKEN env var or run huggingface-cli login", "changed_from_ocr": False, } # Encode image as base64 try: with open(image_path, "rb") as f: image_b64 = base64.b64encode(f.read()).decode("utf-8") image_ext = image_path.suffix.lower().replace(".", "") media_type = f"image/{image_ext if image_ext in ('png','jpg','jpeg','webp') else 'png'}" except Exception as e: return {"error": f"image_load_failed: {e}", "text_clean": raw_text, "uncertain_words": [], "changed_from_ocr": False} user_message = USER_PROMPT_TEMPLATE.format( raw_text=raw_text[:800], # truncate for token budget region_class=region_class, confidence=round(confidence, 3), ) payload = json.dumps({ "model": CONFIG["model"], "messages": [ {"role": "system", "content": SYSTEM_PROMPT}, { "role": "user", "content": [ { "type": "image_url", "image_url": {"url": f"data:{media_type};base64,{image_b64}"} }, {"type": "text", "text": user_message} ] } ], "max_tokens": CONFIG["max_new_tokens"], "temperature": CONFIG["temperature"], }).encode("utf-8") api_url = f"https://api-inference.huggingface.co/models/{CONFIG['model']}/v1/chat/completions" req = urllib.request.Request( api_url, data=payload, headers={ "Authorization": f"Bearer {token}", "Content-Type": "application/json", }, method="POST", ) try: with urllib.request.urlopen(req, timeout=30) as resp: response_data = json.loads(resp.read().decode("utf-8")) raw_response = response_data["choices"][0]["message"]["content"].strip() except urllib.error.HTTPError as e: return {"error": f"http_{e.code}: {e.reason}", "text_clean": raw_text, "uncertain_words": [], "changed_from_ocr": False} except Exception as e: return {"error": str(e), "text_clean": raw_text, "uncertain_words": [], "changed_from_ocr": False} # Parse JSON response — strip markdown fences if model added them try: clean = raw_response.strip() if clean.startswith("```"): clean = clean.split("```")[1] if clean.startswith("json"): clean = clean[4:] result = json.loads(clean.strip()) except json.JSONDecodeError: return { "error": "invalid_json_response", "raw_response": raw_response[:500], "text_clean": raw_text, "uncertain_words": [], "changed_from_ocr": False, } return result # ── Input hash (for audit trail) ────────────────────────────────────────────── def _input_hash(text: str, image_path: Path) -> str: h = hashlib.sha256() h.update(text.encode("utf-8")) if image_path.exists(): with open(image_path, "rb") as f: h.update(f.read(4096)) # first 4kb sufficient for fingerprint return h.hexdigest()[:16] def _output_hash(result: dict) -> str: return hashlib.sha256( json.dumps(result, sort_keys=True).encode("utf-8") ).hexdigest()[:16] # ── Load region data ─────────────────────────────────────────────────────────── def load_ordered_regions(book_id: str) -> Optional[dict]: path = REGIONS_DIR / f"{book_id}_ordered_regions.json" return json.load(open(path)) if path.exists() else None # ── Per-region normalisation ─────────────────────────────────────────────────── def normalise_region( book_id: str, page_num: int, region: dict, render_path: Optional[str], dry_run: bool, ) -> dict: """Normalise a single region via LLM. Returns enriched region dict.""" raw_text = region.get("text", "").strip() region_class = region.get("region_class", "narration") confidence = region.get("confidence", 1.0) region_id = region.get("region_id", f"{book_id}_p{page_num:04d}") if dry_run: return { **region, "text_clean": raw_text, "uncertain_words": [], "confidence_notes": "dry-run", "changed_from_ocr": False, "visible_context_notes": "dry-run", "llm_model": CONFIG["model"], "prompt_version": CONFIG["prompt_version"], "normalised_at": datetime.utcnow().isoformat() + "Z", "dry_run": True, } # Find render image img_path = None if render_path: candidate = ROOT / render_path if not Path(render_path).is_absolute() else Path(render_path) if candidate.exists(): img_path = candidate if img_path is None: # Try to find any render for this book/page book_renders = RENDERS_DIR / book_id if book_renders.exists(): candidates = sorted(book_renders.glob(f"{book_id}_page_{page_num:04d}_*.png")) if candidates: img_path = candidates[0] if img_path is None: return { **region, "text_clean": raw_text, "uncertain_words": [], "confidence_notes": "no_render_available", "changed_from_ocr": False, "error": "no_render_found", } in_hash = _input_hash(raw_text, img_path) llm_result = call_llm_with_image(img_path, raw_text, region_class, confidence) out_hash = _output_hash(llm_result) return { **region, "text_clean": llm_result.get("text_clean", raw_text), "uncertain_words": llm_result.get("uncertain_words", []), "confidence_notes": llm_result.get("confidence_notes", ""), "changed_from_ocr": llm_result.get("changed_from_ocr", False), "visible_context_notes": llm_result.get("visible_context_notes", ""), "llm_model": CONFIG["model"], "prompt_version": CONFIG["prompt_version"], "input_hash": in_hash, "output_hash": out_hash, "llm_error": llm_result.get("error"), "normalised_at": datetime.utcnow().isoformat() + "Z", } # ── Per-book normalisation ──────────────────────────────────────────────────── def normalise_book( book_id: str, filename: str, confidence_threshold: float, max_pages: int, dry_run: bool, ) -> tuple: print(f"\n [{book_id}] {filename}") regions_data = load_ordered_regions(book_id) if regions_data is None: print(f" ✗ No region data. Run 04_region_detector.py first.") return None, {"error": "no_region_data"} pages_to_process = [] for page in regions_data.get("pages", []): page_conf = page.get("page_confidence", 1.0) route = page.get("route", "embedded_text") if route == "embedded_text": continue if page_conf < confidence_threshold: pages_to_process.append(page) if not pages_to_process: print(f" ✓ No pages below confidence threshold {confidence_threshold} — skipping.") return regions_data, {} # Apply page cap if len(pages_to_process) > max_pages: print(f" ⚠️ {len(pages_to_process)} pages need normalisation — capping at {max_pages}") pages_to_process = pages_to_process[:max_pages] print(f" Pages to normalise: {len(pages_to_process)}") total_changed = 0 total_uncertain = 0 total_errors = 0 call_log = [] # Process each page page_index = {p["page_number"]: i for i, p in enumerate(regions_data["pages"])} for page in pages_to_process: page_num = page["page_number"] page_conf = page.get("page_confidence", 0) render_path = None # Find render path from OCR raw data ocr_path = OCR_RAW_DIR / book_id / f"{book_id}_ocr_raw.json" if ocr_path.exists(): ocr_data = json.load(open(ocr_path)) for ocr_page in ocr_data.get("pages", []): if ocr_page.get("page_number") == page_num: render_path = ocr_page.get("render_path") break print(f" Page {page_num:3d} (conf={page_conf:.2f}): ", end="", flush=True) normalised_regions = [] for region in page.get("regions", []): region_class = region.get("region_class", "narration") # Only normalise story-relevant regions if region_class not in CONFIG["story_classes"]: normalised_regions.append(region) continue region_conf = region.get("confidence", 1.0) if region_conf >= confidence_threshold: normalised_regions.append(region) continue result = normalise_region(book_id, page_num, region, render_path, dry_run) if result.get("changed_from_ocr"): total_changed += 1 if result.get("uncertain_words"): total_uncertain += len(result["uncertain_words"]) if result.get("llm_error"): total_errors += 1 call_log.append({ "book_id": book_id, "page_number": page_num, "region_id": region.get("region_id"), "input_hash": result.get("input_hash"), "output_hash": result.get("output_hash"), "changed": result.get("changed_from_ocr", False), "uncertain_count": len(result.get("uncertain_words", [])), "error": result.get("llm_error"), "model": CONFIG["model"], "prompt_version": CONFIG["prompt_version"], }) normalised_regions.append(result) # Update page in regions data idx = page_index.get(page_num) if idx is not None: regions_data["pages"][idx]["regions"] = normalised_regions regions_data["pages"][idx]["normalised"] = True changed_count = sum(1 for r in normalised_regions if r.get("changed_from_ocr")) uncertain_count = sum(len(r.get("uncertain_words", [])) for r in normalised_regions) print(f"changed={changed_count} uncertain_words={uncertain_count}") # ── Save cleaned regions ────────────────────────────────────────────────── if not dry_run: regions_data["normalised_at"] = datetime.utcnow().isoformat() + "Z" regions_data["prompt_version"] = CONFIG["prompt_version"] regions_data["llm_model"] = CONFIG["model"] regions_data["normalisation_log"] = call_log cleaned_path = CLEANED_DIR / f"{book_id}_cleaned_regions.json" with open(cleaned_path, "w", encoding="utf-8") as f: json.dump(regions_data, f, indent=2) print(f" Cleaned regions → {cleaned_path.relative_to(ROOT)}") # Log file log_path = LOGS_DIR / f"llm_calls_{book_id}_{datetime.utcnow().strftime('%Y%m%d')}.json" with open(log_path, "w") as f: json.dump(call_log, f, indent=2) print(f" Total changed={total_changed} | uncertain_words={total_uncertain} | errors={total_errors}") return regions_data, {} # ── Main ─────────────────────────────────────────────────────────────────────── def main(): parser = argparse.ArgumentParser(description="Smoke Signal — Stage 7: LLM Normalisation") parser.add_argument("--book-id", help="Process a single book by ID") parser.add_argument("--batch-id", help="Tag this run with a batch ID") parser.add_argument("--confidence-below", type=float, default=CONFIG["confidence_threshold"], help=f"Only process pages below this confidence (default {CONFIG['confidence_threshold']})") parser.add_argument("--max-pages", type=int, default=CONFIG["max_pages_per_run"], help=f"Max pages per book per run (default {CONFIG['max_pages_per_run']})") parser.add_argument("--dry-run", action="store_true") args = parser.parse_args() run_id = args.batch_id or f"SS-RUN-{datetime.utcnow().strftime('%Y%m%d-%H%M%S')}" dry_run = args.dry_run print(f"\n{'='*60}") print(f" Smoke Signal — Stage 7: LLM Normalisation") print(f" Run ID : {run_id}") print(f" Model : {CONFIG['model']}") print(f" Prompt : {CONFIG['prompt_version']}") print(f" Threshold : conf < {args.confidence_below}") print(f" Max pages : {args.max_pages}") if dry_run: print(f" Mode : DRY RUN") print(f"{'='*60}") # Find books with region data if args.book_id: region_files = [REGIONS_DIR / f"{args.book_id}_ordered_regions.json"] else: region_files = sorted(REGIONS_DIR.glob("*_ordered_regions.json")) if not region_files: print("\n No region files found. Run 04_region_detector.py first.") sys.exit(0) print(f"\n Books to process: {len(region_files)}") results = [] t_start = time.time() for region_file in region_files: book_id = region_file.stem.replace("_ordered_regions", "") filename = book_id # fallback _, error = normalise_book( book_id, filename, args.confidence_below, args.max_pages, dry_run, ) results.append({"book_id": book_id, "error": error}) elapsed = round(time.time() - t_start, 1) succeeded = sum(1 for r in results if not r.get("error")) print(f"\n{'─'*60}") print(f" Books processed : {len(results)}") print(f" Succeeded : {succeeded}") print(f" Time : {elapsed}s") print(f"{'─'*60}") print(f"\n Cleaned regions → {CLEANED_DIR.relative_to(ROOT)}") print(f" Next: Run 06_review_workbench.py (Stage 8)\n") if __name__ == "__main__": main()