Skip to content

Commit 604efda

Browse files
ana-danieleana-daniele
committed
style: apply ruff formatting to example_07
1 parent be822be commit 604efda

1 file changed

Lines changed: 108 additions & 62 deletions

File tree

‎examples/example_07_page_entity_extraction.py‎

Lines changed: 108 additions & 62 deletions
Original file line numberDiff line numberDiff line change
@@ -49,18 +49,18 @@
4949
# ---------------------------------------------------------------------------
5050
# CONFIG — edit these to match your environment
5151
# ---------------------------------------------------------------------------
52-
PAPERS_DIR = Path("./papers") # directory containing input PDFs
53-
RUNS_DIR = Path("./runs") # output root
54-
CACHE_DIR = Path("./docling-cache") # cached Docling conversions
55-
LMS_BIN = Path.home() / ".lmstudio/bin/lms" # LM Studio CLI (optional)
56-
LMS_URL = "http://localhost:1234/v1"
52+
PAPERS_DIR = Path("./papers") # directory containing input PDFs
53+
RUNS_DIR = Path("./runs") # output root
54+
CACHE_DIR = Path("./docling-cache") # cached Docling conversions
55+
LMS_BIN = Path.home() / ".lmstudio/bin/lms" # LM Studio CLI (optional)
56+
LMS_URL = "http://localhost:1234/v1"
5757
CONTEXT_LEN = 32768
5858

5959
# Models to compare — (LM Studio model ID, short slug for filenames)
6060
MODELS = [
6161
("openai/gpt-oss-20b", "gpt-oss-20b"),
62-
("granite-4.1-8b", "granite-4.1-8b"),
63-
("nuextract3", "nuextract3"),
62+
("granite-4.1-8b", "granite-4.1-8b"),
63+
("nuextract3", "nuextract3"),
6464
]
6565

