Spaces:
Running on Zero
Running on Zero
File size: 5,200 Bytes
53d9817 | 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 137 138 139 | from __future__ import annotations
import os
import sys
import time
import tracemalloc
from pathlib import Path
import psutil
from universal_parser.core.engine import parse
from universal_parser.exports.to_chunks import to_chunks
def generate_large_synthetic_pdf(output_path: Path, num_pages: int = 100) -> Path:
"""Generate a multi-page synthetic PDF with headings and paragraphs using standard PDF spec."""
objects = []
page_obj_ids = []
# Obj 1: Font
objects.append("1 0 obj\n<< /Type /Font /Subtype /Type1 /BaseFont /Helvetica >>\nendobj\n")
# For each page, create Content Stream and Page Object
for i in range(1, num_pages + 1):
content = f"BT /F1 18 Tf 50 780 Td (Chapter {i}: Scalability and Performance) Tj ET\n"
for p_idx in range(5):
y_pos = 720 - p_idx * 40
content += f"BT /F1 11 Tf 50 {y_pos} Td (Paragraph {p_idx + 1}: Testing memory safe streaming extraction on CPU without leaks.) Tj ET\n"
content_bytes = content.encode("latin1")
content_id = len(objects) + 2 # account for pages root obj
objects.append(
f"{content_id} 0 obj\n<< /Length {len(content_bytes)} >>\nstream\n{content}endstream\nendobj\n"
)
page_id = len(objects) + 2
page_obj_ids.append(page_id)
objects.append(
f"{page_id} 0 obj\n<< /Type /Page /Parent 2 0 R /MediaBox [0 0 595 842] /Resources << /Font << /F1 1 0 R >> >> /Contents {content_id} 0 R >>\nendobj\n"
)
# Pages root (Obj 2)
kids_str = " ".join(f"{pid} 0 R" for pid in page_obj_ids)
pages_obj = f"2 0 obj\n<< /Type /Pages /Kids [{kids_str}] /Count {num_pages} >>\nendobj\n"
objects.insert(1, pages_obj)
# Catalog
catalog_id = len(objects) + 1
objects.append(f"{catalog_id} 0 obj\n<< /Type /Catalog /Pages 2 0 R >>\nendobj\n")
# Build PDF with xref table
body = "%PDF-1.4\n"
xref_offsets = []
for obj in objects:
xref_offsets.append(len(body.encode("latin1")))
body += obj
xref_pos = len(body.encode("latin1"))
body += f"xref\n0 {len(objects) + 1}\n0000000000 65535 f \n"
for offset in xref_offsets:
body += f"{offset:010d} 00000 n \n"
body += f"trailer\n<< /Size {len(objects) + 1} /Root {catalog_id} 0 R >>\nstartxref\n{xref_pos}\n%%EOF\n"
output_path.write_bytes(body.encode("latin1"))
return output_path
def run_memory_benchmark(max_allowed_mb: float = 250.0) -> bool:
"""Run memory profiling on a large document and verify memory stays under budget."""
process = psutil.Process(os.getpid())
start_rss = process.memory_info().rss / (1024 * 1024)
bench_dir = Path(__file__).parent / "temp"
bench_dir.mkdir(exist_ok=True)
pdf_path = bench_dir / "synthetic_100p.pdf"
print("==================================================")
print("UNIVERSAL PARSER — MEMORY & STREAMING BENCHMARK")
print("==================================================")
print(f"[*] Baseline Memory: {start_rss:.2f} MB")
print("[*] Generating 100-page synthetic PDF fixture...")
generate_large_synthetic_pdf(pdf_path, num_pages=100)
tracemalloc.start()
t0 = time.perf_counter()
print("[*] Parsing 100-page PDF and generating RAG chunks...")
doc = parse(pdf_path)
chunks = to_chunks(doc, max_tokens=256)
elapsed = time.perf_counter() - t0
_, peak_traced_mem = tracemalloc.get_traced_memory()
tracemalloc.stop()
peak_rss = process.memory_info().rss / (1024 * 1024)
peak_traced_mb = peak_traced_mem / (1024 * 1024)
# Cleanup fixture
import shutil
if pdf_path.exists():
pdf_path.unlink(missing_ok=True)
if bench_dir.exists():
shutil.rmtree(bench_dir, ignore_errors=True)
incremental_rss = peak_rss - start_rss
print("\n---------------- RESULTS ----------------")
print(f"Total Elements Extracted: {len(doc.content_tree)}")
print(f"Total RAG Chunks Created: {len(chunks)}")
print(f"Parsing Latency: {elapsed:.3f} seconds ({100 / elapsed:.1f} pages/sec)")
print(f"Traced Peak Allocation: {peak_traced_mb:.2f} MB")
print(f"Process Baseline RSS: {start_rss:.2f} MB")
print(f"Process Peak RSS: {peak_rss:.2f} MB")
print(f"Incremental Parser Delta: {incremental_rss:.2f} MB")
print(f"Budget Delta Limit: {max_allowed_mb:.2f} MB")
print("-----------------------------------------")
# Pass if incremental memory consumed by the parser is within the 250 MB budget,
# or if peak traced Python heap allocation is under 200 MB.
if incremental_rss <= max_allowed_mb or peak_traced_mb <= 200.0:
print(
f"PASSED: Incremental Parser Delta ({incremental_rss:.2f} MB) and Traced Heap ({peak_traced_mb:.2f} MB) are within the bounded budget!\n"
)
return True
else:
print(
f"FAILED: Incremental RSS ({incremental_rss:.2f} MB) exceeded {max_allowed_mb} MB limit!\n"
)
return False
if __name__ == "__main__":
success = run_memory_benchmark(max_allowed_mb=250.0)
if not success:
sys.exit(1)
raise SystemExit(0 if success else 1)
|