Svngoku/PDF2Dataset
7
1from dotenv import load_dotenv2 3load_dotenv()4 5import gradio as gr6from chonkie import RecursiveChunker7from typing import Dict, Any, List, Optional8from dataclasses import dataclass, field9import logging10import re11import base6412import hashlib13import mimetypes14import json15from collections import Counter16from datasets import Dataset, Features, Value, Sequence, load_dataset17from datasets.features import Image as HFImage18from huggingface_hub import HfApi, get_token19import huggingface_hub20import os21from mistralai.client import Mistral22import fitz # pymupdf23from PIL import Image24import io25import tempfile26 27# Configure logging28logging.basicConfig(29 level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s"30)31logger = logging.getLogger(__name__)32 33 34# --- Exceptions ---35 36 37class OCRError(Exception):38 """Raised when OCR processing fails."""39 40 pass41 42 43# --- Mistral Client (lazy init) ---44 45_client: Mistral | None = None46 47 48def get_mistral_client() -> Mistral:49 """Get or initialize the Mistral client."""50 global _client51 if _client is not None:52 return _client53 54 api_key = os.environ.get("MISTRAL_API_KEY")55 if not api_key:56 logger.warning("MISTRAL_API_KEY not set. Attempting to use Hugging Face token.")57 api_key = get_token()58 if api_key:59 logger.info("Using Hugging Face token as MISTRAL_API_KEY.")60 61 if not api_key:62 raise OCRError(63 "No API key found. Set MISTRAL_API_KEY or run `huggingface-cli login`."64 )65 66 _client = Mistral(api_key=api_key)67 logger.info("Mistral client initialized successfully.")68 return _client69 70 71# --- Helper Functions ---72 73 74def encode_image_bytes(image_bytes: bytes) -> str:75 """Encodes image bytes to a base64 string."""76 return base64.b64encode(image_bytes).decode("utf-8")77 78 79def decode_base64_data_uri(data_uri: str) -> Optional[dict]:80 """Decode a base64 data URI to a HF-compatible image bytes dict.81 82 Args:83 data_uri: A string like "data:image/jpeg;base64,/9j/4AAQ..." or raw base64.84 85 Returns:86 Dict with {"bytes": <raw bytes>, "path": None} for datasets.Image feature,87 or None if decoding fails.88 """89 try:90 if data_uri.startswith("data:"):91 # Strip the "data:image/...;base64," prefix92 _, encoded = data_uri.split(",", 1)93 else:94 encoded = data_uri95 raw_bytes = base64.b64decode(encoded)96 # Validate it's a real image by opening it97 img = Image.open(io.BytesIO(raw_bytes))98 # Re-encode as PNG for consistency99 buf = io.BytesIO()100 img.save(buf, format="PNG")101 return {"bytes": buf.getvalue(), "path": None}102 except Exception as e:103 logger.warning(f"Failed to decode base64 image ({len(data_uri)} chars): {e}")104 return None105 106 107def extract_images_from_markdown(markdown_text: str) -> Dict[str, str]:108 """109 Extracts base64 image data URIs from markdown and maps them to reference IDs.110 Returns a dictionary mapping reference IDs to base64 data URIs.111 """112 image_map = {}113 img_refs = re.findall(114 r"!\[.*?\]\((data:image/[a-zA-Z+]+;base64,[A-Za-z0-9+/=]+)\)", markdown_text115 )116 for idx, img_uri in enumerate(img_refs):117 ref_id = f"img_ref_{idx + 1}"118 image_map[ref_id] = img_uri119 return image_map120 121 122def replace_image_references(markdown_text: str, image_map: Dict[str, str]) -> str:123 """124 Replaces base64 image data URIs in markdown with reference IDs (e.g., img_ref_1).125 """126 updated_markdown = markdown_text127 for ref_id, img_uri in image_map.items():128 escaped_uri = re.escape(img_uri)129 pattern = r"(!\[.*?\]\()" + escaped_uri + r"(\))"130 updated_markdown = re.sub(pattern, f"\\1{ref_id}\\2", updated_markdown)131 return updated_markdown132 133 134def get_combined_markdown(ocr_response: Any) -> tuple[str, str, Dict[str, str]]:135 """Combines markdown from OCR pages, replacing image IDs with base64 data URIs."""136 processed_markdowns = []137 raw_markdowns = []138 image_data_map = {}139 140 if not hasattr(ocr_response, "pages") or not ocr_response.pages:141 logger.warning("OCR response has no pages.")142 return "", "", {}143 144 for page_idx, page in enumerate(ocr_response.pages):145 if hasattr(page, "images") and page.images:146 logger.info(f"Page {page_idx}: Found {len(page.images)} images.")147 for img in page.images:148 if (149 hasattr(img, "id")150 and hasattr(img, "image_base64")151 and img.image_base64152 ):153 image_data_map[img.id] = img.image_base64154 else:155 logger.warning(156 f"Page {page_idx}: Image object lacks 'id' or valid 'image_base64'."157 )158 else:159 logger.info(f"Page {page_idx}: No images found.")160 161 if not hasattr(page, "markdown"):162 logger.warning(f"Page {page_idx} lacks 'markdown' attribute. Skipping.")163 continue164 165 current_raw_markdown = page.markdown or ""166 raw_markdowns.append(current_raw_markdown)167 current_processed_markdown = current_raw_markdown168 169 img_refs = re.findall(r"!\[.*?\]\((.*?)\)", current_processed_markdown)170 for img_id in img_refs:171 if img_id in image_data_map:172 base64_data_uri = image_data_map[img_id]173 escaped_img_id = re.escape(img_id)174 pattern = r"(!\[.*?\]\()" + escaped_img_id + r"(\))"175 current_processed_markdown = re.sub(176 pattern,177 r"\1" + base64_data_uri + r"\2",178 current_processed_markdown,179 )180 elif not img_id.startswith(("http:", "https:", "data:")):181 logger.warning(182 f"Page {page_idx}: Image ID '{img_id}' not in image data."183 )184 185 processed_markdowns.append(current_processed_markdown)186 187 logger.info(188 f"Processed {len(processed_markdowns)} pages with {len(image_data_map)} images."189 )190 return "\n\n".join(processed_markdowns), "\n\n".join(raw_markdowns), image_data_map191 192 193def perform_ocr(file_path: str) -> tuple[str, str, Dict[str, str]]:194 """Performs OCR on a file using Mistral API.195 196 Args:197 file_path: Path to the file on disk.198 199 Returns:200 Tuple of (processed_markdown, raw_markdown, image_data_map).201 202 Raises:203 OCRError: If OCR processing fails.204 """205 client = get_mistral_client()206 file_name = os.path.basename(file_path)207 file_ext = os.path.splitext(file_name)[1].lower()208 logger.info(f"Performing OCR on file: {file_name}")209 210 ocr_response = None211 supported_images = {".png", ".jpg", ".jpeg", ".webp", ".bmp"}212 213 if file_ext == ".pdf":214 uploaded_file_id = None215 try:216 with open(file_path, "rb") as f:217 file_content = f.read()218 219 logger.info(f"Uploading PDF {file_name} to Mistral...")220 uploaded_pdf = client.files.upload(221 file={"file_name": file_name, "content": file_content},222 purpose="ocr",223 )224 uploaded_file_id = uploaded_pdf.id225 logger.info(f"PDF uploaded. File ID: {uploaded_file_id}")226 227 signed_url_response = client.files.get_signed_url(file_id=uploaded_file_id)228 ocr_response = client.ocr.process(229 model="mistral-ocr-latest",230 document={231 "type": "document_url",232 "document_url": signed_url_response.url,233 },234 include_image_base64=True,235 )236 finally:237 if uploaded_file_id:238 try:239 client.files.delete(file_id=uploaded_file_id)240 except Exception as delete_err:241 logger.warning(242 f"Failed to delete temporary file {uploaded_file_id}: {delete_err}"243 )244 245 elif file_ext in supported_images:246 with open(file_path, "rb") as f:247 image_bytes = f.read()248 if not image_bytes:249 raise OCRError(f"Uploaded image file '{file_name}' is empty.")250 251 base64_encoded = encode_image_bytes(image_bytes)252 mime_type, _ = mimetypes.guess_type(file_path)253 mime_type = mime_type or "image/jpeg"254 data_uri = f"data:{mime_type};base64,{base64_encoded}"255 ocr_response = client.ocr.process(256 model="mistral-ocr-latest",257 document={"type": "image_url", "image_url": data_uri},258 include_image_base64=True,259 )260 else:261 raise OCRError(f"Unsupported file type: '{file_ext}'")262 263 if not ocr_response:264 raise OCRError(f"OCR returned no response for '{file_name}'.")265 266 processed_md, raw_md, img_map = get_combined_markdown(ocr_response)267 logger.info(f"Processed markdown length: {len(processed_md)}")268 return processed_md, raw_md, img_map269 270 271def _build_header_index(markdown_text: str) -> list[tuple[int, int, str]]:272 """Build a sorted index of (position, level, title) for all markdown headers."""273 headers = []274 for match in re.finditer(r"^(#{1,6})\s+(.+)$", markdown_text, re.MULTILINE):275 level = len(match.group(1))276 title = match.group(2).strip()277 headers.append((match.start(), level, title))278 return headers279 280 281def _get_headers_for_position(282 headers: list[tuple[int, int, str]], position: int283) -> dict[str, str]:284 """Given a character position, find the active chapter/section/subsection.285 286 Maps header levels: H1 -> chapter, H2 -> section, H3+ -> subsection.287 """288 active: dict[int, str] = {}289 for hdr_pos, level, title in headers:290 if hdr_pos > position:291 break292 active[level] = title293 # Clear deeper levels when a higher-level header appears294 for deeper in list(active.keys()):295 if deeper > level:296 del active[deeper]297 298 return {299 "chapter": active.get(1, ""),300 "section": active.get(2, ""),301 "subsection": active.get(3, active.get(4, active.get(5, active.get(6, "")))),302 }303 304 305def _clean_text(text: str) -> str:306 """Remove markdown formatting, image refs, and extra whitespace."""307 cleaned = re.sub(r"!\[.*?\]\(.*?\)", "", text)308 cleaned = re.sub(r"#{1,6}\s+", "", cleaned)309 cleaned = re.sub(r"\*\*(.+?)\*\*", r"\1", cleaned)310 cleaned = re.sub(r"\*(.+?)\*", r"\1", cleaned)311 cleaned = re.sub(r"`(.+?)`", r"\1", cleaned)312 cleaned = re.sub(r"\[(.+?)\]\(.*?\)", r"\1", cleaned)313 cleaned = re.sub(r"\n{3,}", "\n\n", cleaned)314 return cleaned.strip()315 316 317def chunk_markdown(318 markdown_text_with_images: str,319 chunk_size: int = 512,320) -> list[dict]:321 """Chunks markdown text using chonkie's RecursiveChunker with markdown recipe.322 323 Args:324 markdown_text_with_images: Markdown text possibly containing base64 image references.325 chunk_size: Maximum character count per chunk.326 327 Returns:328 List of chunk dicts with the full dataset schema fields.329 """330 if not markdown_text_with_images or not markdown_text_with_images.strip():331 logger.warning("chunk_markdown received empty input.")332 return []333 334 # Extract images and replace with reference IDs335 image_map = extract_images_from_markdown(markdown_text_with_images)336 updated_markdown = replace_image_references(markdown_text_with_images, image_map)337 logger.info(f"Extracted {len(image_map)} images from markdown.")338 339 # Build header index for chapter/section/subsection lookup340 header_index = _build_header_index(updated_markdown)341 342 # Use chonkie's RecursiveChunker with markdown recipe343 chunker = RecursiveChunker.from_recipe(344 "markdown",345 lang="en",346 chunk_size=chunk_size,347 )348 chunks = chunker.chunk(updated_markdown)349 350 if not chunks:351 logger.warning("No chunks created. Treating entire text as one chunk.")352 all_refs = list(image_map.keys())353 all_images = [354 decoded355 for uri in image_map.values()356 if (decoded := decode_base64_data_uri(uri)) is not None357 ]358 headers = _get_headers_for_position(header_index, 0)359 return [360 {361 "text": updated_markdown,362 "text_clean": _clean_text(updated_markdown),363 "chapter": headers["chapter"],364 "section": headers["section"],365 "subsection": headers["subsection"],366 "images": all_images,367 "image_refs": all_refs,368 "num_images": len(all_images),369 "has_images": len(all_images) > 0,370 "start_index": 0,371 "char_count": len(updated_markdown),372 }373 ]374 375 result = []376 for chunk in chunks:377 chunk_img_refs = re.findall(r"!\[.*?\]\((img_ref_\d+)\)", chunk.text)378 chunk_images = [379 decoded380 for ref_id in chunk_img_refs381 if ref_id in image_map382 and (decoded := decode_base64_data_uri(image_map[ref_id])) is not None383 ]384 headers = _get_headers_for_position(header_index, chunk.start_index)385 text_clean = _clean_text(chunk.text)386 387 result.append(388 {389 "text": chunk.text,390 "text_clean": text_clean,391 "chapter": headers["chapter"],392 "section": headers["section"],393 "subsection": headers["subsection"],394 "images": chunk_images,395 "image_refs": chunk_img_refs,396 "num_images": len(chunk_images),397 "has_images": len(chunk_images) > 0,398 "start_index": chunk.start_index,399 "char_count": len(chunk.text),400 }401 )402 403 logger.info(f"Created {len(result)} chunks.")404 return result405 406 407### --- Dataset Builder Pipeline ---408 409 410DATASET_SCHEMA = {411 "chunk_id": str,412 "text": str,413 "text_clean": str,414 "chapter": str,415 "section": str,416 "subsection": str,417 "images": list,418 "image_refs": list,419 "num_images": int,420 "has_images": bool,421 "source_filename": str,422 "start_index": int,423 "char_count": int,424}425 426 427@dataclass428class QualityConfig:429 """Configuration for quality filtering thresholds."""430 431 min_char_count: int = 20432 min_clean_char_count: int = 10433 max_char_count: int = 50_000434 min_word_count: int = 3435 max_image_refs_without_text: int = 0436 remove_empty_text: bool = True437 remove_whitespace_only: bool = True438 439 440@dataclass441class PipelineStats:442 """Statistics collected during pipeline execution."""443 444 total_input_chunks: int = 0445 chunks_after_validation: int = 0446 chunks_after_dedup: int = 0447 chunks_after_quality: int = 0448 duplicates_removed: int = 0449 quality_filtered: int = 0450 validation_errors: List[str] = field(default_factory=list)451 quality_reasons: Counter = field(default_factory=Counter)452 source_file_counts: Counter = field(default_factory=Counter)453 avg_char_count: float = 0.0454 avg_images_per_chunk: float = 0.0455 chapters_found: List[str] = field(default_factory=list)456 457 def summary(self) -> str:458 """Generate a human-readable pipeline summary."""459 lines = [460 "--- Dataset Pipeline Report ---",461 f"Input chunks: {self.total_input_chunks}",462 f"After validation: {self.chunks_after_validation}",463 f"Duplicates removed: {self.duplicates_removed}",464 f"After deduplication: {self.chunks_after_dedup}",465 f"Quality filtered out: {self.quality_filtered}",466 f"Final dataset size: {self.chunks_after_quality}",467 "",468 f"Avg chars/chunk: {self.avg_char_count:.0f}",469 f"Avg images/chunk: {self.avg_images_per_chunk:.2f}",470 ]471 472 if self.source_file_counts:473 lines.append("")474 lines.append("Chunks per source file:")475 for fname, count in sorted(self.source_file_counts.items()):476 lines.append(f" {fname}: {count}")477 478 if self.chapters_found:479 unique_chapters = sorted(set(c for c in self.chapters_found if c))480 if unique_chapters:481 lines.append("")482 lines.append(f"Chapters found ({len(unique_chapters)}):")483 for ch in unique_chapters[:20]:484 lines.append(f" - {ch}")485 if len(unique_chapters) > 20:486 lines.append(f" ... and {len(unique_chapters) - 20} more")487 488 if self.quality_reasons:489 lines.append("")490 lines.append("Quality filter reasons:")491 for reason, count in self.quality_reasons.most_common():492 lines.append(f" {reason}: {count}")493 494 if self.validation_errors:495 lines.append("")496 lines.append(f"Validation errors ({len(self.validation_errors)}):")497 for err in self.validation_errors[:10]:498 lines.append(f" - {err}")499 if len(self.validation_errors) > 10:500 lines.append(f" ... and {len(self.validation_errors) - 10} more")501 502 lines.append("-------------------------------")503 return "\n".join(lines)504 505 506class DatasetBuilder:507 """Pipeline for building high-quality datasets before pushing to HF Hub.508 509 Stages:510 1. Validate -- ensure every chunk matches the expected schema511 2. Deduplicate -- remove chunks with identical content hashes512 3. Quality filter -- remove empty, too-short, or malformed chunks513 4. Statistics -- compute summary stats for review514 5. Push -- incremental append or full overwrite to HF Hub515 """516 517 def __init__(518 self,519 quality_config: Optional[QualityConfig] = None,520 ):521 self.quality_config = quality_config or QualityConfig()522 self.stats = PipelineStats()523 self._chunks: List[Dict[str, Any]] = []524 self._seen_hashes: set = set()525 526 def add_chunks(self, chunks: List[Dict[str, Any]], source_filename: str) -> None:527 """Add raw chunks from a processed file into the pipeline.528 529 Each chunk gets its source_filename attached and is tracked for stats.530 """531 for chunk in chunks:532 chunk_with_source = {**chunk, "source_filename": source_filename}533 self._chunks.append(chunk_with_source)534 self.stats.source_file_counts[source_filename] += len(chunks)535 536 def _content_hash(self, chunk: Dict[str, Any]) -> str:537 """Compute a stable hash of chunk content for deduplication."""538 text = chunk.get("text_clean", chunk.get("text", ""))539 source = chunk.get("source_filename", "")540 return hashlib.sha256(f"{source}::{text}".encode("utf-8")).hexdigest()541 542 # --- Stage 1: Validation ---543 544 def _validate(self, chunks: List[Dict[str, Any]]) -> List[Dict[str, Any]]:545 """Validate every chunk conforms to the expected schema.546 547 Drops chunks with missing required fields and logs errors.548 """549 valid = []550 required_keys = set(DATASET_SCHEMA.keys())551 552 for i, chunk in enumerate(chunks):553 missing = required_keys - set(chunk.keys())554 if missing:555 self.stats.validation_errors.append(556 f"Chunk {i} ({chunk.get('chunk_id', '?')}): missing fields {missing}"557 )558 continue559 560 type_ok = True561 for key, expected_type in DATASET_SCHEMA.items():562 val = chunk[key]563 if not isinstance(val, expected_type):564 self.stats.validation_errors.append(565 f"Chunk {i} ({chunk.get('chunk_id', '?')}): "566 f"field '{key}' expected {expected_type.__name__}, "567 f"got {type(val).__name__}"568 )569 type_ok = False570 break571 572 if type_ok:573 valid.append(chunk)574 575 return valid576 577 # --- Stage 2: Deduplication ---578 579 def _deduplicate(self, chunks: List[Dict[str, Any]]) -> List[Dict[str, Any]]:580 """Remove chunks with identical content hashes."""581 unique = []582 for chunk in chunks:583 h = self._content_hash(chunk)584 if h not in self._seen_hashes:585 self._seen_hashes.add(h)586 unique.append(chunk)587 return unique588 589 # --- Stage 3: Quality Filter ---590 591 def _quality_filter(self, chunks: List[Dict[str, Any]]) -> List[Dict[str, Any]]:592 """Filter out low-quality chunks based on configurable thresholds."""593 cfg = self.quality_config594 passed = []595 596 for chunk in chunks:597 text = chunk.get("text", "")598 text_clean = chunk.get("text_clean", "")599 char_count = chunk.get("char_count", len(text))600 601 # Empty text602 if cfg.remove_empty_text and not text.strip():603 self.stats.quality_reasons["empty_text"] += 1604 continue605 606 # Whitespace only607 if cfg.remove_whitespace_only and not text_clean.strip():608 self.stats.quality_reasons["whitespace_only"] += 1609 continue610 611 # Too short612 if char_count < cfg.min_char_count:613 self.stats.quality_reasons[614 f"below_min_chars({cfg.min_char_count})"615 ] += 1616 continue617 618 # Clean text too short619 if len(text_clean.strip()) < cfg.min_clean_char_count:620 self.stats.quality_reasons[621 f"clean_text_too_short({cfg.min_clean_char_count})"622 ] += 1623 continue624 625 # Too long (likely malformed)626 if char_count > cfg.max_char_count:627 self.stats.quality_reasons[628 f"above_max_chars({cfg.max_char_count})"629 ] += 1630 continue631 632 # Too few words633 word_count = len(text_clean.split())634 if word_count < cfg.min_word_count:635 self.stats.quality_reasons[636 f"below_min_words({cfg.min_word_count})"637 ] += 1638 continue639 640 # Image-only chunk with no text641 if (642 cfg.max_image_refs_without_text == 0643 and chunk.get("num_images", 0) > 0644 and word_count == 0645 ):646 self.stats.quality_reasons["image_only_no_text"] += 1647 continue648 649 passed.append(chunk)650 651 return passed652 653 # --- Stage 4: Compute Stats ---654 655 def _compute_stats(self, chunks: List[Dict[str, Any]]) -> None:656 """Compute summary statistics on the final dataset."""657 if not chunks:658 return659 660 total_chars = sum(c.get("char_count", 0) for c in chunks)661 total_images = sum(c.get("num_images", 0) for c in chunks)662 self.stats.avg_char_count = total_chars / len(chunks)663 self.stats.avg_images_per_chunk = total_images / len(chunks)664 self.stats.chapters_found = [c.get("chapter", "") for c in chunks]665 666 # Update source file counts to reflect final dataset667 final_counts: Counter = Counter()668 for c in chunks:669 final_counts[c.get("source_filename", "unknown")] += 1670 self.stats.source_file_counts = final_counts671 672 # --- Run Full Pipeline ---673 674 def build(self) -> tuple[Dict[str, list], PipelineStats]:675 """Run the full pipeline and return columnar data + stats.676 677 Returns:678 Tuple of (columnar_data_dict, pipeline_stats).679 """680 chunks = list(self._chunks)681 self.stats.total_input_chunks = len(chunks)682 logger.info(f"Pipeline: {len(chunks)} input chunks")683 684 # Stage 1: Validate685 chunks = self._validate(chunks)686 self.stats.chunks_after_validation = len(chunks)687 logger.info(f"Pipeline: {len(chunks)} after validation")688 689 # Stage 2: Deduplicate690 before_dedup = len(chunks)691 chunks = self._deduplicate(chunks)692 self.stats.duplicates_removed = before_dedup - len(chunks)693 self.stats.chunks_after_dedup = len(chunks)694 logger.info(695 f"Pipeline: {len(chunks)} after dedup ({self.stats.duplicates_removed} removed)"696 )697 698 # Stage 3: Quality filter699 before_quality = len(chunks)700 chunks = self._quality_filter(chunks)701 self.stats.quality_filtered = before_quality - len(chunks)702 self.stats.chunks_after_quality = len(chunks)703 logger.info(704 f"Pipeline: {len(chunks)} after quality filter ({self.stats.quality_filtered} removed)"705 )706 707 # Stage 4: Stats708 self._compute_stats(chunks)709 710 # Convert to columnar format711 all_data: Dict[str, list] = {key: [] for key in DATASET_SCHEMA.keys()}712 for chunk in chunks:713 for key in DATASET_SCHEMA.keys():714 all_data[key].append(chunk[key])715 716 return all_data, self.stats717 718 # --- Push to Hub ---719 720 @staticmethod721 def push(722 all_data: Dict[str, list],723 repo_name: str,724 hf_token: str,725 stats: PipelineStats,726 append: bool = False,727 ) -> str:728 """Push the built dataset to Hugging Face Hub.729 730 Args:731 all_data: Columnar data dict from build().732 repo_name: HF repo in 'username/dataset-name' format.733 hf_token: Hugging Face API token.734 stats: Pipeline stats for the dataset card.735 append: If True, append to existing dataset instead of overwriting.736 737 Returns:738 Status message string.739 """740 if not all_data or not all_data.get("chunk_id"):741 return "Error: No data to push after pipeline."742 743 api = HfApi(token=hf_token)744 745 try:746 user_info = api.whoami()747 logger.info(f"Authenticated as: {user_info['name']}")748 except Exception as auth_err:749 return f"Error: Invalid HF token - authentication failed: {auth_err}"750 751 # Create repo if needed752 try:753 api.repo_info(repo_id=repo_name, repo_type="dataset")754 logger.info(f"Repository '{repo_name}' exists.")755 except huggingface_hub.utils.RepositoryNotFoundError:756 api.create_repo(repo_id=repo_name, repo_type="dataset", private=False)757 logger.info(f"Created repository '{repo_name}'.")758 759 if append:760 # Incremental append: load existing, concatenate, push761 try:762 existing_ds = load_dataset(repo_name, token=hf_token, split="train")763 existing_data = existing_ds.to_dict()764 for key in all_data:765 if key in existing_data:766 existing_data[key].extend(all_data[key])767 else:768 existing_data[key] = all_data[key]769 # Deduplicate by chunk_id across old + new770 seen_ids = set()771 deduped: Dict[str, list] = {key: [] for key in existing_data}772 for i, cid in enumerate(existing_data["chunk_id"]):773 if cid not in seen_ids:774 seen_ids.add(cid)775 for key in existing_data:776 deduped[key].append(existing_data[key][i])777 merged_dataset = Dataset.from_dict(deduped)778 total_chunks = len(deduped["chunk_id"])779 new_chunks = total_chunks - len(existing_ds)780 commit_msg = f"Append {new_chunks} new chunks (total: {total_chunks})"781 except Exception as e:782 logger.warning(783 f"Could not load existing dataset for append, doing full push: {e}"784 )785 merged_dataset = Dataset.from_dict(all_data)786 total_chunks = len(all_data["chunk_id"])787 commit_msg = f"Add {total_chunks} chunks"788 else:789 merged_dataset = Dataset.from_dict(all_data)790 total_chunks = len(all_data["chunk_id"])791 commit_msg = f"Add OCR data: {total_chunks} chunks"792 793 # Cast the images column so the HF Dataset Viewer renders actual images794 # instead of showing raw base64 strings795 try:796 merged_dataset = merged_dataset.cast_column("images", Sequence(HFImage()))797 logger.info(798 "Cast 'images' column to Sequence(Image()) for viewer rendering."799 )800 except Exception as e:801 logger.warning(f"Could not cast images column to Image feature: {e}")802 803 merged_dataset.push_to_hub(804 repo_name,805 token=hf_token,806 commit_message=commit_msg,807 )808 809 # Generate and upload dataset card810 card_content = DatasetBuilder._generate_dataset_card(repo_name, stats, all_data)811 try:812 api.upload_file(813 path_or_fileobj=card_content.encode("utf-8"),814 path_in_repo="README.md",815 repo_id=repo_name,816 repo_type="dataset",817 commit_message="Update dataset card with pipeline stats",818 )819 except Exception as e:820 logger.warning(f"Failed to update dataset card: {e}")821 822 repo_url = f"https://huggingface.co/datasets/{repo_name}"823 return f"Success! {total_chunks} chunks pushed to: {repo_url}"824 825 @staticmethod826 def _generate_dataset_card(827 repo_name: str,828 stats: PipelineStats,829 all_data: Dict[str, list],830 ) -> str:831 """Generate a dataset card (README.md) with schema and stats."""832 total = stats.chunks_after_quality833 sources = sorted(stats.source_file_counts.items())834 unique_chapters = sorted(set(c for c in stats.chapters_found if c))835 836 card = f"""---837license: mit838task_categories:839 - text-generation840 - question-answering841language:842 - en843tags:844 - pdf2dataset845 - ocr846 - chunked847size_categories:848 - {"1K<n<10K" if total >= 1000 else "n<1K"}849---850 851# {repo_name.split("/")[-1]}852 853Dataset created with [PDF2Dataset](https://github.com/svngoku/PDF2Dataset) -- OCR + structure-aware chunking pipeline.854 855## Dataset Summary856 857| Metric | Value |858|---|---|859| Total chunks | {total} |860| Avg chars/chunk | {stats.avg_char_count:.0f} |861| Avg images/chunk | {stats.avg_images_per_chunk:.2f} |862| Source files | {len(sources)} |863| Duplicates removed | {stats.duplicates_removed} |864| Quality filtered | {stats.quality_filtered} |865 866## Schema867 868| Column | Type | Description |869|---|---|---|870| `chunk_id` | `string` | Unique identifier: `filename_chunk_N` |871| `text` | `string` | Raw markdown chunk with image refs |872| `text_clean` | `string` | Cleaned text without markdown formatting |873| `chapter` | `string` | H1 header active at chunk position |874| `section` | `string` | H2 header active at chunk position |875| `subsection` | `string` | H3+ header active at chunk position |876| `images` | `list[Image]` | Rendered images extracted from chunk (viewable in Dataset Viewer) |877| `image_refs` | `list[string]` | Image reference IDs in chunk text |878| `num_images` | `int` | Number of images in chunk |879| `has_images` | `bool` | Whether chunk contains images |880| `source_filename` | `string` | Original source file name |881| `start_index` | `int` | Character offset in source document |882| `char_count` | `int` | Character count of chunk text |883 884## Source Files885 886| File | Chunks |887|---|---|888"""889 for fname, count in sources:890 card += f"| `{fname}` | {count} |\n"891 892 if unique_chapters:893 card += "\n## Document Structure\n\n"894 card += "Chapters found in the source documents:\n\n"895 for ch in unique_chapters[:30]:896 card += f"- {ch}\n"897 if len(unique_chapters) > 30:898 card += f"- ... and {len(unique_chapters) - 30} more\n"899 900 card += """901## Pipeline902 903This dataset was processed through the PDF2Dataset pipeline:904 9051. **OCR** -- Mistral OCR extracts text and images from PDF/image files9062. **Chunking** -- Structure-aware recursive splitting preserves document hierarchy9073. **Validation** -- Schema validation ensures every chunk has required fields9084. **Deduplication** -- Content-hash based dedup removes identical chunks9095. **Quality Filtering** -- Removes empty, too-short, or malformed chunks910"""911 return card912 913 914def get_hf_token(explicit_token: str | None = None) -> str | None:915 """Retrieve Hugging Face token with fallback mechanisms."""916 if explicit_token and explicit_token.strip() and explicit_token.startswith("hf_"):917 return explicit_token.strip()918 919 env_token = os.environ.get("HF_TOKEN")920 if env_token and env_token.startswith("hf_"):921 return env_token922 923 try:924 stored_token = huggingface_hub.get_token()925 if stored_token:926 return stored_token927 except Exception as e:928 logger.warning(f"Could not retrieve token from Hugging Face config: {e}")929 930 return None931 932 933def process_files(934 file_paths: list[str],935 chunk_size: int,936 hf_token: str,937 repo_name: str,938 append_mode: bool = False,939 min_chunk_chars: int = 20,940 min_words: int = 3,941) -> str:942 """Orchestrates OCR, chunking, pipeline processing, and push to HF Hub.943 944 Pipeline stages:945 1. OCR each file with Mistral946 2. Chunk markdown with structure-aware splitting947 3. Validate schema on every chunk948 4. Deduplicate by content hash949 5. Quality-filter (min chars, min words, empty, etc.)950 6. Compute statistics and generate report951 7. Push to HF Hub (overwrite or append)952 953 Args:954 file_paths: List of file paths to process.955 chunk_size: Maximum character count per chunk.956 hf_token: Explicit HF token (optional).957 repo_name: HF dataset repository in 'username/dataset-name' format.958 append_mode: If True, append to existing dataset instead of replacing.959 min_chunk_chars: Minimum characters per chunk for quality filter.960 min_words: Minimum words per chunk for quality filter.961 962 Returns:963 Status message string with pipeline report.964 """965 if not file_paths:966 return "Error: No files uploaded."967 968 if not repo_name or "/" not in repo_name:969 return "Error: Invalid repository name (use 'username/dataset-name')."970 971 chunk_size = max(0, chunk_size)972 973 effective_hf_token = get_hf_token(hf_token)974 if not effective_hf_token:975 return (976 "Error: No valid Hugging Face token found.\n"977 "Please either:\n"978 "1. Provide a token in the input field (starts with 'hf_')\n"979 "2. Set HF_TOKEN environment variable\n"980 "3. Run `huggingface-cli login` in your terminal"981 )982 983 try:984 # Initialize pipeline with quality config985 quality_cfg = QualityConfig(986 min_char_count=min_chunk_chars,987 min_word_count=min_words,988 )989 builder = DatasetBuilder(quality_config=quality_cfg)990 991 files_processed = 0992 error_messages = []993 994 for file_idx, file_path in enumerate(file_paths, 1):995 source_filename = os.path.basename(file_path)996 logger.info(997 f"--- Processing file {file_idx}/{len(file_paths)}: {source_filename} ---"998 )999 1000 try:1001 processed_markdown, raw_markdown, img_map = perform_ocr(file_path)1002 except OCRError as e:1003 error_messages.append(f"File '{source_filename}': {e}")1004 logger.error(f"Failed to process file {source_filename}: {e}")1005 continue1006 1007 chunks = chunk_markdown(processed_markdown, chunk_size)1008 if not chunks:1009 error_messages.append(1010 f"File '{source_filename}': Failed to chunk the document."1011 )1012 logger.error(f"Failed to chunk file {source_filename}")1013 continue1014 1015 # Assign chunk_id before adding to pipeline1016 for i, chunk in enumerate(chunks):1017 chunk["chunk_id"] = f"{source_filename}_chunk_{i}"1018 1019 builder.add_chunks(chunks, source_filename)1020 files_processed += 11021 logger.info(1022 f"File {source_filename}: queued {len(chunks)} chunks for pipeline"1023 )1024 1025 if files_processed == 0:1026 return "Error: No files were processed successfully.\n" + "\n".join(1027 error_messages1028 )1029 1030 # Run the pipeline1031 all_data, stats = builder.build()1032 1033 if not all_data or not all_data.get("chunk_id"):1034 return (1035 "Error: All chunks were filtered out by the pipeline.\n"1036 + stats.summary()1037 + (1038 "\n\nOCR Errors:\n" + "\n".join(error_messages)1039 if error_messages1040 else ""1041 )1042 )1043 1044 # Push to Hub1045 push_result = DatasetBuilder.push(1046 all_data=all_data,1047 repo_name=repo_name,1048 hf_token=effective_hf_token,1049 stats=stats,1050 append=append_mode,1051 )1052 1053 # Build final report1054 report_parts = [push_result, "", stats.summary()]1055 if error_messages:1056 report_parts.append(f"\nOCR Errors ({len(error_messages)}):")1057 report_parts.extend(f" - {e}" for e in error_messages)1058 1059 return "\n".join(report_parts)1060 1061 except huggingface_hub.utils.HfHubHTTPError as hf_http_err:1062 status = getattr(hf_http_err.response, "status_code", "Unknown")1063 if status == 401:1064 return "Error: Invalid or unauthorized Hugging Face token."1065 elif status == 403:1066 return "Error: Token lacks write permission."1067 return f"Error: Hugging Face Hub Error (Status {status}): {hf_http_err}"1068 except Exception as e:1069 logger.error(f"Unexpected error: {e}", exc_info=True)1070 return f"Unexpected error: {e}"1071 1072 1073# --- Preview ---1074 1075 1076def render_preview(file_objs) -> list[Image.Image]:1077 """Render uploaded files as preview images.1078 1079 PDFs are rendered page-by-page using PyMuPDF. Images are returned directly.1080 """1081 if not file_objs:1082 return []1083 if not isinstance(file_objs, list):1084 file_objs = [file_objs]1085 1086 images = []1087 for file_obj in file_objs:1088 file_path = file_obj.name if hasattr(file_obj, "name") else str(file_obj)1089 ext = os.path.splitext(file_path)[1].lower()1090 1091 if ext == ".pdf":1092 try:1093 doc = fitz.open(file_path)1094 for page in doc:1095 pix = page.get_pixmap(dpi=150)1096 img = Image.open(io.BytesIO(pix.tobytes("png")))1097 images.append(img)1098 doc.close()1099 except Exception as e:1100 logger.error(f"Failed to render PDF preview: {e}")1101 elif ext in {".png", ".jpg", ".jpeg", ".webp", ".bmp"}:1102 try:1103 images.append(Image.open(file_path))1104 except Exception as e:1105 logger.error(f"Failed to open image preview: {e}")1106 1107 return images1108 1109 1110# --- Gradio Interface ---1111 1112 1113MISTRAL_CSS = """1114@import url('https://fonts.googleapis.com/css2?family=JetBrains+Mono:wght@400;500;600&family=Inter:wght@400;500;600;700;800&display=swap');1115 1116:root {1117 --mistral-bg: #FFFAEB;1118 --mistral-bg-grid: #E9E2CB;1119 --mistral-panel: #FFFAEB;1120 --mistral-panel-warm: #FFF0C3;1121 --mistral-border: #E9E2CB;1122 --mistral-text: #1E1E1E;1123 --mistral-muted: #444444;1124 --mistral-soft-muted: #766B54;1125 --mistral-accent: #FF8205;1126 --mistral-accent-hover: #E67200;1127 --mistral-shadow: rgba(30, 30, 30, 0.08);1128 --mistral-grid-opacity: 0.05;1129}1130 1131html,1132body,1133gradio-app,1134.gradio-container,1135.app,1136main {1137 background-color: var(--mistral-bg) !important;1138 background-image:1139 linear-gradient(var(--mistral-bg-grid) 1px, transparent 1px),1140 linear-gradient(90deg, var(--mistral-bg-grid) 1px, transparent 1px) !important;1141 background-size: 40px 40px !important;1142 color: var(--mistral-text) !important;1143 font-family: 'Inter', sans-serif !important;1144}1145 1146html,1147body,1148gradio-app {1149 min-height: 100% !important;1150 width: 100% !important;1151}1152 1153.gradio-container {1154 box-sizing: border-box !important;1155 margin: 0 auto !important;1156 max-width: 1440px !important;1157 padding: clamp(1rem, 2.5vw, 2rem) !important;1158 width: 100% !important;1159}1160 1161.app,1162main {1163 min-height: 100vh !important;1164 max-width: 100% !important;1165 width: 100% !important;1166}1167 1168footer {1169 display: none !important;1170}1171 1172#app-shell {1173 gap: 1rem !important;1174}1175 1176#brand-hero {1177 background: linear-gradient(135deg, var(--mistral-panel) 0%, var(--mistral-panel-warm) 100%) !important;1178 border: 2px solid var(--mistral-border) !important;1179 border-top: 5px solid var(--mistral-accent) !important;1180 box-shadow: 0 8px 32px var(--mistral-shadow) !important;1181 padding: clamp(1.25rem, 3vw, 2rem) !important;1182}1183 1184#brand-hero h1 {1185 color: var(--mistral-text) !important;1186 font-size: clamp(2rem, 4vw, 4.5rem) !important;1187 font-weight: 800 !important;1188 line-height: 0.95 !important;1189 letter-spacing: 0 !important;1190 margin: 0 0 0.75rem !important;1191}1192 1193#brand-hero p {1194 color: var(--mistral-muted) !important;1195 font-size: clamp(1rem, 1.6vw, 1.25rem) !important;1196 line-height: 1.55 !important;1197 margin: 0 !important;1198 max-width: 62rem !important;1199}1200