6666
# Entity types and extraction prompt
@@ -83,6 +83,7 @@
8383
If no entities found, return [].
8484
Example: [{"mention":"BERT","name":"BERT","type":"MODEL"},{"mention":"F1 score","name":"F1","type":"KPI"}]"""
8585

86+
8687
# ---------------------------------------------------------------------------
8788
# Helpers
8889
# ---------------------------------------------------------------------------
@@ -102,7 +103,8 @@ def switch_model(model_id: str) -> None:
102103
subprocess.run([str(LMS_BIN), "unload", parts[0]], capture_output=True)
103104
subprocess.run(
104105
[str(LMS_BIN), "load", model_id, "--context-length", str(CONTEXT_LEN)],
105-
check=True, capture_output=True,
106+
check=True,
107+
capture_output=True,
106108
)
107109
log(f" Loaded {model_id} (ctx={CONTEXT_LEN})")
108110

@@ -144,7 +146,7 @@ def call_model(client: OpenAI, model_id: str, page_text: str) -> str:
144146
model=model_id,
145147
messages=[
146148
{"role": "system", "content": SYSTEM_PROMPT},
147-
{"role": "user", "content": page_text},
149+
{"role": "user", "content": page_text},
148150
],
149151
temperature=0,
150152
timeout=120,
@@ -169,8 +171,8 @@ def parse_entities(raw: str) -> list[dict]:
169171
if not isinstance(e, dict):
170172
continue
171173
mention = str(e.get("mention") or e.get("text") or "").strip()
172-
name = str(e.get("name") or e.get("canonical") or mention).strip()
173-
etype = str(e.get("type") or e.get("label") or "").strip().upper()
174+
name = str(e.get("name") or e.get("canonical") or mention).strip()
175+
etype = str(e.get("type") or e.get("label") or "").strip().upper()
174176
if mention and etype in ENTITY_TYPES:
175177
out.append({"mention": mention, "name": name, "type": etype})
176178
return out
@@ -258,11 +260,7 @@ def run_model(
258260

259261
log(f" {stem}")
260262
doc = convert_pdf(pdf)
261-
page_numbers = sorted({
262-
prov.page_no
263-
for item, _ in doc.iterate_items()
264-
for prov in getattr(item, "prov", [])
265-
})
263+
page_numbers = sorted({prov.page_no for item, _ in doc.iterate_items() for prov in getattr(item, "prov", [])})
266264
log(f" {len(page_numbers)} pages")
267265

268266
pages: dict[int, list[dict]] = {}
@@ -275,8 +273,7 @@ def run_model(
275273

276274
if not page_text.strip():
277275
pages[page_no] = []
278-
_write_empty_row(csv_writer, stem, page_no, model_slug,
279-
"page skipped — no serializable content", 0.0)
276+
_write_empty_row(csv_writer, stem, page_no, model_slug, "page skipped — no serializable content", 0.0)
280277
continue
281278

282279
try:
@@ -300,13 +297,19 @@ def run_model(
300297
elapsed = round(time.time() - t0, 1)
301298
log(f" → {total_entities} entities in {elapsed}s ({hallucinated} unverified)")
302299

303-
out_json.write_text(json.dumps({
304-
"document": stem, "model": model_slug,
305-
"pages": {str(k): v for k, v in pages.items()},
306-
"total_entities": total_entities,
307-
"hallucinated": hallucinated,
308-
"processing_time_s": elapsed,
309-
}, indent=2))
300+
out_json.write_text(
301+
json.dumps(
302+
{
303+
"document": stem,
304+
"model": model_slug,
305+
"pages": {str(k): v for k, v in pages.items()},
306+
"total_entities": total_entities,
307+
"hallucinated": hallucinated,
308+
"processing_time_s": elapsed,
309+
},
310+
indent=2,
311+
)
312+
)
310313
out_html.write_text(build_html(stem, model_slug, pages))
311314
with open(run_dir / "run.log", "a") as f:
312315
f.write(f"[{datetime.now()}] {stem}: {total_entities} entities, {elapsed}s\n")
@@ -315,9 +318,19 @@ def run_model(
315318

316319

317320
def _write_empty_row(w, doc, page, model, note, t):
318-
w.writerow({"document": doc, "page": page, "model": model,
319-
"entity_name": "", "entity_type": "", "entity_mentions": "",
320-
"verified": "", "note": note, "processing_time_s": t})
321+
w.writerow(
322+
{
323+
"document": doc,
324+
"page": page,
325+
"model": model,
326+
"entity_name": "",
327+
"entity_type": "",
328+
"entity_mentions": "",
329+
"verified": "",
330+
"note": note,
331+
"processing_time_s": t,
332+
}
333+
)
321334

322335

323336
def _write_page_rows(w, doc, page_no, model, entities, note, page_time):
@@ -331,13 +344,19 @@ def _write_page_rows(w, doc, page_no, model, entities, note, page_time):
331344
grouped.setdefault(k, []).append(e["mention"])
332345
vmap[k] = vmap.get(k, False) or e.get("verified", True)
333346
for (name, etype), mentions in grouped.items():
334-
w.writerow({
335-
"document": doc, "page": page_no, "model": model,
336-
"entity_name": name, "entity_type": etype,
337-
"entity_mentions": "; ".join(dict.fromkeys(mentions)),
338-
"verified": vmap[(name, etype)], "note": note,
339-
"processing_time_s": page_time,
340-
})
347+
w.writerow(
348+
{
349+
"document": doc,
350+
"page": page_no,
351+
"model": model,
352+
"entity_name": name,
353+
"entity_type": etype,
354+
"entity_mentions": "; ".join(dict.fromkeys(mentions)),
355+
"verified": vmap[(name, etype)],
356+
"note": note,
357+
"processing_time_s": page_time,
358+
}
359+
)
341360

342361

343362
# ---------------------------------------------------------------------------
@@ -348,9 +367,14 @@ def build_complete_csv(raw_csv: Path, out_csv: Path) -> None:
348367
with open(raw_csv, newline="", encoding="utf-8") as f:
349368
rows = list(csv.DictReader(f))
350369

351-
stats: dict[tuple, dict] = defaultdict(lambda: {
352-
"total_entities": 0, "unverified": 0, "total_time_s": 0.0, "_pages": set(),
353-
})
370+
stats: dict[tuple, dict] = defaultdict(
371+
lambda: {
372+
"total_entities": 0,
373+
"unverified": 0,
374+
"total_time_s": 0.0,
375+
"_pages": set(),
376+
}
377+
)
354378
for r in rows:
355379
s = stats[(r["document"], r["model"])]
356380
if r["page"] not in s["_pages"]:
@@ -362,11 +386,19 @@ def build_complete_csv(raw_csv: Path, out_csv: Path) -> None:
362386
s["unverified"] += 1
363387

364388
fields = [
365-
"document", "page", "model",
366-
"entity_name", "entity_type", "entity_mentions",
367-
"hallucinated", "note", "page_processing_time_s",
368-
"doc_model_total_entities", "doc_model_hallucinations",
369-
"doc_model_hallucination_rate_pct", "doc_model_total_time_s",
389+
"document",
390+
"page",
391+
"model",
392+
"entity_name",
393+
"entity_type",
394+
"entity_mentions",
395+
"hallucinated",
396+
"note",
397+
"page_processing_time_s",
398+
"doc_model_total_entities",
399+
"doc_model_hallucinations",
400+
"doc_model_hallucination_rate_pct",
401+
"doc_model_total_time_s",
370402
]
371403
with open(out_csv, "w", newline="", encoding="utf-8") as f:
372404
w = csv.DictWriter(f, fieldnames=fields)
@@ -375,18 +407,23 @@ def build_complete_csv(raw_csv: Path, out_csv: Path) -> None:
375407
s = stats[(r["document"], r["model"])]
376408
total, unver = s["total_entities"], s["unverified"]
377409
v = str(r["verified"]).lower()
378-
w.writerow({
379-
"document": r["document"], "page": r["page"], "model": r["model"],
380-
"entity_name": r["entity_name"], "entity_type": r["entity_type"],
381-
"entity_mentions": r["entity_mentions"],
382-
"hallucinated": "yes" if v == "false" else "no" if v == "true" else "",
383-
"note": r["note"],
384-
"page_processing_time_s": r["processing_time_s"],
385-
"doc_model_total_entities": total,
386-
"doc_model_hallucinations": unver,
387-
"doc_model_hallucination_rate_pct": round(100 * unver / total, 1) if total else 0.0,
388-
"doc_model_total_time_s": round(s["total_time_s"], 1),
389-
})
410+
w.writerow(
411+
{
412+
"document": r["document"],
413+
"page": r["page"],
414+
"model": r["model"],
415+
"entity_name": r["entity_name"],
416+
"entity_type": r["entity_type"],
417+
"entity_mentions": r["entity_mentions"],
418+
"hallucinated": "yes" if v == "false" else "no" if v == "true" else "",
419+
"note": r["note"],
420+
"page_processing_time_s": r["processing_time_s"],
421+
"doc_model_total_entities": total,
422+
"doc_model_hallucinations": unver,
423+
"doc_model_hallucination_rate_pct": round(100 * unver / total, 1) if total else 0.0,
424+
"doc_model_total_time_s": round(s["total_time_s"], 1),
425+
}
426+
)
390427
log(f"Complete CSV → {out_csv} ({len(rows)} rows)")
391428

392429

@@ -396,13 +433,13 @@ def build_complete_csv(raw_csv: Path, out_csv: Path) -> None:
396433
def main() -> None:
397434
parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
398435
parser.add_argument("--papers", type=Path, default=PAPERS_DIR, help="Directory of input PDFs")
399-
parser.add_argument("--out", type=Path, default=RUNS_DIR, help="Output root directory")
400-
parser.add_argument("--url", default=LMS_URL, help="LM Studio base URL")
401-
parser.add_argument("--test", action="store_true", help="Run on first PDF only")
436+
parser.add_argument("--out", type=Path, default=RUNS_DIR, help="Output root directory")
437+
parser.add_argument("--url", default=LMS_URL, help="LM Studio base URL")
438+
parser.add_argument("--test", action="store_true", help="Run on first PDF only")
402439
args = parser.parse_args()
403440

404441
papers_dir = args.papers
405-
runs_dir = args.out
442+
runs_dir = args.out
406443
runs_dir.mkdir(parents=True, exist_ok=True)
407444

408445
pdfs = sorted(papers_dir.glob("*.pdf"))
@@ -414,11 +451,20 @@ def main() -> None:
414451
date = datetime.now().strftime("%Y-%m-%d")
415452
log(f"Papers: {len(pdfs)} | Models: {len(MODELS)}")
416453

417-
raw_csv = runs_dir / f"{date}_entities_raw.csv"
454+
raw_csv = runs_dir / f"{date}_entities_raw.csv"
418455
final_csv = runs_dir / f"{date}_entities.csv"
419456

420-
raw_fields = ["document", "page", "model", "entity_name",
421-
"entity_type", "entity_mentions", "verified", "note", "processing_time_s"]
457+
raw_fields = [
458+
"document",
459+
"page",
460+
"model",
461+
"entity_name",
462+
"entity_type",
463+
"entity_mentions",
464+
"verified",
465+
"note",
466+
"processing_time_s",
467+
]
422468

423469
with open(raw_csv, "w", newline="", encoding="utf-8") as csv_file:
424470
writer = csv.DictWriter(csv_file, fieldnames=raw_fields)

0 commit comments

Comments
 (0)