Spaces:
Running
Running
File size: 5,012 Bytes
54437b9 26d4279 54437b9 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 | """Shared singletons and the background ingestion job runner.
Heavy objects (TrOCR pipeline, embedder) load lazily on first use. Ingestion runs
on a single-worker thread pool so CPU-bound OCR never blocks the event loop and
free-CPU pages process serially; the frontend polls GET /jobs/{id} for progress.
"""
from __future__ import annotations
import json
import logging
import threading
import traceback
from concurrent.futures import ThreadPoolExecutor
from ..config import htr as htr_cfg
from ..htr.pipeline import HTRPipeline
from ..rag.index import RagIndex
from ..schemas import JobStatus, Stage
from ..store import DocumentStore
logger = logging.getLogger("inkference.ingest")
# Map a pipeline Stage -> the job status shown in the queue UI.
_STAGE_STATUS = {
Stage.SEGMENTATION: JobStatus.SEGMENTING,
Stage.RECOGNITION: JobStatus.RECOGNIZING,
Stage.CONFIDENCE: JobStatus.SCORING,
Stage.CORRECTION: JobStatus.CORRECTING,
}
_store: DocumentStore | None = None
_pipeline: HTRPipeline | None = None
_index: RagIndex | None = None
_pipeline_lock = threading.Lock()
_executor = ThreadPoolExecutor(max_workers=1, thread_name_prefix="ingest")
def get_store() -> DocumentStore:
global _store
if _store is None:
_store = DocumentStore()
return _store
def get_index() -> RagIndex:
global _index
if _index is None:
_index = RagIndex()
return _index
def get_pipeline() -> HTRPipeline:
global _pipeline
if _pipeline is None:
with _pipeline_lock:
if _pipeline is None:
_pipeline = HTRPipeline(htr_cfg)
return _pipeline
def run_agent(doc_id: int, question: str, **kwargs):
"""One turn of the LangGraph research agent (POST /documents/{id}/agent).
Imported lazily: langgraph is only needed by this endpoint, so the seeders and
HTR scripts keep booting without it."""
from ..agent.runner import run_agent as _run
return _run(doc_id, question, get_store(), get_index(), **kwargs)
def submit_ingest(doc_id: int, page_specs: list[tuple[int, int, str]], job_id: int) -> None:
"""page_specs = [(page_id, page_number, image_path), ...]"""
_executor.submit(_run_ingest, doc_id, page_specs, job_id)
def _run_ingest(doc_id: int, page_specs: list[tuple[int, int, str]], job_id: int) -> None:
store = get_store()
total = len(page_specs)
logger.info("ingest job %s: doc=%s, %d page(s) queued", job_id, doc_id, total)
store.update_job(job_id, status=JobStatus.QUEUED, total_pages=total, done_pages=0)
try:
pipeline = get_pipeline()
for done, (page_id, page_number, image_path) in enumerate(page_specs):
def progress(stage: Stage, frac: float, msg: str, _done=done, _pn=page_number) -> None:
overall = (_done + frac) / total
logger.debug("job %s page %s: %s %.0f%% — %s",
job_id, _pn, stage.value, frac * 100, msg)
store.update_job(
job_id,
status=_STAGE_STATUS.get(stage, JobStatus.RECOGNIZING),
stage=stage,
progress=round(overall, 4),
message=f"Page {page_number} — {msg}",
)
def on_segmented(bboxes, w, h, _pn=page_number) -> None:
# Surface line boxes the instant segmentation finishes (before the slow
# recognition/correction) so the UI can overlay them right away.
store.update_job(job_id, seg_preview=json.dumps(
{"page_number": _pn, "width": w, "height": h, "boxes": bboxes}))
logger.info("job %s: processing page %s (%d/%d)", job_id, page_number, done + 1, total)
store.set_page_status(page_id, "processing")
result = pipeline.process_path(image_path, page_number, progress, on_segmented)
store.save_page_result(page_id, result)
store.update_job(job_id, done_pages=done + 1)
logger.info("job %s: page %s done — %d lines, avg conf %.2f",
job_id, page_number, len(result.lines), result.avg_confidence)
# Incrementally index only the newly-added pages (a full rebuild re-embeds the
# whole corpus, ~90s for 900 pages, and needlessly delays "Complete").
page_numbers = [pn for (_pid, pn, _path) in page_specs]
n_chunks = get_index().add_pages(doc_id, store, page_numbers)
store.update_job(
job_id, status=JobStatus.COMPLETE, progress=1.0, message="Complete"
)
logger.info("ingest job %s complete; RAG index updated (+%d page(s), %s chunks)",
job_id, len(page_numbers), n_chunks)
except Exception as exc: # surface failure to the job poller
logger.exception("ingest job %s FAILED: %s", job_id, exc)
store.update_job(
job_id, status=JobStatus.FAILED,
error=f"{exc}\n{traceback.format_exc()}", message=str(exc),
)
|