Spaces:
Sleeping
Sleeping
File size: 4,549 Bytes
d86db02 | 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 126 127 128 129 130 131 132 133 134 135 136 | import asyncio
from contextlib import asynccontextmanager
from dataclasses import asdict
from functools import partial
import easyocr
from fastapi import FastAPI, File, HTTPException, UploadFile
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import FileResponse, JSONResponse, StreamingResponse
import json
from fastapi.staticfiles import StaticFiles
from pathlib import Path
from google import genai
import os
from dotenv import load_dotenv
from src.classifier import classify
load_dotenv()
ALLOWED_MIME_TYPES = {"image/png", "image/jpeg", "image/jpg", "image/webp"}
# Cap how many documents are classified concurrently per request.
# Tuned for a 4 GiB replica: each easyocr + Gemini call can spike to ~400-600 MB,
# so >5 in flight risks OOM. Override with CLASSIFY_CONCURRENCY env var.
MAX_CONCURRENT_CLASSIFICATIONS = int(os.environ.get("CLASSIFY_CONCURRENCY", "5"))
# ── metrics store (in-memory, reset on restart) ───────────────────────────────
_metrics: dict = {
"total_requests": 0,
"total_documents": 0,
"total_input_tokens": 0,
"total_output_tokens": 0,
"by_method": {"rules": 0, "ocr": 0, "llm": 0},
"by_doc_type": {"bill": 0, "kyc": 0, "image": 0},
}
# ── shared ML resources (loaded once at startup) ──────────────────────────────
_ocr_reader: easyocr.Reader | None = None
_llm_client: genai.Client | None = None
@asynccontextmanager
async def lifespan(app: FastAPI):
global _ocr_reader, _llm_client
_ocr_reader = easyocr.Reader(["en"], gpu=False)
_llm_client = genai.Client(api_key=os.environ["GOOGLE_API_KEY"])
yield
# cleanup (nothing needed for these clients)
app = FastAPI(
title="MediShield Document Classifier",
description="Classifies scanned insurance documents using rules, OCR, and Gemini LLM.",
version="1.0.0",
lifespan=lifespan,
)
app.add_middleware(
CORSMiddleware,
allow_origins=["*"],
allow_methods=["*"],
allow_headers=["*"],
)
FRONTEND_DIR = Path(__file__).resolve().parent.parent / "frontend"
if FRONTEND_DIR.is_dir():
@app.get("/", include_in_schema=False)
def index():
return FileResponse(FRONTEND_DIR / "index.html")
app.mount("/static", StaticFiles(directory=FRONTEND_DIR), name="static")
@app.get("/health")
def health():
return {"status": "ok"}
@app.get("/metrics")
def metrics():
return _metrics
@app.post("/classify")
async def classify_documents(files: list[UploadFile] = File(...)):
if not files:
raise HTTPException(status_code=422, detail="At least one file is required.")
for f in files:
if f.content_type not in ALLOWED_MIME_TYPES:
raise HTTPException(
status_code=422,
detail=f"Unsupported file type '{f.content_type}' for '{f.filename}'. "
f"Accepted: {', '.join(ALLOWED_MIME_TYPES)}",
)
_metrics["total_requests"] += 1
_metrics["total_documents"] += len(files)
# Read all file bytes concurrently first
contents = await asyncio.gather(*[f.read() for f in files])
loop = asyncio.get_running_loop()
sem = asyncio.Semaphore(MAX_CONCURRENT_CLASSIFICATIONS)
async def _classify_one(filename: str, image_bytes: bytes):
# Semaphore caps in-flight classifications so a 50-doc upload doesn't OOM
# the worker. Excess docs queue here until a slot frees up.
async with sem:
fn = partial(
classify,
filename=filename,
image_bytes=image_bytes,
ocr_reader=_ocr_reader,
llm_client=_llm_client,
)
return await loop.run_in_executor(None, fn)
tasks = [
asyncio.create_task(_classify_one(f.filename or "unknown.png", data))
for f, data in zip(files, contents)
]
async def stream():
for coro in asyncio.as_completed(tasks):
result = await coro
_metrics["total_input_tokens"] += result.input_tokens
_metrics["total_output_tokens"] += result.output_tokens
_metrics["by_method"][result.method] = _metrics["by_method"].get(result.method, 0) + 1
_metrics["by_doc_type"][result.doc_type] = _metrics["by_doc_type"].get(result.doc_type, 0) + 1
yield json.dumps(asdict(result)) + "\n"
return StreamingResponse(stream(), media_type="application/x-ndjson")
|