feat: support partial dataset downloads and cached OCR pair exports with file hashing

Author: fbkaragoz

ocr_pipeline/routers/documents.py

@@ -34,19 +34,20 @@ async def update_documents_bulk(payload: BulkDocumentUpdate):
     if not payload.document_ids:
         raise HTTPException(status_code=400, detail="document_ids must not be empty")
     db = get_db()
-    for document_id in payload.document_ids:
-        try:
-            await db.update_document_metadata(
-                document_id,
-                group_path=payload.group_path,
-            )
-        except KeyError:
-            raise HTTPException(
-                status_code=404, detail=f"Document not found: {document_id}"
-            )
+    try:
+        await db.update_documents_metadata(
+            payload.document_ids,
+            group_path=payload.group_path,
+        )
+    except KeyError as exc:
+        raise HTTPException(
+            status_code=404, detail=f"Document not found: {exc.args[0]}"
+        )
     documents = await db.list_documents(limit=1000)
-    selected = {document_id for document_id in payload.document_ids}
-    return [_document_summary(doc) for doc in documents if doc["id"] in selected]
+    by_id = {doc["id"]: doc for doc in documents}
+    return [
+        _document_summary(by_id[document_id]) for document_id in payload.document_ids
+    ]
 
 
 @router.get("/api/documents/{document_id}", response_model=DocumentSummary)

ocr_pipeline/routers/runs.py

@@ -1,4 +1,5 @@
 import asyncio
+import hashlib
 import json
 from io import BytesIO
 from pathlib import Path
@@ -310,13 +311,25 @@ async def download_ocr_pairs(
     run_id: str = ID,
     dpi: int = Query(160, ge=50, le=400),
     text_mode: str = Query("clean", pattern="^(clean|raw)$"),
+    document_ids: str | None = Query(None),
 ):
     db = get_db()
     run = await _require_run(run_id)
     if run["status"] != "completed":
         raise HTTPException(status_code=409, detail="Run is not yet completed")
 
     documents = await db.list_run_documents(run_id)
