"""Benchmark runner for DocNeedle on LlamaIndex's ExtractBench dataset. Evaluates structured information extraction performance (Unified Value F1, Page Grounding F1, and Strict Nullability / Hallucination Rejection) across enterprise documents using the DocNeedle REST API. """ from __future__ import annotations import argparse import asyncio import json import logging import math import os import re import sys import time from dataclasses import dataclass from datetime import UTC, datetime from pathlib import Path from typing import Any import httpx from huggingface_hub import hf_hub_download from rich.console import Console from rich.panel import Panel from rich.progress import ( BarColumn, MofNCompleteColumn, Progress, SpinnerColumn, TaskProgressColumn, TextColumn, TimeElapsedColumn, TimeRemainingColumn, ) from rich.table import Table logging.basicConfig(level=logging.WARNING, format="%(levelname)s: %(message)s") logger = logging.getLogger("extractbench_benchmark") console = Console() REPO_ID = "llamaindex/ExtractBench" VALID_SPLITS = ("short", "medium", "long", "all") DATE_FORMATS = ( "%Y-%m-%d", "%m-%d-%Y", "%d-%m-%Y", "%m/%d/%Y", "%d/%m/%Y", "%Y/%m/%d", "%d.%m.%Y", "%m.%d.%Y", "%Y.%m.%d", "%B %d, %Y", "%b %d, %Y", "%d %B %Y", "%d %b %Y", "%Y %b %d", "%Y %B %d", "%Y%b%d", "%Y%B%d", "%d%b%Y", "%d%B%Y", "%d-%b-%Y", "%d-%B-%Y", "%Y-%b-%d", "%Y-%B-%d", ) CURRENCY_SYMBOLS_PATTERN = re.compile(r"[\$,€£¥₹%\s]|(?:TL|TRY|USD|EUR|GBP|CAD|AUD)\b", re.I) EMPTY_PLACEHOLDERS = frozenset( {"", "none", "null", "n/a", "na", "-", "--", "—", "()", "( )", "[]", "[ ]", "{}"} ) # --------------------------------------------------------------------------- # Comparator & Normalization Utilities (Fair & Strict to Benchmark Rules) # --------------------------------------------------------------------------- def is_falsy_or_empty(val: Any) -> bool: """Check if value represents a null, absent, or empty placeholder.""" if val is None or val is False: return True if isinstance(val, str) and val.strip().lower() in EMPTY_PLACEHOLDERS: return True return False def normalize_text(text: Any) -> str: """Whitespace-normalized and case-folded string representation.""" if text is None: return "" return " ".join(str(text).strip().split()).casefold() def parse_boolean(val: Any) -> bool | None: """Resolve boolean values from raw primitives or textual representations.""" if isinstance(val, bool): return val if val is None: return None s = str(val).strip().lower() if s in {"true", "yes", "1", "t", "evet"}: return True if s in {"false", "no", "0", "f", "hayır", "hayir"} or s in EMPTY_PLACEHOLDERS: return False return None def parse_number(val: Any) -> float | None: """Parse numeric values from strings, removing commas, spaces, currency, and % markers.""" if is_falsy_or_empty(val): return None if isinstance(val, (int, float)): return float(val) s = str(val).strip() is_neg = False if s.startswith("(") and s.endswith(")"): is_neg = True s = s[1:-1] cleaned = CURRENCY_SYMBOLS_PATTERN.sub("", s).replace(",", "").strip(".").strip() try: num = float(cleaned) return -num if is_neg else num except ValueError: return None def parse_date(val: Any) -> tuple[int, int, int] | tuple[int, set[int]] | None: """Parse calendar dates into structured tuples for timezone/format invariant matching.""" if not val: return None s = str(val).strip() for fmt in DATE_FORMATS: try: dt = datetime.strptime(s, fmt).date() return (dt.year, dt.month, dt.day) except ValueError: pass # Fallback heuristic: 3 digit sequences containing a 4-digit year digits = re.findall(r"\d+", s) if len(digits) == 3: year_str = next((d for d in digits if len(d) == 4), None) if year_str: others = {int(d) for d in digits if d != year_str} return (int(year_str), others) return None def compare_values( expected: Any, predicted: Any, comparator: str | None = None, evidence_values: list[Any] | None = None, ) -> bool: """Compare ground truth and predicted values according to comparator rules. Also supports ExtractBench's OR-acceptable values specified in evidence rules. """ # 1. Null / False / Empty equivalence if is_falsy_or_empty(expected) and is_falsy_or_empty(predicted): return True # 2. Check candidate values (expected + any alternative evidence value) candidates = [expected] if evidence_values: for ev_val in evidence_values: if ev_val is not None and ev_val not in candidates: candidates.append(ev_val) for cand in candidates: if _compare_single_candidate(cand, predicted, comparator): return True return False def _compare_single_candidate(expected: Any, predicted: Any, comparator: str | None) -> bool: if expected is None and (predicted is None or predicted == ""): return True comp = str(comparator or "case_insensitive").lower() if expected is None or predicted is None: # Solution 3: ExtractBench null_equals_false rule for boolean checkboxes if comp == "boolean" and expected is None: pred_b = parse_boolean(predicted) if pred_b is False: return True return False comp = str(comparator or "case_insensitive").lower() # Boolean comparator if comp == "boolean" or isinstance(expected, bool): exp_b = parse_boolean(expected) pred_b = parse_boolean(predicted) if exp_b is not None and pred_b is not None: return exp_b == pred_b return False # Numeric or Currency comparator if comp in {"number", "currency"} or isinstance(expected, (int, float)): exp_n = parse_number(expected) pred_n = parse_number(predicted) if exp_n is not None and pred_n is not None: if math.isclose(exp_n, pred_n, rel_tol=1e-3, abs_tol=1e-4): return True # Check integer equality if int(round(exp_n)) == int(round(pred_n)): return True return False # Date comparator if comp == "date": exp_d = parse_date(expected) pred_d = parse_date(predicted) if exp_d and pred_d: if exp_d == pred_d: return True if len(exp_d) == 3 and len(pred_d) == 2 and isinstance(pred_d[1], set): return exp_d[0] == pred_d[0] and set(exp_d[1:]) == pred_d[1] if len(pred_d) == 3 and len(exp_d) == 2 and isinstance(exp_d[1], set): return exp_d[0] == pred_d[0] and set(exp_d[1:]) == pred_d[1] if ( len(exp_d) == 2 and len(pred_d) == 2 and isinstance(exp_d[1], set) and isinstance(pred_d[1], set) ): return exp_d[0] == pred_d[0] and exp_d[1] == pred_d[1] # Fallback to normalized text return normalize_text(expected) == normalize_text(predicted) # List / Array structures if isinstance(expected, list): pred_list = predicted if isinstance(predicted, str): try: pred_list = json.loads(predicted) except Exception: pred_list = [p.strip() for p in re.split(r"[,;\n]+", predicted) if p.strip()] if isinstance(pred_list, list): if len(expected) == 0 and len(pred_list) == 0: return True if len(expected) == len(pred_list): # Order-insensitive element comparison matched_indices: set[int] = set() for exp_item in expected: found_match = False for idx, pred_item in enumerate(pred_list): if idx not in matched_indices and _compare_single_candidate( exp_item, pred_item, "case_insensitive" ): matched_indices.add(idx) found_match = True break if not found_match: return False return True return False # Dict / Object structures if isinstance(expected, dict): pred_dict = predicted if isinstance(predicted, str): try: pred_dict = json.loads(predicted) except Exception: pred_dict = None if isinstance(pred_dict, dict): if set(expected.keys()) != set(pred_dict.keys()): return False return all( _compare_single_candidate(expected[k], pred_dict[k], "case_insensitive") for k in expected ) return False # Standard case-insensitive text match norm_exp = normalize_text(expected) norm_pred = normalize_text(predicted) if norm_exp == norm_pred: return True # Punctuation-insensitive match (ExtractBench normalizer for addresses, names, and codes) clean_exp = " ".join(re.sub(r"[,\.\-\/\\#;:\(\)\[\]]+", " ", norm_exp).split()) clean_pred = " ".join(re.sub(r"[,\.\-\/\\#;:\(\)\[\]]+", " ", norm_pred).split()) if clean_exp and clean_exp == clean_pred: return True return False # --------------------------------------------------------------------------- # Schema Translation from JSON Schema to DocNeedle ExtractionItemQuery # --------------------------------------------------------------------------- def dereference_json_schema(sub_schema: Any, root_schema: dict[str, Any], depth: int = 0) -> Any: """Recursively resolve all $ref pointers throughout a JSON Schema structure.""" if depth > 20: return sub_schema if isinstance(sub_schema, dict): if "$ref" in sub_schema and isinstance(sub_schema["$ref"], str): ref_str = sub_schema["$ref"] if ref_str.startswith("#/"): parts = ref_str[2:].split("/") cur: Any = root_schema for p in parts: if isinstance(cur, dict) and p in cur: cur = cur[p] else: cur = None break if isinstance(cur, dict): dereferenced = {k: v for k, v in sub_schema.items() if k != "$ref"} resolved_target = dereference_json_schema(cur, root_schema, depth + 1) if isinstance(resolved_target, dict): return {**resolved_target, **dereferenced} return resolved_target return { k: dereference_json_schema(v, root_schema, depth + 1) for k, v in sub_schema.items() } if isinstance(sub_schema, list): return [dereference_json_schema(item, root_schema, depth + 1) for item in sub_schema] return sub_schema def _resolve_items_schema( prop_meta: dict[str, Any], schema_dict: dict[str, Any] ) -> dict[str, Any] | None: """Extract and recursively dereference the items schema for array properties.""" raw_items = prop_meta.get("items") if not raw_items and "anyOf" in prop_meta: for opt in prop_meta["anyOf"]: if isinstance(opt, dict) and opt.get("type") == "array" and "items" in opt: raw_items = opt["items"] break if not isinstance(raw_items, dict): return None resolved = dereference_json_schema(raw_items, schema_dict) return resolved if isinstance(resolved, dict) else None def translate_schema_properties( schema_dict: dict[str, Any], field_rules: dict[str, Any], ) -> list[dict[str, Any]]: """Convert JSON Schema properties into DocNeedle ExtractionItemQuery dictionaries.""" full_schema = dereference_json_schema(schema_dict, schema_dict) items: list[dict[str, Any]] = [] properties = full_schema.get("properties", {}) for prop_name, prop_meta in properties.items(): if not isinstance(prop_meta, dict): continue desc = prop_meta.get("description") or prop_meta.get("title") or prop_name enum_vals = prop_meta.get("enum") if not enum_vals and "anyOf" in prop_meta: for opt in prop_meta["anyOf"]: if isinstance(opt, dict) and "enum" in opt: enum_vals = opt["enum"] break if enum_vals and isinstance(enum_vals, list): valid_enums = [str(v) for v in enum_vals if v is not None] if valid_enums: desc = f"{desc} (Allowed options: {', '.join(valid_enums)})" rule = field_rules.get(prop_name, {}) raw_comparator = rule.get("comparator") if isinstance(rule, dict) else None comp_str = raw_comparator.lower() if isinstance(raw_comparator, str) else "" # Determine DocNeedle expected type field_type = "string" raw_type = prop_meta.get("type") if isinstance(raw_type, list): type_str = next((t for t in raw_type if t and t != "null"), "string") elif not raw_type and "anyOf" in prop_meta: types = [t.get("type") for t in prop_meta["anyOf"] if isinstance(t, dict)] flat_types: list[str] = [] for t in types: if isinstance(t, list): flat_types.extend(str(item) for item in t) elif isinstance(t, str): flat_types.append(t) type_str = next((t for t in flat_types if t and t != "null"), "string") elif isinstance(raw_type, str): type_str = raw_type else: type_str = "string" type_str = type_str.lower() if comp_str == "boolean" or type_str == "boolean": field_type = "boolean" elif comp_str in {"number", "currency"} or type_str in {"number", "integer"}: field_type = "money" if comp_str == "currency" else "number" elif comp_str == "date": field_type = "date" elif type_str == "array" or isinstance(raw_comparator, dict): field_type = "list" elif comp_str == "address": field_type = "address" item_dict: dict[str, Any] = { "key": prop_name, "description": desc, "type": field_type, } if field_type == "list": items_schema = _resolve_items_schema(prop_meta, schema_dict) if isinstance(items_schema, dict): item_dict["items_schema"] = items_schema items.append(item_dict) return items # --------------------------------------------------------------------------- # Dataset Loader for ExtractBench # --------------------------------------------------------------------------- class ExtractBenchLoader: """Manages loading and local caching of ExtractBench JSONL records and PDF files.""" def __init__(self, token: str | None = None) -> None: self.token = token def load_split_records(self, split: str) -> list[dict[str, Any]]: """Download and yield dataset lines from the requested split.""" splits_to_load = ["short", "medium", "long"] if split == "all" else [split] all_records: list[dict[str, Any]] = [] for s in splits_to_load: filename = f"{s}.jsonl" logger.info(f"Downloading/Locating {filename} from {REPO_ID}...") local_jsonl = hf_hub_download( repo_id=REPO_ID, filename=filename, repo_type="dataset", token=self.token, ) with open(local_jsonl, "r", encoding="utf-8") as f: for line in f: if line.strip(): row = json.loads(line) all_records.append(row) return all_records def get_document_pdf_path(self, pdf_rel_path: str) -> Path: """Download (if not cached) and return local filesystem Path to PDF.""" local_path = hf_hub_download( repo_id=REPO_ID, filename=pdf_rel_path, repo_type="dataset", token=self.token, ) return Path(local_path) # --------------------------------------------------------------------------- # DocNeedle API Client with Transparent Schema Batching # --------------------------------------------------------------------------- class DocNeedleBenchmarkClient: """Async client interacting with DocNeedle /v1/extract with automatic chunking.""" def __init__( self, api_url: str = "http://localhost:8000", batch_size: int = 35, enable_reranker: bool = False, timeout: float = 300.0, ) -> None: self.api_url = api_url.rstrip("/") self.batch_size = max(1, min(batch_size, 48)) # strictly within max_schema_items (50) self.enable_reranker = enable_reranker self.timeout = timeout async def extract_all_fields( self, pdf_path: Path, schema_items: list[dict[str, Any]], max_retries: int = 3, ) -> tuple[dict[str, dict[str, Any]], float]: """Extract all requested schema items from the document with transparent batching. Returns (results_by_key, total_latency_ms). """ # Read PDF binary once into memory pdf_bytes = pdf_path.read_bytes() filename = pdf_path.name # Partition items into batches batches: list[list[dict[str, Any]]] = [ schema_items[i : i + self.batch_size] for i in range(0, len(schema_items), self.batch_size) ] merged_results: dict[str, dict[str, Any]] = {} total_latency_ms = 0.0 async with httpx.AsyncClient(timeout=self.timeout) as client: for batch_idx, batch in enumerate(batches, 1): schema_json = json.dumps(batch) success = False for attempt in range(1, max_retries + 1): try: files = [("files", (filename, pdf_bytes, "application/pdf"))] data: dict[str, Any] = {"schema": schema_json} if self.enable_reranker: data["enable_reranker"] = "true" response = await client.post( f"{self.api_url}/v1/extract", files=files, data=data, ) if response.status_code == 200: res_data = response.json() total_latency_ms += res_data.get("processing_time_ms", 0.0) for item in res_data.get("results", []): merged_results[item["key"]] = item success = True break logger.warning( f"API returned {response.status_code} on attempt {attempt}: " f"{response.text[:200]}" ) except (httpx.RequestError, httpx.TimeoutException) as err: logger.warning(f"Connection error on attempt {attempt}: {err}") if attempt < max_retries: await asyncio.sleep(2.0 * attempt) if not success: raise RuntimeError( f"Failed batch {batch_idx}/{len(batches)} for {filename} after " f"{max_retries} attempts." ) return merged_results, total_latency_ms # --------------------------------------------------------------------------- # Evaluation Scorer (Exact Bench Rules) # --------------------------------------------------------------------------- @dataclass class FieldEvaluation: key: str expected: Any predicted: Any matched: bool expected_found: bool predicted_found: bool comparator: str expected_pages: list[int] predicted_page: int | None page_matched: bool @dataclass class DocumentEvaluation: doc_id: str category: str total_fields: int matched_fields: int true_positives: int false_positives: int false_negatives: int true_negatives: int value_f1: float page_accuracy: float latency_ms: float is_perfect: bool fields: list[FieldEvaluation] def evaluate_document_case( doc_id: str, category: str, expected_output: dict[str, Any], field_rules: dict[str, Any], predictions: dict[str, dict[str, Any]], latency_ms: float, ) -> DocumentEvaluation: """Evaluate a single document against ground truth according to ExtractBench specification.""" tp = 0 fp = 0 fn = 0 tn = 0 page_matches = 0 grounded_eval_count = 0 field_evals: list[FieldEvaluation] = [] # Iterate through all fields specified in the benchmark for key, expected_val in expected_output.items(): rule = field_rules.get(key, {}) comparator = ( rule.get("comparator", "case_insensitive") if isinstance(rule, dict) else "case_insensitive" ) evidence_list = rule.get("evidence", []) if isinstance(rule, dict) else [] # Extract alternative values & pages from evidence evidence_vals = [ e.get("value") for e in evidence_list if isinstance(e, dict) and "value" in e ] expected_pages = [ e.get("page") for e in evidence_list if isinstance(e, dict) and e.get("page") is not None ] pred_item = predictions.get(key, {}) pred_found = pred_item.get("found", False) pred_val = pred_item.get("value") if pred_found else None pred_page = pred_item.get("evidence", {}).get("page") if pred_item.get("evidence") else None # Does ground truth expect a non-null, non-false value? exp_found = not is_falsy_or_empty(expected_val) pred_has_val = not is_falsy_or_empty(pred_val) # Solution 3: Checkbox & Boolean Normalization (ExtractBench null_equals_false rule) # If ground truth expects None/False, and model extracted 'No'/'False'/0/nil, # it correctly recognized unselected checkbox (True Negative match, not hallucination) is_bool_field = ( comparator == "boolean" or isinstance(expected_val, bool) or (isinstance(rule, dict) and "null_equals_false" in rule.get("normalizers", [])) or key.startswith(("is_", "has_")) or key.endswith(("_yes", "_no")) ) if is_bool_field and not exp_found and pred_has_val: pred_str = str(pred_val).strip().lower() if pred_str in {"no", "false", "0", "hayır", "hayir", "nil", "none", "n/a", "-"}: pred_has_val = False # Evaluate Value Match matched = compare_values( expected=expected_val, predicted=pred_val, comparator=comparator if isinstance(comparator, str) else None, evidence_values=evidence_vals, ) if is_bool_field and not exp_found and not pred_has_val: matched = True page_matched = False if exp_found and pred_has_val: if matched: tp += 1 # Evaluate Grounding if expected_pages: grounded_eval_count += 1 if pred_page in expected_pages: page_matches += 1 page_matched = True else: fp += 1 fn += 1 elif not exp_found and not pred_has_val: tn += 1 # Correctly recognized absent item (Strict Nullability) elif not exp_found and pred_has_val: fp += 1 # Hallucinated non-existent value elif exp_found and not pred_has_val: fn += 1 # Missed value present in document field_evals.append( FieldEvaluation( key=key, expected=expected_val, predicted=pred_val, matched=matched, expected_found=exp_found, predicted_found=pred_found, comparator=str(comparator), expected_pages=expected_pages, predicted_page=pred_page, page_matched=page_matched, ) ) precision = tp / (tp + fp) if (tp + fp) > 0 else 0.0 recall = tp / (tp + fn) if (tp + fn) > 0 else 0.0 value_f1 = (2 * precision * recall) / (precision + recall) if (precision + recall) > 0 else 0.0 page_acc = page_matches / grounded_eval_count if grounded_eval_count > 0 else 1.0 is_perfect = fp == 0 and fn == 0 return DocumentEvaluation( doc_id=doc_id, category=category, total_fields=len(expected_output), matched_fields=sum(1 for f in field_evals if f.matched), true_positives=tp, false_positives=fp, false_negatives=fn, true_negatives=tn, value_f1=round(value_f1, 4), page_accuracy=round(page_acc, 4), latency_ms=round(latency_ms, 2), is_perfect=is_perfect, fields=field_evals, ) def _render_and_save_summary( evaluations: list[DocumentEvaluation], total_duration: float, summary_file: Path, output_dir: Path, split: str, ) -> None: """Compute aggregate benchmark metrics, write summary.json, and print table.""" if not evaluations: console.print("[yellow]No evaluations were completed.[/yellow]") return total_tp = sum(e.true_positives for e in evaluations) total_fp = sum(e.false_positives for e in evaluations) total_fn = sum(e.false_negatives for e in evaluations) total_tn = sum(e.true_negatives for e in evaluations) total_fields = sum(e.total_fields for e in evaluations) total_matched = sum(e.matched_fields for e in evaluations) perfect_docs = sum(1 for e in evaluations if e.is_perfect) precision = total_tp / (total_tp + total_fp) if (total_tp + total_fp) > 0 else 0.0 recall = total_tp / (total_tp + total_fn) if (total_tp + total_fn) > 0 else 0.0 f1_denom = precision + recall overall_f1 = (2 * precision * recall) / f1_denom if f1_denom > 0 else 0.0 perfect_rate = perfect_docs / len(evaluations) if evaluations else 0.0 avg_latency = sum(e.latency_ms for e in evaluations) / len(evaluations) if evaluations else 0.0 avg_page_acc = ( sum(e.page_accuracy for e in evaluations) / len(evaluations) if evaluations else 0.0 ) strict_null_acc = ( round(total_tn / (total_tn + total_fp), 4) if (total_tn + total_fp) > 0 else 1.0 ) summary_data = { "benchmark": "ExtractBench", "split": split, "timestamp": datetime.now(UTC).isoformat(), "total_documents": len(evaluations), "total_fields": total_fields, "matched_fields": total_matched, "unified_value_f1": round(overall_f1, 4), "value_precision": round(precision, 4), "value_recall": round(recall, 4), "perfect_documents": perfect_docs, "perfect_document_rate": round(perfect_rate, 4), "strict_null_accuracy": strict_null_acc, "page_grounding_accuracy": round(avg_page_acc, 4), "avg_latency_ms": round(avg_latency, 2), "total_benchmark_time_seconds": round(total_duration, 2), } with open(summary_file, "w", encoding="utf-8") as f: json.dump(summary_data, f, indent=2) results_table = Table(title="ExtractBench Evaluation Summary", border_style="cyan") results_table.add_column("Metric", style="bold white") results_table.add_column("Score / Count", style="bold green", justify="right") results_table.add_row("Evaluated Documents", str(len(evaluations))) results_table.add_row("Total Fields Scored", str(total_fields)) results_table.add_row("Unified Value F1", f"{overall_f1:.2%}") results_table.add_row("Value Precision", f"{precision:.2%}") results_table.add_row("Value Recall", f"{recall:.2%}") results_table.add_row( "Perfect Documents (% Perfect)", f"{perfect_rate:.2%} ({perfect_docs}/{len(evaluations)})", ) results_table.add_row("Strict Nullability (Anti-Hallucination)", f"{strict_null_acc:.2%}") results_table.add_row("Page Grounding Accuracy", f"{avg_page_acc:.2%}") results_table.add_row("Avg Document Latency", f"{avg_latency / 1000.0:.2f} s") results_table.add_row("Total Execution Time", f"{total_duration:.1f} s") console.print("\n") console.print(results_table) console.print( f"\n[bold green]✓ Detailed outputs saved to:[/bold green] " f"[underline]{output_dir}[/underline]\n" ) # --------------------------------------------------------------------------- # Benchmark Runner Orchestrator # --------------------------------------------------------------------------- async def run_benchmark( api_url: str, split: str, batch_size: int, output_dir: Path, limit: int | None = None, offset: int = 0, target_doc_id: str | None = None, enable_reranker: bool = False, hf_token: str | None = None, dry_run: bool = False, resume: bool = False, rescore: bool = False, timeout: float = 600.0, ) -> None: """Execute the end-to-end ExtractBench benchmark suite against DocNeedle API.""" output_dir.mkdir(parents=True, exist_ok=True) summary_file = output_dir / f"summary_{split}.json" detailed_file = output_dir / f"detailed_{split}.jsonl" if rescore: if not detailed_file.exists(): console.print( f"[bold red]Hata:[/bold red] {detailed_file} bulunamadı, rescore yapılamaz." ) return console.print( Panel.fit( f"[bold cyan]DocNeedle Benchmark Re-Scorer[/bold cyan]\n" f"[dim]Re-evaluating:[/dim] {detailed_file} with updated normalizers", title="ExtractBench Re-Score", border_style="cyan", ) ) loader = ExtractBenchLoader(token=hf_token) records_map = {r["id"]: r for r in loader.load_split_records(split)} updated_evals = [] tmp_detailed = output_dir / f"detailed_{split}.tmp.jsonl" with ( open(detailed_file, "r", encoding="utf-8") as rf, open(tmp_detailed, "w", encoding="utf-8") as wf, ): for line in rf: if not line.strip(): continue sdoc = json.loads(line) doc_id = sdoc["doc_id"] rec = records_map.get(doc_id) if not rec: continue exp = json.loads(rec["expected_output"]) rules = json.loads(rec["field_rules"]) preds = { f["key"]: { "found": f["predicted_found"], "value": f["predicted"], "evidence": {"page": 1 if f.get("page_matched") else None}, } for f in sdoc["fields"] } ev = evaluate_document_case( doc_id=doc_id, category=sdoc.get("category", split), expected_output=exp, field_rules=rules, predictions=preds, latency_ms=sdoc.get("latency_ms", 0.0), ) updated_evals.append(ev) doc_dict = { "doc_id": ev.doc_id, "category": ev.category, "total_fields": ev.total_fields, "matched_fields": ev.matched_fields, "true_positives": ev.true_positives, "false_positives": ev.false_positives, "false_negatives": ev.false_negatives, "true_negatives": ev.true_negatives, "value_f1": ev.value_f1, "page_accuracy": ev.page_accuracy, "latency_ms": ev.latency_ms, "is_perfect": ev.is_perfect, "fields": [ { "key": f.key, "matched": f.matched, "expected": f.expected, "predicted": f.predicted, "expected_found": f.expected_found, "predicted_found": f.predicted_found, "page_matched": f.page_matched, } for f in ev.fields ], } wf.write(json.dumps(doc_dict) + "\n") tmp_detailed.replace(detailed_file) evaluations = updated_evals total_duration = sum(e.latency_ms for e in evaluations) / 1000.0 # Calculate summary and return _render_and_save_summary(evaluations, total_duration, summary_file, output_dir, split) return console.print( Panel.fit( f"[bold cyan]DocNeedle Benchmark Runner[/bold cyan]\n" f"[dim]Dataset:[/dim] {REPO_ID} | [dim]Split:[/dim] [yellow]{split}[/yellow]\n" f"[dim]Target API:[/dim] [green]{api_url}[/green] | " f"[dim]Batch Size:[/dim] {batch_size}\n" f"[dim]Output Directory:[/dim] {output_dir}", title="ExtractBench Evaluation", border_style="cyan", ) ) loader = ExtractBenchLoader(token=hf_token) client = DocNeedleBenchmarkClient( api_url=api_url, batch_size=batch_size, enable_reranker=enable_reranker, timeout=timeout, ) # 1. Load dataset records with console.status( f"[bold yellow]Loading records for split '{split}' from Hugging Face...[/bold yellow]" ): records = loader.load_split_records(split) if target_doc_id: records = [r for r in records if r.get("id") == target_doc_id] if not records: console.print(f"[bold red]Error:[/bold red] Document ID '{target_doc_id}' not found.") return # Apply offset and limit if offset > 0: records = records[offset:] if limit is not None and limit > 0: records = records[:limit] total_docs = len(records) console.print( f"[bold green]✓[/bold green] Prepared [bold]{total_docs}[/bold] documents for evaluation.\n" ) # Pre-flight check: verify API server availability before running if not dry_run: try: async with httpx.AsyncClient(timeout=4.0) as check_client: health_resp = await check_client.get(f"{api_url}/healthz") if health_resp.status_code != 200: console.print( f"[bold red]Hata:[/bold red] API sağlık kontrolü başarısız " f"({health_resp.status_code}): {api_url}/healthz" ) return except Exception: console.print( f"\n[bold red]Hata: DocNeedle API ({api_url}) bağlanılamadı![/bold red]\n" f"[yellow]Lütfen önce ayrı bir terminalde API sunucusunu başlatın:[/yellow]\n" f" [bold cyan]uv run uvicorn docneedle.main:app --host 0.0.0.0 --port 8000" f"[/bold cyan]\n" f"[dim]veya Gradio arayüzü çalışıyorsa (port 7860):[/dim]\n" f" [bold cyan]uv run python scripts/benchmark_extractbench.py " f"--api-url http://localhost:7860 --split {split}[/bold cyan]\n" ) return evaluations: list[DocumentEvaluation] = [] already_evaluated_ids: set[str] = set() if resume and detailed_file.exists(): with open(detailed_file, "r", encoding="utf-8") as f: for line in f: if not line.strip(): continue try: d = json.loads(line) already_evaluated_ids.add(d["doc_id"]) evaluations.append( DocumentEvaluation( doc_id=d["doc_id"], category=d.get("category", split), total_fields=d["total_fields"], matched_fields=d["matched_fields"], true_positives=d.get("true_positives", d["matched_fields"]), false_positives=d.get("false_positives", 0), false_negatives=d.get("false_negatives", 0), true_negatives=d.get("true_negatives", 0), value_f1=d["value_f1"], page_accuracy=d.get("page_accuracy", 1.0), latency_ms=d.get("latency_ms", 0.0), is_perfect=d.get("is_perfect", False), fields=[], ) ) except Exception: pass if already_evaluated_ids: console.print( f"[bold cyan]ℹ Resume etkin:[/bold cyan] " f"{len(already_evaluated_ids)} belge önceden tamamlanmış, atlanıyor.\n" ) records = [r for r in records if r.get("id") not in already_evaluated_ids] start_time = time.perf_counter() # Setup progress bar progress = Progress( SpinnerColumn(), TextColumn("[bold blue]{task.description}"), BarColumn(), TaskProgressColumn(), MofNCompleteColumn(), TimeElapsedColumn(), TimeRemainingColumn(), console=console, ) with progress: task_id = progress.add_task("[cyan]Evaluating documents...", total=total_docs) for idx, row in enumerate(records, 1): doc_id = row["id"] pdf_rel_path = row["pdf"] schema_dict = json.loads(row["data_schema"]) expected_output = json.loads(row["expected_output"]) field_rules = json.loads(row["field_rules"]) progress.update(task_id, description=f"[cyan]Doc {idx}/{total_docs}: {doc_id}") try: # 1. Ensure PDF is available locally pdf_path = loader.get_document_pdf_path(pdf_rel_path) # 2. Translate JSON Schema to DocNeedle ExtractionItemQuery list schema_items = translate_schema_properties(schema_dict, field_rules) if dry_run: # Dry run mode: check parsing and skip API call predictions: dict[str, dict[str, Any]] = {} latency_ms = 0.0 else: # 3. Call DocNeedle API (with transparent batching) predictions, latency_ms = await client.extract_all_fields( pdf_path, schema_items ) # 4. Evaluate predictions eval_res = evaluate_document_case( doc_id=doc_id, category=row.get("category", split), expected_output=expected_output, field_rules=field_rules, predictions=predictions, latency_ms=latency_ms, ) evaluations.append(eval_res) # Write line to detailed log immediately with open(detailed_file, "a", encoding="utf-8") as f: doc_dict = { "doc_id": eval_res.doc_id, "category": eval_res.category, "total_fields": eval_res.total_fields, "matched_fields": eval_res.matched_fields, "true_positives": eval_res.true_positives, "false_positives": eval_res.false_positives, "false_negatives": eval_res.false_negatives, "true_negatives": eval_res.true_negatives, "value_f1": eval_res.value_f1, "page_accuracy": eval_res.page_accuracy, "latency_ms": eval_res.latency_ms, "is_perfect": eval_res.is_perfect, "fields": [ { "key": f.key, "matched": f.matched, "expected": f.expected, "predicted": f.predicted, "expected_found": f.expected_found, "predicted_found": f.predicted_found, "page_matched": f.page_matched, } for f in eval_res.fields ], } f.write(json.dumps(doc_dict) + "\n") except Exception as e: console.print(f"[bold red]Failed on {doc_id}:[/bold red] {e}") logger.exception(f"Error evaluating {doc_id}") progress.advance(task_id) total_duration = time.perf_counter() - start_time # --------------------------------------------------------------------------- # Aggregate Metrics Calculation & Final Reporting # --------------------------------------------------------------------------- _render_and_save_summary(evaluations, total_duration, summary_file, output_dir, split) # --------------------------------------------------------------------------- # CLI Argument Parsing & Entry Point # --------------------------------------------------------------------------- def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser( description="ExtractBench benchmark runner for DocNeedle RAG extraction API.", formatter_class=argparse.ArgumentDefaultsHelpFormatter, ) parser.add_argument( "--api-url", type=str, default="http://localhost:8000", help="DocNeedle REST API base URL.", ) parser.add_argument( "--split", type=str, default="short", choices=VALID_SPLITS, help="Dataset split to evaluate (short: 252, medium: 98, long: 20, all: 370).", ) parser.add_argument( "--limit", type=int, default=None, help="Limit number of documents to evaluate (default: None, runs entire split).", ) parser.add_argument( "--offset", type=int, default=0, help="Skip first N documents in split (useful for resuming).", ) parser.add_argument( "--doc-id", type=str, default=None, help="Target a specific document ID (e.g. 'short/00581-2011-p0050').", ) parser.add_argument( "--batch-size", type=int, default=35, help="Schema items per API call (strictly <= 48 to stay within server 50-item limit).", ) parser.add_argument( "--timeout", type=float, default=600.0, help="HTTP request timeout in seconds per batch (default: 600.0).", ) parser.add_argument( "--output-dir", type=Path, default=Path("benchmark_results"), help="Directory to store evaluation reports and logs.", ) parser.add_argument( "--enable-reranker", action="store_true", help="Enable Qwen3 cross-encoder reranker in retrieval.", ) parser.add_argument( "--hf-token", type=str, default=os.getenv("HF_TOKEN"), help="Hugging Face API token (if needed for rate limits).", ) parser.add_argument( "--dry-run", action="store_true", help="Load, parse, and validate dataset records without sending API requests.", ) parser.add_argument( "--resume", action="store_true", help="Resume previous run by skipping already evaluated documents.", ) parser.add_argument( "--rescore", action="store_true", help="Re-score detailed_{split}.jsonl results with updated normalizers without inference.", ) return parser.parse_args() def main() -> None: args = parse_args() try: asyncio.run( run_benchmark( api_url=args.api_url, split=args.split, batch_size=args.batch_size, output_dir=args.output_dir, limit=args.limit, offset=args.offset, target_doc_id=args.doc_id, enable_reranker=args.enable_reranker, hf_token=args.hf_token, dry_run=args.dry_run, resume=args.resume, rescore=args.rescore, timeout=args.timeout, ) ) except KeyboardInterrupt: console.print("\n[yellow]Benchmark cancelled by user.[/yellow]") sys.exit(130) if __name__ == "__main__": main()