123 lines
4.3 KiB
Python
123 lines
4.3 KiB
Python
import os
|
|
import argparse
|
|
import hashlib
|
|
from datetime import datetime, timezone
|
|
from typing import List, Dict, Any
|
|
|
|
from config import (
|
|
MODEL_NAME, EMBEDDING_DIM, CHUNK_SIZE, OVERLAP, BATCH_SIZE,
|
|
EMB_DIR, MANIFEST_PATH, AVG_CHARS_PER_TOKEN, PRICE_PER_1K_TOKENS
|
|
)
|
|
from pdf_text import extract_text_from_pdf
|
|
from chunker import basic_clean, chunk_text
|
|
from embedder import get_client, embed_texts
|
|
from writer import write_parquet, append_manifest
|
|
|
|
|
|
def _estimate_tokens(texts: List[str]) -> tuple[int, str]:
|
|
"""Estimate total tokens for a list of texts. Prefer tiktoken; fallback to char heuristic."""
|
|
try:
|
|
import tiktoken # type: ignore
|
|
enc = tiktoken.get_encoding("cl100k_base")
|
|
total = sum(len(enc.encode(t)) for t in texts)
|
|
return total, "tiktoken"
|
|
except Exception:
|
|
total_chars = sum(len(t) for t in texts)
|
|
approx = int(total_chars / AVG_CHARS_PER_TOKEN)
|
|
return approx, "chars-heuristic"
|
|
|
|
|
|
def sha1(text: str) -> str:
|
|
return hashlib.sha1(text.encode("utf-8", errors="ignore")).hexdigest()
|
|
|
|
|
|
def build_for_pdf(pdf_path: str, doc_id: str) -> str:
|
|
if not os.path.isfile(pdf_path):
|
|
raise FileNotFoundError(f"PDF not found: {pdf_path}")
|
|
|
|
print(f"[INFO] Reading PDF: {pdf_path}")
|
|
raw_text, total_pages = extract_text_from_pdf(pdf_path)
|
|
cleaned = basic_clean(raw_text)
|
|
|
|
print(f"[INFO] Chunking (size={CHUNK_SIZE}, overlap={OVERLAP})…")
|
|
chunks = chunk_text(cleaned, CHUNK_SIZE, OVERLAP)
|
|
if not chunks:
|
|
raise RuntimeError("No chunks produced; check PDF text extraction or chunking settings.")
|
|
|
|
texts = [c[2] for c in chunks]
|
|
positions = [(c[0], c[1]) for c in chunks]
|
|
|
|
# --- FACT LOGS ---
|
|
print(f"[INFO] PAGES: {total_pages}")
|
|
print(f"[INFO] CHUNKS: {len(texts)}")
|
|
|
|
# --- ESTIMATION LOGS ---
|
|
est_tokens, method = _estimate_tokens(texts)
|
|
est_cost = (est_tokens / 1000.0) * PRICE_PER_1K_TOKENS
|
|
print(f"[INFO] ESTIMATE: tokens≈{est_tokens:,} (via {method}) | cost≈${est_cost:.3f} | "
|
|
f"avg_chars_per_token={AVG_CHARS_PER_TOKEN}")
|
|
|
|
print(f"[INFO] Embedding {len(texts)} chunks with model={MODEL_NAME} (batch={BATCH_SIZE})…")
|
|
client = get_client()
|
|
vectors = embed_texts(texts, client=client, model=MODEL_NAME, batch_size=BATCH_SIZE)
|
|
|
|
if len(vectors) != len(texts):
|
|
raise RuntimeError("Mismatch between vectors and texts length.")
|
|
|
|
# ISO 8601 UTC without microseconds
|
|
created_at = datetime.now(timezone.utc).replace(microsecond=0).isoformat()
|
|
source_path = os.path.abspath(pdf_path)
|
|
|
|
# Prepare Parquet rows
|
|
rows: List[Dict[str, Any]] = []
|
|
for idx, ((start, end), txt, vec) in enumerate(zip(positions, texts, vectors)):
|
|
rows.append({
|
|
"doc_id": doc_id,
|
|
"source_path": source_path,
|
|
"page_start": None, # Optional: could be inferred with a more advanced pager
|
|
"page_end": None,
|
|
"chunk_id": idx,
|
|
"start_char": start,
|
|
"end_char": end,
|
|
"text": txt,
|
|
"embedding": vec,
|
|
"model_name": MODEL_NAME,
|
|
"dim": EMBEDDING_DIM,
|
|
"chunk_size": CHUNK_SIZE,
|
|
"overlap": OVERLAP,
|
|
"hash": sha1(txt),
|
|
"created_at": created_at
|
|
})
|
|
|
|
out_path = os.path.join(EMB_DIR, f"{doc_id}.parquet")
|
|
write_parquet(rows, out_path)
|
|
|
|
manifest_entry = {
|
|
"doc_id": doc_id,
|
|
"source_path": source_path,
|
|
"num_chunks": len(rows),
|
|
"model_name": MODEL_NAME,
|
|
"dim": EMBEDDING_DIM,
|
|
"chunk_size": CHUNK_SIZE,
|
|
"overlap": OVERLAP,
|
|
"created_at": created_at,
|
|
"parquet_path": out_path
|
|
}
|
|
append_manifest([manifest_entry])
|
|
|
|
print(f"[INFO] Manifest appended at: {MANIFEST_PATH}")
|
|
return out_path
|
|
|
|
|
|
if __name__ == "__main__":
|
|
parser = argparse.ArgumentParser(description="Build embeddings Parquet for a single PDF.")
|
|
parser.add_argument("--pdf", required=True, help="Path to the PDF file.")
|
|
parser.add_argument("--doc-id", required=True, help="Document ID, e.g., DOC-001-Employee-Handbook.")
|
|
args = parser.parse_args()
|
|
|
|
try:
|
|
path = build_for_pdf(args.pdf, args.doc_id)
|
|
print(f"[DONE] Embeddings built: {path}")
|
|
except Exception as e:
|
|
print(f"[ERROR] {e}")
|
|
raise |