""" Parent-Child Hierarchical Chunker ===================================================== Parent = full judgment text (used for context retrieval) Child = fixed-token sub-chunks of parent (used for embedding + search) Reads: data/html/extracted_judgments.jsonl (metadata) data/pdfs/*.pdf (judgment text) Writes: data/chunks/parent_child/.json Architecture: Query hits a child chunk (small, precise, embedded) ↓ Child carries parent_chunk_id ↓ Fetch parent for full context window ↓ Send parent text to LLM for answer generation """ import json import os import re import fitz import logging import tiktoken import concurrent.futures from pathlib import Path from datetime import datetime # --------------------------------------------------------------------------- # Logging # --------------------------------------------------------------------------- logging.basicConfig( level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s", datefmt="%H:%M:%S", handlers=[ logging.StreamHandler(), logging.FileHandler("data/parent_child_chunker.log", encoding="utf-8"), ] ) log = logging.getLogger(__name__) # --------------------------------------------------------------------------- # Config # --------------------------------------------------------------------------- METADATA_FILE = os.path.join("data", "html", "extracted_judgments.jsonl") PDF_DIR = Path("data", "pdfs") OUTPUT_DIR = Path("data", "chunks", "parent_child") ERROR_LOG = Path("data", "chunks", "parent_child_errors.jsonl") CHILD_CHUNK_SIZE = 512 # tokens per child chunk CHILD_CHUNK_OVERLAP = 100 # token overlap between children TOKENIZER_MODEL = "cl100k_base" MAX_WORKERS = min(8, (os.cpu_count() or 4)) # Parent size threshold — if judgment is smaller than this, skip children MIN_TOKENS_FOR_CHILDREN = CHILD_CHUNK_SIZE OUTPUT_DIR.mkdir(parents=True, exist_ok=True) # --------------------------------------------------------------------------- # Globals (shared across threads — all read-only after init) # --------------------------------------------------------------------------- TOKENIZER = tiktoken.get_encoding(TOKENIZER_MODEL) # Pre-compiled regex RE_BACKSPACE = re.compile(r"\x08") RE_AUTHOR = re.compile(r"\n\*\s*Author\n\d+\n") RE_PAGE_NUMS = re.compile(r"\n\s*\d{1,3}\s*\n") RE_FOOTER = re.compile(r"\nJudgment\s*/\s*Order of the Supreme Court\n?", re.I) RE_HEADER = re.compile(r"\n(Supreme Court of India|IN THE SUPREME COURT OF INDIA)\n", re.I) RE_NEWLINES = re.compile(r"\n{3,}") RE_SPACES = re.compile(r"[ \t]+") # --------------------------------------------------------------------------- # Step 1: Load metadata index keyed by PDF filename # --------------------------------------------------------------------------- def load_metadata_index(jsonl_path: str) -> dict: """ Returns: { "2026_INSC_479.pdf": { ...full metadata record... }, ... } """ index = {} missing = 0 with open(jsonl_path, "r", encoding="utf-8") as f: for line in f: line = line.strip() if not line: continue try: record = json.loads(line) except json.JSONDecodeError: continue pdf_path = record.get("pdf_path", "") if pdf_path: filename = Path(pdf_path).name # "2026_INSC_479.pdf" index[filename] = record else: # Fallback: derive from neutral citation nc = record.get("neutral_citation", "").strip() if nc: filename = nc.replace(" ", "_") + ".pdf" index[filename] = record missing += 1 log.info(f"Loaded {len(index)} metadata records " f"({missing} used neutral_citation fallback).") return index # --------------------------------------------------------------------------- # Step 2: PDF text extraction # --------------------------------------------------------------------------- def extract_pdf_text(pdf_path: Path) -> str: """Extract full text from PDF using pymupdf.""" parts = [] try: doc = fitz.open(str(pdf_path)) for page in doc: parts.append(page.get_text()) doc.close() except Exception as e: log.error(f"PDF read failed [{pdf_path.name}]: {e}") return "\n".join(parts) # --------------------------------------------------------------------------- # Step 3: Text cleaning # --------------------------------------------------------------------------- def clean_text(text: str) -> str: """Remove PDF artifacts, headers, footers, page numbers.""" if not text: return "" text = RE_BACKSPACE.sub("", text) text = RE_AUTHOR.sub("\n", text) text = RE_PAGE_NUMS.sub("\n\n", text) text = RE_FOOTER.sub("\n", text) text = RE_HEADER.sub("\n", text) text = RE_NEWLINES.sub("\n\n", text) text = RE_SPACES.sub(" ", text) return text.strip() # --------------------------------------------------------------------------- # Step 4: Build lean metadata (for child chunks) # --------------------------------------------------------------------------- def build_lean_metadata(record: dict) -> dict: """ Child chunks carry only the fields needed for Qdrant payload filtering. Full metadata lives on the parent — fetched at answer-generation time. """ return { "case_name": record.get("case_name", ""), "neutral_citation": record.get("neutral_citation", ""), "date": record.get("date", ""), "court": record.get("court", "Supreme Court"), "case_type": record.get("case_type", ""), "outcome": record.get("outcome", ""), "acts": record.get("acts", []), "keywords": record.get("keywords", []), } # --------------------------------------------------------------------------- # Step 5: Build full metadata (for parent chunk) # --------------------------------------------------------------------------- def build_full_metadata(record: dict) -> dict: return { "case_name": record.get("case_name", ""), "neutral_citation": record.get("neutral_citation", ""), "appeal_no": record.get("appeal_no", ""), "citation": record.get("citation", ""), "date": record.get("date", ""), "court": record.get("court", "Supreme Court"), "lower_court": record.get("lower_court", ""), "jurisdiction": record.get("jurisdiction", "India"), "bench": record.get("bench", []), "author_judge": record.get("author_judge", ""), "outcome": record.get("outcome", ""), "case_type": record.get("case_type", ""), "acts": record.get("acts", []), "sections": record.get("sections", []), "cases_cited": record.get("cases_cited", []), "keywords": record.get("keywords", []), "issue": record.get("issue", ""), "short_summary": record.get("short_summary", ""), "full_headnote": record.get("full_headnote", ""), "source_url": record.get("source_url", ""), "scraped_at": record.get("scraped_at", ""), } # --------------------------------------------------------------------------- # Step 6: Core parent-child chunking logic # --------------------------------------------------------------------------- def build_parent_child( neutral_citation: str, cleaned_text: str, full_metadata: dict, lean_metadata: dict, ) -> dict: """ Returns: { "parent": { single parent chunk with full text + full metadata }, "children": [ child chunks with sub-text + lean metadata ] } """ safe_nc = neutral_citation.replace(" ", "_") parent_id = f"{safe_nc}__parent" tokens = TOKENIZER.encode(cleaned_text) num_tokens = len(tokens) # --- Parent chunk --- parent = { "chunk_id": parent_id, "chunk_type": "parent", "document_id": neutral_citation, "text": cleaned_text, "token_count": num_tokens, "char_count": len(cleaned_text), "metadata": full_metadata, "chunked_at": datetime.utcnow().isoformat(), } # --- Skip children if text is too small --- if num_tokens <= MIN_TOKENS_FOR_CHILDREN: parent["child_count"] = 0 return {"parent": parent, "children": []} # --- Child chunks --- step = CHILD_CHUNK_SIZE - CHILD_CHUNK_OVERLAP children = [] idx = 0 for start in range(0, num_tokens, step): end = min(start + CHILD_CHUNK_SIZE, num_tokens) chunk_tokens = tokens[start:end] chunk_text = TOKENIZER.decode(chunk_tokens) children.append({ "chunk_id": f"{safe_nc}__child_{idx:04d}", "chunk_type": "child", "parent_chunk_id": parent_id, "document_id": neutral_citation, "child_index": idx, "token_count": len(chunk_tokens), "char_count": len(chunk_text), "start_token": start, "end_token": end, "text": chunk_text, "metadata": lean_metadata, }) idx += 1 if end == num_tokens: break parent["child_count"] = len(children) return {"parent": parent, "children": children} # --------------------------------------------------------------------------- # Step 7: Process single PDF (called by thread pool) # --------------------------------------------------------------------------- def process_pdf(pdf_path: Path, metadata_index: dict) -> dict: """ Returns a result dict: { "status": "success" | "skipped" | "error", "file": pdf filename, "message": description, "chunks": { parent, children } or None } """ filename = pdf_path.name # --- Idempotency: skip if already processed --- record = metadata_index.get(filename) if not record: return {"status": "unmatched", "file": filename, "message": f"No metadata found for {filename}"} neutral_citation = record.get("neutral_citation", "") safe_nc = neutral_citation.replace(" ", "_") output_file = OUTPUT_DIR / f"{safe_nc}.json" if output_file.exists(): return {"status": "skipped", "file": filename, "message": f"Already processed: {output_file.name}"} # --- Extract + clean text --- raw_text = extract_pdf_text(pdf_path) cleaned = clean_text(raw_text) if not cleaned: return {"status": "error", "file": filename, "message": "Empty text after cleaning"} # --- Build metadata --- full_meta = build_full_metadata(record) lean_meta = build_lean_metadata(record) # --- Build parent-child structure --- result = build_parent_child(neutral_citation, cleaned, full_meta, lean_meta) # --- Write output atomically --- # Write to temp file first, then rename — prevents corrupt files on crash temp_file = output_file.with_suffix(".tmp") try: with open(temp_file, "w", encoding="utf-8") as f: json.dump(result, f, ensure_ascii=False, indent=2) temp_file.rename(output_file) except Exception as e: if temp_file.exists(): temp_file.unlink() return {"status": "error", "file": filename, "message": str(e)} return { "status": "success", "file": filename, "message": f"{result['parent']['child_count']} children created", "children": result["parent"]["child_count"], } # --------------------------------------------------------------------------- # Step 8: Main pipeline # --------------------------------------------------------------------------- def main(): log.info("=" * 60) log.info("Parent-Child Chunker — Starting") log.info("=" * 60) # Load metadata metadata_index = load_metadata_index(METADATA_FILE) # Discover PDFs pdf_files = sorted(PDF_DIR.glob("*.pdf")) log.info(f"Found {len(pdf_files)} PDFs in {PDF_DIR}") if not pdf_files: log.error("No PDFs found. Check PDF_DIR path.") return # Counters counts = {"success": 0, "skipped": 0, "unmatched": 0, "error": 0} total_children = 0 errors = [] # Process concurrently log.info(f"Processing with {MAX_WORKERS} workers...") with concurrent.futures.ThreadPoolExecutor(max_workers=MAX_WORKERS) as executor: futures = { executor.submit(process_pdf, pdf_path, metadata_index): pdf_path for pdf_path in pdf_files } for i, future in enumerate(concurrent.futures.as_completed(futures), 1): pdf_path = futures[future] try: result = future.result() status = result["status"] counts[status] = counts.get(status, 0) + 1 if status == "success": total_children += result.get("children", 0) if i % 50 == 0: log.info( f"Progress: {i}/{len(pdf_files)} | " f"Success: {counts['success']} | " f"Skipped: {counts['skipped']} | " f"Errors: {counts['error']}" ) elif status == "error": log.warning(f"[ERROR] {result['file']}: {result['message']}") errors.append(result) elif status == "unmatched": log.warning(f"[UNMATCHED] {result['file']}") errors.append(result) except Exception as exc: counts["error"] += 1 log.error(f"[EXCEPTION] {pdf_path.name}: {exc}") errors.append({"file": pdf_path.name, "message": str(exc)}) # Write error log if errors: with open(ERROR_LOG, "w", encoding="utf-8") as f: for e in errors: f.write(json.dumps(e, ensure_ascii=False) + "\n") log.info(f"Error details → {ERROR_LOG}") # Final summary log.info("=" * 60) log.info("PIPELINE COMPLETE") log.info(f" Successful : {counts['success']}") log.info(f" Skipped : {counts['skipped']} (already processed)") log.info(f" Unmatched : {counts['unmatched']} (no metadata)") log.info(f" Errors : {counts['error']}") log.info(f" Total children created : {total_children}") log.info(f" Output dir : {OUTPUT_DIR.resolve()}") log.info("=" * 60) if __name__ == "__main__": main()