+    selected_ids = {
+        part.strip() for part in (document_ids or "").split(",") if part.strip()
+    } or None
+    if selected_ids:
+        available_ids = {doc["document_id"] for doc in documents}
+        missing = selected_ids - available_ids
+        if missing:
+            raise HTTPException(
+                status_code=404,
+                detail=f"Document not found in run: {sorted(missing)[0]}",
+            )
     pages_by_document = {
         doc["document_id"]: await db.list_pages(run_id, doc["document_id"])
         for doc in documents
@@ -325,13 +338,29 @@ async def download_ocr_pairs(
         doc["document_id"]: await db.get_document(doc["document_id"]) or {}
         for doc in documents
     }
-    exporter = OCRPairExporter(settings.runs_dir / run_id / "dataset" / "ocr_pairs")
+    scope = "all"
+    if selected_ids:
+        scope = hashlib.sha256(
+            ",".join(sorted(selected_ids)).encode("utf-8")
+        ).hexdigest()[:12]
+    export_dir = (
+        settings.runs_dir / run_id / "dataset" / f"ocr_pairs_{text_mode}_{dpi}_{scope}"
+    )
+    cached_bundle = export_dir.with_suffix(".zip")
+    if cached_bundle.exists():
+        return FileResponse(
+            cached_bundle,
+            media_type="application/zip",
+            filename=f"{run_id}-ocr-pairs.zip",
+        )
+    exporter = OCRPairExporter(export_dir)
     result = await asyncio.to_thread(
         exporter.export_run,
         run=run,
         documents=documents,
         pages_by_document=pages_by_document,
         catalog_by_document=catalog_by_document,
+        document_ids=selected_ids,
         dpi=dpi,
         text_mode=text_mode,
     )

ocr_pipeline/routers/ui.py

@@ -1,6 +1,7 @@
 """Static-friendly endpoints for input file management. Output/dataset listing
 moved to /api/runs."""
 
+import hashlib
 from pathlib import Path
 
 from fastapi import APIRouter, HTTPException, UploadFile
@@ -25,12 +26,18 @@ async def upload_pdf(file: UploadFile):
         raise HTTPException(status_code=400, detail="Invalid filename")
 
     settings.input_dir.mkdir(parents=True, exist_ok=True)
-    dest = settings.input_dir / safe_name
     content = await file.read()
+    digest = hashlib.sha256(content).hexdigest()
+    dest = settings.input_dir / f"{digest[:16]}__{safe_name}"
     dest.write_bytes(content)
     await catalog_pdf(get_db(), dest, filename=safe_name)
 
-    return {"filename": safe_name, "size": len(content), "path": str(dest)}
+    return {
+        "filename": safe_name,
+        "stored_filename": dest.name,
+        "size": len(content),
+        "path": str(dest),
+    }
 
 
 @router.get("/api/files/input", response_model=list[FileInfo])
@@ -40,13 +47,17 @@ async def list_input_files():
     if not input_dir.exists():
         return []
 
+    documents_by_path = {
+        doc["source_path"]: doc for doc in await get_db().list_documents(limit=1000)
+    }
     files = []
     for p in sorted(input_dir.iterdir()):
         if p.is_file() and p.suffix.lower() == ".pdf":
             stat = p.stat()
+            document = documents_by_path.get(str(p))
             files.append(
                 FileInfo(
-                    name=p.name,
+                    name=document["filename"] if document else p.name,
                     size=stat.st_size,
                     modified=stat.st_mtime,
                     path=str(p),

ocr_pipeline/services/db.py

@@ -431,6 +431,40 @@ async def update_document_metadata(
             raise KeyError(document_id)
         return doc
 
+    async def update_documents_metadata(
+        self, document_ids: list[str], **fields: Any
+    ) -> list[dict[str, Any]]:
+        if not document_ids:
+            return []
+        clean = {k: v for k, v in fields.items() if k in DOCUMENT_METADATA_FIELDS}
+        placeholders = ", ".join("?" for _ in document_ids)
+        async with self.conn.execute(
+            f"SELECT id FROM documents WHERE id IN ({placeholders})",
+            document_ids,
+        ) as cur:
+            existing = {row["id"] for row in await cur.fetchall()}
+        missing = [
+            document_id for document_id in document_ids if document_id not in existing
+        ]
+        if missing:
+            raise KeyError(missing[0])
+        if clean:
+            clean["catalog_updated_at"] = _now()
+            cols = ", ".join(f"{k} = ?" for k in clean)
+            values = [*clean.values()]
+            for document_id in document_ids:
+                await self.conn.execute(
+                    f"UPDATE documents SET {cols} WHERE id = ?",
+                    [*values, document_id],
+                )
+            await self.conn.commit()
+        docs = []
+        for document_id in document_ids:
+            doc = await self.get_document(document_id)
+            if doc:
+                docs.append(doc)
+        return docs
+
     async def list_document_runs(self, document_id: str) -> list[dict[str, Any]]:
         async with self.conn.execute(
             """

ocr_pipeline/services/ocr_pair_exporter.py

@@ -1,6 +1,7 @@
 import hashlib
 import json
 import shutil
+import tempfile
 import zipfile
 from dataclasses import dataclass
 from pathlib import Path
@@ -45,6 +46,8 @@ def _split_name(stable_key: str) -> str:
 
     @staticmethod
     def _json_list(raw: str | None) -> list[str]:
+        if isinstance(raw, list):
+            return [str(item) for item in raw]
         if not raw:
             return []
         try:
@@ -68,15 +71,19 @@ def export_run(
         documents: list[dict],
         pages_by_document: dict[str, list[dict]],
         catalog_by_document: dict[str, dict],
+        document_ids: set[str] | None = None,
         dpi: int = 160,
         text_mode: str = "clean",
     ) -> OCRPairExportResult:
         if text_mode not in {"clean", "raw"}:
             raise ValueError("text_mode must be clean or raw")
 
-        if self.export_dir.exists():
-            shutil.rmtree(self.export_dir)
-        images_dir = self.export_dir / "images"
+        tmp_parent = self.export_dir.parent
+        tmp_parent.mkdir(parents=True, exist_ok=True)
+        tmp_path = Path(
+            tempfile.mkdtemp(prefix=f"{self.export_dir.name}.", dir=tmp_parent)
+        )
+        images_dir = tmp_path / "images"
         images_dir.mkdir(parents=True, exist_ok=True)
 
         split_rows: dict[str, list[dict]] = {"train": [], "validation": [], "test": []}
@@ -86,6 +93,8 @@ def export_run(
             if doc.get("status") != "completed":
                 continue
             document_id = doc["document_id"]
+            if document_ids is not None and document_id not in document_ids:
+                continue
             catalog = catalog_by_document.get(document_id, {})
             pdf_path_str = doc.get("artifact_source_pdf") or doc.get(
                 "document_source_path"
@@ -121,16 +130,24 @@ def export_run(
                 raw_text = raw_pages[page_num - 1]
                 clean_text = clean_pages[page_num - 1]
                 text = clean_text if text_mode == "clean" else raw_text
-                split = self._split_name(page_id)
+                split_key = doc.get("file_sha256") or document_id
+                split = self._split_name(split_key)
+                image_path = images_dir / f"{page_id}.png"
+                image_hash = hashlib.sha256(image_path.read_bytes()).hexdigest()
 
                 split_rows[split].append(
                     {
                         "id": page_id,
+                        "run_id": run["id"],
                         "image": image_rel,
                         "text": text,
                         "raw_text": raw_text,
                         "clean_text": clean_text,
                         "text_mode": text_mode,
+                        "label_source": "cleaned_machine_ocr"
+                        if text_mode == "clean"
+                        else "machine_ocr",
+                        "review_status": "unreviewed",
                         "document_id": document_id,
                         "document_name": doc.get("document_filename"),
                         "page": page_num,
@@ -155,6 +172,10 @@ def export_run(
                         "extraction_mode": page_meta.get("extraction_mode"),
                         "extraction_attempt": page_meta.get("extraction_attempt"),
                         "dpi_used": page_meta.get("dpi_used"),
+                        "render_dpi": dpi,
+                        "image_width": image.width,
+                        "image_height": image.height,
+                        "image_sha256": image_hash,
                         "source_file": doc.get("document_filename"),
                         "source_pdf_sha256": doc.get("file_sha256"),
                         "ocr_model": run.get("model_used"),
@@ -163,8 +184,11 @@ def export_run(
                 )
                 pages_count += 1
 
-        self._write_jsonl(split_rows)
-        self._write_manifest(run, pages_count, dpi, text_mode)
+        self._write_jsonl(tmp_path, split_rows)
+        self._write_manifest(tmp_path, run, pages_count, dpi, text_mode)
+        if self.export_dir.exists():
+            shutil.rmtree(self.export_dir)
+        tmp_path.replace(self.export_dir)
         bundle = self.export_dir.with_suffix(".zip")
         if bundle.exists():
             bundle.unlink()
@@ -181,16 +205,16 @@ def _read_text(path_str: str | None) -> str:
         path = Path(path_str)
         return path.read_text(encoding="utf-8") if path.exists() else ""
 
-    def _write_jsonl(self, split_rows: dict[str, list[dict]]) -> None:
+    def _write_jsonl(self, export_dir: Path, split_rows: dict[str, list[dict]]) -> None:
         for split, rows in split_rows.items():
-            path = self.export_dir / f"{split}.jsonl"
+            path = export_dir / f"{split}.jsonl"
             path.write_text(
                 "".join(json.dumps(row, ensure_ascii=False) + "\n" for row in rows),
                 encoding="utf-8",
             )
 
     def _write_manifest(
-        self, run: dict, pages_count: int, dpi: int, text_mode: str
+        self, export_dir: Path, run: dict, pages_count: int, dpi: int, text_mode: str
     ) -> None:
         payload = {
             "export_type": "ocr_pairs",
@@ -200,13 +224,23 @@ def _write_manifest(
             "image_format": "png",
             "dpi": dpi,
             "text_mode": text_mode,
+            "dataset_purpose": "ocr_audit",
+            "label_source": "cleaned_machine_ocr"
+            if text_mode == "clean"
+            else "machine_ocr",
+            "review_status": "unreviewed",
             "schema_version": 1,
+            "split_strategy": {
+                "method": "sha256_bucket",
+                "key": "source_pdf_sha256",
+                "ratios": {"train": 0.90, "validation": 0.05, "test": 0.05},
+            },
             "ocr_model": run.get("model_used") or settings.model_name,
             "pipeline_version": run.get("pipeline_version")
             or settings.pipeline_version,
             "splits": ["train", "validation", "test"],
         }
-        (self.export_dir / "manifest.json").write_text(
+        (export_dir / "manifest.json").write_text(
             json.dumps(payload, indent=2, ensure_ascii=False),
             encoding="utf-8",
         )

ocr_pipeline/static/index.html

@@ -170,7 +170,9 @@ <h2>Documents</h2>
                     <span class="pill pill-sm"
                           :class="doc.metadata_complete ? 'pill-success' : 'pill-warn'"
                           x-text="doc.metadata_complete ? 'ready' : 'missing'"></span>
-                    <span x-text="doc.latest_run_status || 'never'"></span>
+                    <span class="pill pill-sm"
+                          :class="documentProcessClass(doc)"
+                          x-text="documentProcessLabel(doc)"></span>
                   </div>
                 </template>
               </div>
@@ -258,7 +260,7 @@ <h2>Run <code x-text="selectedRun?.id"></code></h2>
           ↓ Dataset bundle (zip)
         </button>
         <button class="btn btn-ghost" @click="downloadOCRPairs">
-          ↓ OCR pairs (zip)
+          <span x-text="selectedRunDocumentIds.length ? `↓ OCR pairs (${selectedRunDocumentIds.length})` : '↓ OCR pairs (zip)'"></span>
         </button>
         <button class="btn btn-primary"
                 @click="openHFModal"
@@ -280,7 +282,12 @@ <h3 class="section-title">Documents</h3>
                 :class="{ active: inspector.documentId === doc.document_id }"
                 @click="openDocument(doc.document_id)">
               <div class="doc-head">
-                <span class="doc-name" x-text="doc.filename"></span>
+                <label class="checkbox-row" @click.stop>
+                  <input type="checkbox"
+                         :checked="selectedRunDocumentIds.includes(doc.document_id)"
+                         @change="toggleRunDocument(doc.document_id)">
+                  <span class="doc-name" x-text="doc.filename"></span>
+                </label>
                 <span class="pill pill-sm"
                       :class="runStatusClass(doc.status)"
                       x-text="doc.status"></span>

ocr_pipeline/static/js/api.js

@@ -143,8 +143,10 @@ const API = {
     return `/api/runs/${encodeURIComponent(runId)}/dataset/download`;
   },
 
-  ocrPairsDownloadUrl(runId, { dpi = 160, textMode = 'clean' } = {}) {
-    return `/api/runs/${encodeURIComponent(runId)}/ocr-pairs/download?dpi=${dpi}&text_mode=${encodeURIComponent(textMode)}`;
+  ocrPairsDownloadUrl(runId, { dpi = 160, textMode = 'clean', documentIds = [] } = {}) {
+    const params = new URLSearchParams({ dpi: String(dpi), text_mode: textMode });
+    if (documentIds.length > 0) params.set('document_ids', documentIds.join(','));
+    return `/api/runs/${encodeURIComponent(runId)}/ocr-pairs/download?${params.toString()}`;
   },
 
   async publishToHF(runId, payload) {