67 lines
2.4 KiB
Python
67 lines
2.4 KiB
Python
# ------------------------------------------------------------
|
|
# prepare_pdf.py
|
|
# ------------------------------------------------------------
|
|
# Utility script for running PDF→Parquet conversions
|
|
# either for a predefined list of documents or for a single file.
|
|
#
|
|
# Mode 1 (recommended):
|
|
# Fill the CHAPTERS/DOCUMENTS list and run:
|
|
# python prepare_pdf.py
|
|
#
|
|
# Mode 2 (single PDF from CLI):
|
|
# python prepare_pdf.py --pdf "/path/to/notebook.pdf" --doc-id "DOC-001-File-Notebook"
|
|
#
|
|
# Output: Each run produces Parquet embeddings and a JSONL manifest
|
|
#
|
|
# Requirements:
|
|
# Set the same API key environment variable as used in the project:
|
|
# export EMBEDDING_API_KEY=...
|
|
# ------------------------------------------------------------
|
|
|
|
import argparse
|
|
from typing import List, Tuple
|
|
from build_embeddings import build_for_pdf
|
|
|
|
# ====== List of documents with (doc_id, path) ======
|
|
CHAPTERS: List[Tuple[str, str]] = [
|
|
# Example entries:
|
|
# ("DOC-001-Employee-Handbook", r"/path/to/Employee-Handbook.pdf"),
|
|
# ("DOC-002-Product-Manual", r"/path/to/Product-Manual.pdf"),
|
|
# ("DOC-003-Research-Paper", r"/path/to/Research-Paper.pdf"),
|
|
|
|
# Leave empty if running with --pdf and --doc-id arguments only.
|
|
]
|
|
|
|
def run_single(pdf: str, doc_id: str) -> None:
|
|
print(f"[RUN] {doc_id} ← {pdf}")
|
|
out = build_for_pdf(pdf, doc_id)
|
|
print(f"[OK ] {doc_id} → {out}")
|
|
|
|
def run_batch(pairs: List[Tuple[str, str]]) -> None:
|
|
if not pairs:
|
|
print("[WARN] CHAPTERS list is empty. Fill it or use --pdf/--doc-id parameters.")
|
|
return
|
|
errors = 0
|
|
for doc_id, pdf in pairs:
|
|
try:
|
|
run_single(pdf, doc_id)
|
|
except Exception as e:
|
|
errors += 1
|
|
print(f"[FAIL] {doc_id}: {e}")
|
|
print(f"\n[DONE] Run complete. Total: {len(pairs)}, Success: {len(pairs) - errors}, Failures: {errors}")
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(description="Prepare PDFs → Parquet (single or batch).")
|
|
parser.add_argument("--pdf", help="Path to a single PDF file.")
|
|
parser.add_argument("--doc-id", help="Document identifier, e.g. DOC-001-Employee-Handbook.")
|
|
args = parser.parse_args()
|
|
|
|
if args.pdf and args.doc_id:
|
|
run_single(args.pdf, args.doc_id)
|
|
elif args.pdf or args.doc_id:
|
|
parser.error("Both --pdf and --doc-id must be provided together, or neither (use CHAPTERS list).")
|
|
else:
|
|
run_batch(CHAPTERS)
|
|
|
|
if __name__ == "__main__":
|
|
main() |