Spaces:
Running
Running
File size: 4,837 Bytes
79b0bef | 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 | """
scripts/run_ner_batch.py
────────────────────────────────────────────────────────────────
Batch NER runner — the missing link between the ETL loader and
the entities table. Reads stored clinical notes from the
database (Supabase Postgres or local SQLite, whichever
DATABASE_URL points at), runs the configured NER pipeline
(src/nlp/ner.py) over notes that don't have entities yet, and
writes the results back via load_entities() (src/etl/load.py).
Idempotent / resumable: a note is only processed if it has zero
rows in the entities table, so interrupting the run (Ctrl+C,
crash, overnight power loss) and re-running later just picks up
where it left off — no duplicate entities.
Usage
─────
python scripts/run_ner_batch.py # all pending notes
python scripts/run_ner_batch.py --limit 250 # quick demo subset
python scripts/run_ner_batch.py --batch-size 16 # smaller spaCy batches
────────────────────────────────────────────────────────────────
"""
from __future__ import annotations
import argparse
import time
from sqlalchemy import select
from src.db.connection import get_session
from src.db.models import ClinicalNote, Entity as EntityRow
from src.etl.load import load_entities
from src.nlp.ner import build_ner_pipeline
from src.utils.logger import get_logger
logger = get_logger(__name__)
def _fetch_pending_notes(limit: int | None) -> list[tuple[int, str]]:
"""Return (id, transcription) pairs for notes with no entities yet.
Args:
limit: Maximum number of notes to return, or None for all
pending notes.
Returns:
List of (note_id, transcription) tuples, ordered by note id
so a --limit run always covers the same notes on re-run.
"""
with get_session() as session:
stmt = (
select(ClinicalNote.id, ClinicalNote.transcription)
.outerjoin(EntityRow, EntityRow.note_id == ClinicalNote.id)
.where(EntityRow.id.is_(None))
.order_by(ClinicalNote.id)
)
if limit:
stmt = stmt.limit(limit)
return [(row.id, row.transcription) for row in session.execute(stmt)]
def run(limit: int | None, batch_size: int) -> None:
"""Run NER over pending notes and save the extracted entities.
Args:
limit: Maximum number of pending notes to process, or None
for every note that doesn't have entities yet.
batch_size: Number of notes per spaCy nlp.pipe() batch, and
per load_entities() write.
"""
notes = _fetch_pending_notes(limit)
total = len(notes)
if total == 0:
logger.info("No pending notes — every note already has entities.")
return
logger.info("Running NER on %d notes (batch_size=%d)...", total, batch_size)
pipeline = build_ner_pipeline()
start = time.time()
processed = 0
total_entities = 0
for chunk_start in range(0, total, batch_size):
chunk = notes[chunk_start : chunk_start + batch_size]
ids = [note_id for note_id, _ in chunk]
texts = [text or "" for _, text in chunk]
batches = pipeline.extract_batch(texts, batch_size=batch_size)
entity_dicts = []
for note_id, entities in zip(ids, batches):
for ent in entities:
d = ent.to_dict()
d["note_id"] = note_id
entity_dicts.append(d)
total_entities += load_entities(entity_dicts)
processed += len(chunk)
elapsed = time.time() - start
rate = processed / elapsed if elapsed > 0 else 0
eta_seconds = (total - processed) / rate if rate > 0 else float("inf")
logger.info(
" %d/%d notes | %d entities so far | %.0fs elapsed | ETA %.0fs",
processed, total, total_entities, elapsed, eta_seconds,
)
logger.info(
"Done: %d notes processed, %d entities inserted in %.0fs",
processed, total_entities, time.time() - start,
)
def main() -> None:
parser = argparse.ArgumentParser(
description="Run NER over stored clinical notes and save the extracted entities."
)
parser.add_argument(
"--limit", type=int, default=None,
help="Max number of pending notes to process (default: all pending notes)",
)
parser.add_argument(
"--batch-size", type=int, default=32,
help="Notes per spaCy nlp.pipe() batch (default: 32)",
)
args = parser.parse_args()
run(limit=args.limit, batch_size=args.batch_size)
if __name__ == "__main__":
main()
|