""" Generic Pre-Processing Pipeline (GPP) for Document Intelligence This module handles: 1. Parsing PDFs via MinerU Python API (OCR/text modes) 2. Extracting markdown, images, and content_list JSON 3. Chunking multimodal content (text, tables, images), ensuring tables/images are in single chunks 4. Parsing markdown tables into JSON 2D structures for dense tables 5. Narration of tables/images via LLM 6. Semantic enhancements (deduplication, coreference, metadata summarization) 7. Embedding computation for in-memory use Each step is modular to support swapping components (e.g. different parsers or stores). """ import os import json from typing import List, Dict, Any, Optional import re from magic_pdf.data.data_reader_writer import FileBasedDataWriter, FileBasedDataReader from magic_pdf.data.dataset import PymuDocDataset from magic_pdf.model.doc_analyze_by_custom_model import doc_analyze from magic_pdf.config.enums import SupportedPdfParseMethod from langchain.text_splitter import RecursiveCharacterTextSplitter from sentence_transformers import SentenceTransformer from rank_bm25 import BM25Okapi import numpy as np import hnswlib from src.config import EmbeddingConfig from src.utils import OpenAIEmbedder # LLM client abstraction from src.utils import LLMClient, logger def parse_markdown_table(md: str) -> Optional[Dict[str, Any]]: """ Parses a markdown table into a JSON-like dict: { headers: [...], rows: [[...], ...] } Handles multi-level headers by nesting lists if needed. """ lines = [l for l in md.strip().splitlines() if l.strip().startswith("|")] if len(lines) < 2: return None header_line = lines[0] sep_line = lines[1] # Validate separator line if not re.match(r"^\|?\s*:?-+:?\s*(\|\s*:?-+:?\s*)+\|?", sep_line): return None def split_row(line): parts = [cell.strip() for cell in line.strip().strip("|").split("|")] return parts headers = split_row(header_line) rows = [split_row(r) for r in lines[2:]] return {"headers": headers, "rows": rows} class GPPConfig: """ Configuration for GPP pipeline. """ CHUNK_TOKEN_SIZE = 256 DEDUP_SIM_THRESHOLD = 0.9 EXPANSION_SIM_THRESHOLD = 0.85 COREF_CONTEXT_SIZE = 3 HNSW_EF_CONSTRUCTION = int(os.getenv("HNSW_EF_CONSTRUCTION", "200")) HNSW_M = int(os.getenv("HNSW_M", "16")) HNSW_EF_SEARCH = int(os.getenv("HNSW_EF_SEARCH", "50")) class GPP: def __init__(self, config: GPPConfig): self.config = config # Embedding models if EmbeddingConfig.PROVIDER == "openai": self.text_embedder = OpenAIEmbedder(EmbeddingConfig.TEXT_MODEL) self.meta_embedder = OpenAIEmbedder(EmbeddingConfig.META_MODEL) else: self.text_embedder = SentenceTransformer( EmbeddingConfig.TEXT_MODEL, use_auth_token=True ) self.meta_embedder = SentenceTransformer( EmbeddingConfig.META_MODEL, use_auth_token=True ) self.bm25 = None def parse_pdf(self, pdf_path: str, output_dir: str) -> Dict[str, Any]: """ Uses MinerU API to parse PDF in OCR/text mode, dumps markdown, images, layout PDF, content_list JSON. Returns parsed data plus file paths for UI traceability. """ name = os.path.splitext(os.path.basename(pdf_path))[0] img_dir = os.path.join(output_dir, "images") os.makedirs(img_dir, exist_ok=True) os.makedirs(output_dir, exist_ok=True) writer_imgs = FileBasedDataWriter(img_dir) writer_md = FileBasedDataWriter(output_dir) reader = FileBasedDataReader("") pdf_bytes = reader.read(pdf_path) ds = PymuDocDataset(pdf_bytes) if ds.classify() == SupportedPdfParseMethod.OCR: infer = ds.apply(doc_analyze, ocr=True) pipe = infer.pipe_ocr_mode(writer_imgs) else: infer = ds.apply(doc_analyze, ocr=False) pipe = infer.pipe_txt_mode(writer_imgs) # Visual layout pipe.draw_layout(os.path.join(output_dir, f"{name}_layout.pdf")) # Dump markdown & JSON pipe.dump_md(writer_md, f"{name}.md", os.path.basename(img_dir)) pipe.dump_content_list( writer_md, f"{name}_content_list.json", os.path.basename(img_dir) ) content_list_path = os.path.join(output_dir, f"{name}_content_list.json") with open(content_list_path, "r", encoding="utf-8") as f: blocks = json.load(f) # UI traceability paths return { "blocks": blocks, "md_path": os.path.join(output_dir, f"{name}.md"), "images_dir": img_dir, "layout_pdf": os.path.join(output_dir, f"{name}_layout.pdf"), "spans_pdf": os.path.join(output_dir, f"{name}_spans.pdf"), } def chunk_blocks(self, blocks: List[Dict[str, Any]]) -> List[Dict[str, Any]]: """ Creates chunks of ~CHUNK_TOKEN_SIZE tokens, but ensures any table/image block becomes its own chunk (unsplittable), flushing current text chunk as needed. """ chunks, current, token_count = [], {"text": "", "type": None, "blocks": []}, 0 for blk in blocks: btype = blk.get("type") text = blk.get("text", "") if btype in ("table", "img_path"): # Flush existing text chunk if current["blocks"]: chunks.append(current) current = {"text": "", "type": None, "blocks": []} token_count = 0 # Create isolated chunk for the table/image tbl_chunk = {"text": text, "type": btype, "blocks": [blk]} # Parse markdown table into JSON structure if applicable if btype == "table": tbl_struct = parse_markdown_table(text) tbl_chunk["table_structure"] = tbl_struct chunks.append(tbl_chunk) continue # Standard text accumulation count = len(text.split()) if token_count + count > self.config.CHUNK_TOKEN_SIZE and current["blocks"]: chunks.append(current) current = {"text": "", "type": None, "blocks": []} token_count = 0 current["text"] += text + "\n" current["type"] = current["type"] or btype current["blocks"].append(blk) token_count += count # Flush remaining if current["blocks"]: chunks.append(current) logger.info(f"Chunked into {len(chunks)} pieces (with tables/images isolated).") return chunks def narrate_multimodal(self, chunks: List[Dict[str, Any]]) -> None: """ For table/image chunks, generate LLM narration. Preserve table_structure in metadata. """ for c in chunks: if c["type"] in ("table", "img_path"): prompt = f"Describe this {c['type']} concisely:\n{c['text']}" c["narration"] = LLMClient.generate(prompt) else: c["narration"] = c["text"] def deduplicate(self, chunks: List[Dict[str, Any]]) -> List[Dict[str, Any]]: try: # embs = self.text_embedder.encode([c.get('narration', '') for c in chunks], convert_to_tensor=True) narrations = [c.get("narration", "") for c in chunks] if EmbeddingConfig.PROVIDER == "openai": embs = self.text_embedder.embed(narrations) else: embs = self.text_embedder.encode(narrations) keep = [] for i, emb in enumerate(embs): if not any( (emb @ embs[j]).item() / (np.linalg.norm(emb) * np.linalg.norm(embs[j]) + 1e-8) > self.config.DEDUP_SIM_THRESHOLD for j in keep ): keep.append(i) deduped = [chunks[i] for i in keep] logger.info(f"Deduplicated: {len(chunks)}→{len(deduped)}") return deduped except Exception as e: logger.error(f"Deduplication failed: {e}") return chunks def coref_resolution(self, chunks: List[Dict[str, Any]]) -> None: for idx, c in enumerate(chunks): start = max(0, idx - self.config.COREF_CONTEXT_SIZE) ctx = "\n".join(chunks[i].get("narration", "") for i in range(start, idx)) prompt = f"Context:\n{ctx}\nRewrite pronouns in:\n{c.get('narration', '')}" try: c["narration"] = LLMClient.generate(prompt) except Exception as e: logger.error(f"Coref resolution failed for chunk {idx}: {e}") def metadata_summarization(self, chunks: List[Dict[str, Any]]) -> None: sections: Dict[str, List[Dict[str, Any]]] = {} for c in chunks: sec = c.get("section", "default") sections.setdefault(sec, []).append(c) for sec, items in sections.items(): blob = "\n".join(i.get("narration", "") for i in items) try: summ = LLMClient.generate(f"Summarize this section:\n{blob}") for i in items: i.setdefault("metadata", {})["section_summary"] = summ except Exception as e: logger.error(f"Metadata summarization failed for section {sec}: {e}") def build_bm25(self, chunks: List[Dict[str, Any]]) -> None: """ Build BM25 index on token lists for sparse retrieval. """ tokenized = [c["narration"].split() for c in chunks] self.bm25 = BM25Okapi(tokenized) def compute_and_store(self, chunks: List[Dict[str, Any]], output_dir: str) -> None: """ 1. Compute embeddings for each chunk's narration (text_vec) and section_summary (meta_vec). 2. Build two HNSWlib indices (one for text_vecs, one for meta_vecs). 3. Save both indices to disk. 4. Dump human-readable chunk metadata (incl. section_summary) for traceability in the UI. """ # --- 1. Prepare embedder --- if EmbeddingConfig.PROVIDER.lower() == "openai": embedder = OpenAIEmbedder(EmbeddingConfig.TEXT_MODEL) embed_fn = embedder.embed else: st_model = SentenceTransformer( EmbeddingConfig.TEXT_MODEL, use_auth_token=True ) embed_fn = lambda texts: st_model.encode( texts, show_progress_bar=False ).tolist() # Batch compute text & meta embeddings --- narrations = [c["narration"] for c in chunks] meta_texts = [c.get("section_summary", "") for c in chunks] logger.info( "computing_embeddings", provider=EmbeddingConfig.PROVIDER, num_chunks=len(chunks), ) text_vecs = embed_fn(narrations) meta_vecs = embed_fn(meta_texts) if len(text_vecs) != len(chunks) or len(meta_vecs) != len(chunks): raise RuntimeError( f"Embedding count mismatch: text_vecs={len(text_vecs)}, meta_vecs={len(meta_vecs)}, chunks={len(chunks)}" ) # Convert to numpy arrays text_matrix = np.vstack(text_vecs).astype(np.float32) meta_matrix = np.vstack(meta_vecs).astype(np.float32) # Build HNSW indices --- dim = text_matrix.shape[1] text_index = hnswlib.Index(space="cosine", dim=dim) text_index.init_index( max_elements=len(chunks), ef_construction=GPPConfig.HNSW_EF_CONSTRUCTION, M=GPPConfig.HNSW_M, ) ids = [c["id"] for c in chunks] text_index.add_items(text_matrix, ids) text_index.set_ef(GPPConfig.HNSW_EF_SEARCH) logger.info("text_hnsw_built", elements=len(chunks)) # Meta index (same dim) meta_index = hnswlib.Index(space="cosine", dim=dim) meta_index.init_index( max_elements=len(chunks), ef_construction=GPPConfig.HNSW_EF_CONSTRUCTION, M=GPPConfig.HNSW_M, ) meta_index.add_items(meta_matrix, ids) meta_index.set_ef(GPPConfig.HNSW_EF_SEARCH) logger.info("meta_hnsw_built", elements=len(chunks)) # Persist indices to disk --- text_idx_path = os.path.join(output_dir, "hnsw_text_index.bin") meta_idx_path = os.path.join(output_dir, "hnsw_meta_index.bin") text_index.save_index(text_idx_path) meta_index.save_index(meta_idx_path) logger.info( "hnsw_indices_saved", text_index=text_idx_path, meta_index=meta_idx_path ) # Dump chunk metadata for UI traceability --- meta_path = os.path.join(output_dir, "chunk_metadata.json") metadata = { str(c["id"]): { "text": c.get("text", ""), "narration": c["narration"], "type": c.get("type", ""), "section_summary": c.get("section_summary", ""), } for c in chunks } with open(meta_path, "w", encoding="utf-8") as f: json.dump(metadata, f, ensure_ascii=False, indent=2) logger.info("chunk_metadata_saved", path=meta_path) def run(self, pdf_path: str, output_dir: str) -> Dict[str, Any]: """ Executes full GPP: parse -> chunk -> narrate -> enhance -> index. Returns parse output dict augmented with `chunks` for downstream processes. """ parsed = self.parse_pdf(pdf_path, output_dir) blocks = parsed.get("blocks", []) chunks = self.chunk_blocks(blocks) # assigning ID's to chuncks for traceability for idx, chunk in enumerate(chunks): chunk["id"] = idx self.narrate_multimodal(chunks) chunks = self.deduplicate(chunks) self.coref_resolution(chunks) self.metadata_summarization(chunks) self.build_bm25(chunks) self.compute_and_store(chunks, output_dir) parsed["chunks"] = chunks logger.info("GPP pipeline complete.") return parsed