| """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", "-", "--", "—", "()", "( )", "[]", "[ ]", "{}"} |
| ) |
|
|
|
|
| |
| |
| |
|
|
|
|
| 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 |
| |
| 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. |
| """ |
| |
| if is_falsy_or_empty(expected) and is_falsy_or_empty(predicted): |
| return True |
|
|
| |
| 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: |
| |
| 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() |
|
|
| |
| 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 |
|
|
| |
| 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 |
| |
| if int(round(exp_n)) == int(round(pred_n)): |
| return True |
| return False |
|
|
| |
| 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] |
| |
| return normalize_text(expected) == normalize_text(predicted) |
|
|
| |
| 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): |
| |
| 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 |
|
|
| |
| 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 |
|
|
| |
| norm_exp = normalize_text(expected) |
| norm_pred = normalize_text(predicted) |
| if norm_exp == norm_pred: |
| return True |
|
|
| |
| 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 |
|
|
|
|
| |
| |
| |
|
|
|
|
| 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 "" |
|
|
| |
| 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 |
|
|
|
|
| |
| |
| |
|
|
|
|
| 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) |
|
|
|
|
| |
| |
| |
|
|
|
|
| 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)) |
| 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). |
| """ |
| |
| pdf_bytes = pdf_path.read_bytes() |
| filename = pdf_path.name |
|
|
| |
| 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 |
|
|
|
|
| |
| |
| |
|
|
|
|
| @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] = [] |
|
|
| |
| 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 [] |
|
|
| |
| 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 |
|
|
| |
| exp_found = not is_falsy_or_empty(expected_val) |
| pred_has_val = not is_falsy_or_empty(pred_val) |
|
|
| |
| |
| |
| 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 |
|
|
| |
| 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 |
| |
| 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 |
| elif not exp_found and pred_has_val: |
| fp += 1 |
| elif exp_found and not pred_has_val: |
| fn += 1 |
|
|
| 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" |
| ) |
|
|
|
|
| |
| |
| |
|
|
|
|
| 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 |
| |
| _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, |
| ) |
|
|
| |
| 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 |
|
|
| |
| 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" |
| ) |
|
|
| |
| 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() |
|
|
| |
| 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: |
| |
| pdf_path = loader.get_document_pdf_path(pdf_rel_path) |
|
|
| |
| schema_items = translate_schema_properties(schema_dict, field_rules) |
|
|
| if dry_run: |
| |
| predictions: dict[str, dict[str, Any]] = {} |
| latency_ms = 0.0 |
| else: |
| |
| predictions, latency_ms = await client.extract_all_fields( |
| pdf_path, schema_items |
| ) |
|
|
| |
| 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) |
|
|
| |
| 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 |
|
|
| |
| |
| |
| _render_and_save_summary(evaluations, total_duration, summary_file, output_dir, split) |
|
|
|
|
| |
| |
| |
|
|
|
|
| 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() |
|
|
