Spaces:
Sleeping
Sleeping
File size: 2,522 Bytes
13779e8 | 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 | """Re-extract just the still-failing PDFs (those with _raw.txt but no .json).
Uses asyncio + Semaphore(3) — 3 concurrent NIM calls at a time. NIM's
40 req/min budget = ~6 concurrent in-flight comfortably; 3 leaves headroom
for the eval/sweep that might run alongside.
Reuses extract_one from rag.extract so behaviour matches the canonical
extraction path. Logs OK/FAIL per PDF + final tally.
"""
import asyncio, sys, time
from pathlib import Path
sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
from rag.extract import (
extract_one, find_pdfs, load_manifest, init_db,
EXTRACTED_DIR, ROOT,
)
from rag.ingest import policy_id_for
from backend.providers.nvidia_nim_llm import get_brain_llm
from backend.providers.sarvam_llm import SarvamLLM
async def run_one(pdf, manifest, primary, fallback, sem, idx, total):
async with sem:
rel = str(pdf.relative_to(ROOT))
entry = manifest.get(rel, {})
t0 = time.time()
print(f"[{idx}/{total}] {pdf.parent.name}__{pdf.stem[:50]}")
try:
result = await asyncio.wait_for(
extract_one(pdf, entry, primary, fallback),
timeout=180, # 3 min hard cap per PDF
)
elapsed = time.time() - t0
return ("OK" if result else "FAIL", pdf, elapsed)
except Exception as e:
return (f"FAIL ({type(e).__name__})", pdf, time.time() - t0)
async def main():
init_db()
pdfs = find_pdfs()
manifest = load_manifest()
# Filter to those with _raw.txt but no .json (the failing set)
targets = []
for pdf in pdfs:
pid = policy_id_for(pdf)
json_path = EXTRACTED_DIR / f"{pid}.json"
raw_path = EXTRACTED_DIR / f"{pid}._raw.txt"
if not json_path.exists() and raw_path.exists():
targets.append(pdf)
print(f"Re-extracting {len(targets)} failed PDFs (concurrency=3)...\n")
primary = get_brain_llm()
fallback = SarvamLLM()
sem = asyncio.Semaphore(3)
tasks = [
run_one(pdf, manifest, primary, fallback, sem, i+1, len(targets))
for i, pdf in enumerate(targets)
]
results = await asyncio.gather(*tasks)
ok = sum(1 for r in results if r[0] == "OK")
print(f"\n=== Final: {ok}/{len(targets)} OK ===")
for status, pdf, elapsed in results:
if not status.startswith("OK"):
print(f" FAIL | {pdf.parent.name}__{pdf.stem[:60]} | {status} ({elapsed:.0f}s)")
if __name__ == "__main__":
asyncio.run(main())
|