from dotenv import load_dotenv load_dotenv() import gradio as gr from chonkie import RecursiveChunker from typing import Dict, Any, List, Optional from dataclasses import dataclass, field import logging import re import base64 import hashlib import mimetypes import json from collections import Counter from datasets import Dataset, Features, Value, Sequence, load_dataset from datasets.features import Image as HFImage from huggingface_hub import HfApi, get_token import huggingface_hub import os from mistralai import Mistral import fitz # pymupdf from PIL import Image import io import tempfile # Configure logging logging.basicConfig( level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s" ) logger = logging.getLogger(__name__) # --- Exceptions --- class OCRError(Exception): """Raised when OCR processing fails.""" pass # --- Mistral Client (lazy init) --- _client: Mistral | None = None def get_mistral_client() -> Mistral: """Get or initialize the Mistral client.""" global _client if _client is not None: return _client api_key = os.environ.get("MISTRAL_API_KEY") if not api_key: logger.warning("MISTRAL_API_KEY not set. Attempting to use Hugging Face token.") api_key = get_token() if api_key: logger.info("Using Hugging Face token as MISTRAL_API_KEY.") if not api_key: raise OCRError( "No API key found. Set MISTRAL_API_KEY or run `huggingface-cli login`." ) _client = Mistral(api_key=api_key) logger.info("Mistral client initialized successfully.") return _client # --- Helper Functions --- def encode_image_bytes(image_bytes: bytes) -> str: """Encodes image bytes to a base64 string.""" return base64.b64encode(image_bytes).decode("utf-8") def decode_base64_data_uri(data_uri: str) -> Optional[dict]: """Decode a base64 data URI to a HF-compatible image bytes dict. Args: data_uri: A string like "data:image/jpeg;base64,/9j/4AAQ..." or raw base64. Returns: Dict with {"bytes": , "path": None} for datasets.Image feature, or None if decoding fails. """ try: if data_uri.startswith("data:"): # Strip the "data:image/...;base64," prefix _, encoded = data_uri.split(",", 1) else: encoded = data_uri raw_bytes = base64.b64decode(encoded) # Validate it's a real image by opening it img = Image.open(io.BytesIO(raw_bytes)) # Re-encode as PNG for consistency buf = io.BytesIO() img.save(buf, format="PNG") return {"bytes": buf.getvalue(), "path": None} except Exception as e: logger.warning(f"Failed to decode base64 image ({len(data_uri)} chars): {e}") return None def extract_images_from_markdown(markdown_text: str) -> Dict[str, str]: """ Extracts base64 image data URIs from markdown and maps them to reference IDs. Returns a dictionary mapping reference IDs to base64 data URIs. """ image_map = {} img_refs = re.findall( r"!\[.*?\]\((data:image/[a-zA-Z+]+;base64,[A-Za-z0-9+/=]+)\)", markdown_text ) for idx, img_uri in enumerate(img_refs): ref_id = f"img_ref_{idx + 1}" image_map[ref_id] = img_uri return image_map def replace_image_references(markdown_text: str, image_map: Dict[str, str]) -> str: """ Replaces base64 image data URIs in markdown with reference IDs (e.g., img_ref_1). """ updated_markdown = markdown_text for ref_id, img_uri in image_map.items(): escaped_uri = re.escape(img_uri) pattern = r"(!\[.*?\]\()" + escaped_uri + r"(\))" updated_markdown = re.sub(pattern, f"\\1{ref_id}\\2", updated_markdown) return updated_markdown def get_combined_markdown(ocr_response: Any) -> tuple[str, str, Dict[str, str]]: """Combines markdown from OCR pages, replacing image IDs with base64 data URIs.""" processed_markdowns = [] raw_markdowns = [] image_data_map = {} if not hasattr(ocr_response, "pages") or not ocr_response.pages: logger.warning("OCR response has no pages.") return "", "", {} for page_idx, page in enumerate(ocr_response.pages): if hasattr(page, "images") and page.images: logger.info(f"Page {page_idx}: Found {len(page.images)} images.") for img in page.images: if ( hasattr(img, "id") and hasattr(img, "image_base64") and img.image_base64 ): image_data_map[img.id] = img.image_base64 else: logger.warning( f"Page {page_idx}: Image object lacks 'id' or valid 'image_base64'." ) else: logger.info(f"Page {page_idx}: No images found.") if not hasattr(page, "markdown"): logger.warning(f"Page {page_idx} lacks 'markdown' attribute. Skipping.") continue current_raw_markdown = page.markdown or "" raw_markdowns.append(current_raw_markdown) current_processed_markdown = current_raw_markdown img_refs = re.findall(r"!\[.*?\]\((.*?)\)", current_processed_markdown) for img_id in img_refs: if img_id in image_data_map: base64_data_uri = image_data_map[img_id] escaped_img_id = re.escape(img_id) pattern = r"(!\[.*?\]\()" + escaped_img_id + r"(\))" current_processed_markdown = re.sub( pattern, r"\1" + base64_data_uri + r"\2", current_processed_markdown, ) elif not img_id.startswith(("http:", "https:", "data:")): logger.warning( f"Page {page_idx}: Image ID '{img_id}' not in image data." ) processed_markdowns.append(current_processed_markdown) logger.info( f"Processed {len(processed_markdowns)} pages with {len(image_data_map)} images." ) return "\n\n".join(processed_markdowns), "\n\n".join(raw_markdowns), image_data_map def perform_ocr(file_path: str) -> tuple[str, str, Dict[str, str]]: """Performs OCR on a file using Mistral API. Args: file_path: Path to the file on disk. Returns: Tuple of (processed_markdown, raw_markdown, image_data_map). Raises: OCRError: If OCR processing fails. """ client = get_mistral_client() file_name = os.path.basename(file_path) file_ext = os.path.splitext(file_name)[1].lower() logger.info(f"Performing OCR on file: {file_name}") ocr_response = None supported_images = {".png", ".jpg", ".jpeg", ".webp", ".bmp"} if file_ext == ".pdf": uploaded_file_id = None try: with open(file_path, "rb") as f: file_content = f.read() logger.info(f"Uploading PDF {file_name} to Mistral...") uploaded_pdf = client.files.upload( file={"file_name": file_name, "content": file_content}, purpose="ocr", ) uploaded_file_id = uploaded_pdf.id logger.info(f"PDF uploaded. File ID: {uploaded_file_id}") signed_url_response = client.files.get_signed_url(file_id=uploaded_file_id) ocr_response = client.ocr.process( model="mistral-ocr-latest", document={ "type": "document_url", "document_url": signed_url_response.url, }, include_image_base64=True, ) finally: if uploaded_file_id: try: client.files.delete(file_id=uploaded_file_id) except Exception as delete_err: logger.warning( f"Failed to delete temporary file {uploaded_file_id}: {delete_err}" ) elif file_ext in supported_images: with open(file_path, "rb") as f: image_bytes = f.read() if not image_bytes: raise OCRError(f"Uploaded image file '{file_name}' is empty.") base64_encoded = encode_image_bytes(image_bytes) mime_type, _ = mimetypes.guess_type(file_path) mime_type = mime_type or "image/jpeg" data_uri = f"data:{mime_type};base64,{base64_encoded}" ocr_response = client.ocr.process( model="mistral-ocr-latest", document={"type": "image_url", "image_url": data_uri}, include_image_base64=True, ) else: raise OCRError(f"Unsupported file type: '{file_ext}'") if not ocr_response: raise OCRError(f"OCR returned no response for '{file_name}'.") processed_md, raw_md, img_map = get_combined_markdown(ocr_response) logger.info(f"Processed markdown length: {len(processed_md)}") return processed_md, raw_md, img_map def _build_header_index(markdown_text: str) -> list[tuple[int, int, str]]: """Build a sorted index of (position, level, title) for all markdown headers.""" headers = [] for match in re.finditer(r"^(#{1,6})\s+(.+)$", markdown_text, re.MULTILINE): level = len(match.group(1)) title = match.group(2).strip() headers.append((match.start(), level, title)) return headers def _get_headers_for_position( headers: list[tuple[int, int, str]], position: int ) -> dict[str, str]: """Given a character position, find the active chapter/section/subsection. Maps header levels: H1 -> chapter, H2 -> section, H3+ -> subsection. """ active: dict[int, str] = {} for hdr_pos, level, title in headers: if hdr_pos > position: break active[level] = title # Clear deeper levels when a higher-level header appears for deeper in list(active.keys()): if deeper > level: del active[deeper] return { "chapter": active.get(1, ""), "section": active.get(2, ""), "subsection": active.get(3, active.get(4, active.get(5, active.get(6, "")))), } def _clean_text(text: str) -> str: """Remove markdown formatting, image refs, and extra whitespace.""" cleaned = re.sub(r"!\[.*?\]\(.*?\)", "", text) cleaned = re.sub(r"#{1,6}\s+", "", cleaned) cleaned = re.sub(r"\*\*(.+?)\*\*", r"\1", cleaned) cleaned = re.sub(r"\*(.+?)\*", r"\1", cleaned) cleaned = re.sub(r"`(.+?)`", r"\1", cleaned) cleaned = re.sub(r"\[(.+?)\]\(.*?\)", r"\1", cleaned) cleaned = re.sub(r"\n{3,}", "\n\n", cleaned) return cleaned.strip() def chunk_markdown( markdown_text_with_images: str, chunk_size: int = 512, ) -> list[dict]: """Chunks markdown text using chonkie's RecursiveChunker with markdown recipe. Args: markdown_text_with_images: Markdown text possibly containing base64 image references. chunk_size: Maximum character count per chunk. Returns: List of chunk dicts with the full dataset schema fields. """ if not markdown_text_with_images or not markdown_text_with_images.strip(): logger.warning("chunk_markdown received empty input.") return [] # Extract images and replace with reference IDs image_map = extract_images_from_markdown(markdown_text_with_images) updated_markdown = replace_image_references(markdown_text_with_images, image_map) logger.info(f"Extracted {len(image_map)} images from markdown.") # Build header index for chapter/section/subsection lookup header_index = _build_header_index(updated_markdown) # Use chonkie's RecursiveChunker with markdown recipe chunker = RecursiveChunker.from_recipe( "markdown", lang="en", chunk_size=chunk_size, ) chunks = chunker.chunk(updated_markdown) if not chunks: logger.warning("No chunks created. Treating entire text as one chunk.") all_refs = list(image_map.keys()) all_images = [ decoded for uri in image_map.values() if (decoded := decode_base64_data_uri(uri)) is not None ] headers = _get_headers_for_position(header_index, 0) return [ { "text": updated_markdown, "text_clean": _clean_text(updated_markdown), "chapter": headers["chapter"], "section": headers["section"], "subsection": headers["subsection"], "images": all_images, "image_refs": all_refs, "num_images": len(all_images), "has_images": len(all_images) > 0, "start_index": 0, "char_count": len(updated_markdown), } ] result = [] for chunk in chunks: chunk_img_refs = re.findall(r"!\[.*?\]\((img_ref_\d+)\)", chunk.text) chunk_images = [ decoded for ref_id in chunk_img_refs if ref_id in image_map and (decoded := decode_base64_data_uri(image_map[ref_id])) is not None ] headers = _get_headers_for_position(header_index, chunk.start_index) text_clean = _clean_text(chunk.text) result.append( { "text": chunk.text, "text_clean": text_clean, "chapter": headers["chapter"], "section": headers["section"], "subsection": headers["subsection"], "images": chunk_images, "image_refs": chunk_img_refs, "num_images": len(chunk_images), "has_images": len(chunk_images) > 0, "start_index": chunk.start_index, "char_count": len(chunk.text), } ) logger.info(f"Created {len(result)} chunks.") return result ### --- Dataset Builder Pipeline --- DATASET_SCHEMA = { "chunk_id": str, "text": str, "text_clean": str, "chapter": str, "section": str, "subsection": str, "images": list, "image_refs": list, "num_images": int, "has_images": bool, "source_filename": str, "start_index": int, "char_count": int, } @dataclass class QualityConfig: """Configuration for quality filtering thresholds.""" min_char_count: int = 20 min_clean_char_count: int = 10 max_char_count: int = 50_000 min_word_count: int = 3 max_image_refs_without_text: int = 0 remove_empty_text: bool = True remove_whitespace_only: bool = True @dataclass class PipelineStats: """Statistics collected during pipeline execution.""" total_input_chunks: int = 0 chunks_after_validation: int = 0 chunks_after_dedup: int = 0 chunks_after_quality: int = 0 duplicates_removed: int = 0 quality_filtered: int = 0 validation_errors: List[str] = field(default_factory=list) quality_reasons: Counter = field(default_factory=Counter) source_file_counts: Counter = field(default_factory=Counter) avg_char_count: float = 0.0 avg_images_per_chunk: float = 0.0 chapters_found: List[str] = field(default_factory=list) def summary(self) -> str: """Generate a human-readable pipeline summary.""" lines = [ "--- Dataset Pipeline Report ---", f"Input chunks: {self.total_input_chunks}", f"After validation: {self.chunks_after_validation}", f"Duplicates removed: {self.duplicates_removed}", f"After deduplication: {self.chunks_after_dedup}", f"Quality filtered out: {self.quality_filtered}", f"Final dataset size: {self.chunks_after_quality}", "", f"Avg chars/chunk: {self.avg_char_count:.0f}", f"Avg images/chunk: {self.avg_images_per_chunk:.2f}", ] if self.source_file_counts: lines.append("") lines.append("Chunks per source file:") for fname, count in sorted(self.source_file_counts.items()): lines.append(f" {fname}: {count}") if self.chapters_found: unique_chapters = sorted(set(c for c in self.chapters_found if c)) if unique_chapters: lines.append("") lines.append(f"Chapters found ({len(unique_chapters)}):") for ch in unique_chapters[:20]: lines.append(f" - {ch}") if len(unique_chapters) > 20: lines.append(f" ... and {len(unique_chapters) - 20} more") if self.quality_reasons: lines.append("") lines.append("Quality filter reasons:") for reason, count in self.quality_reasons.most_common(): lines.append(f" {reason}: {count}") if self.validation_errors: lines.append("") lines.append(f"Validation errors ({len(self.validation_errors)}):") for err in self.validation_errors[:10]: lines.append(f" - {err}") if len(self.validation_errors) > 10: lines.append(f" ... and {len(self.validation_errors) - 10} more") lines.append("-------------------------------") return "\n".join(lines) class DatasetBuilder: """Pipeline for building high-quality datasets before pushing to HF Hub. Stages: 1. Validate -- ensure every chunk matches the expected schema 2. Deduplicate -- remove chunks with identical content hashes 3. Quality filter -- remove empty, too-short, or malformed chunks 4. Statistics -- compute summary stats for review 5. Push -- incremental append or full overwrite to HF Hub """ def __init__( self, quality_config: Optional[QualityConfig] = None, ): self.quality_config = quality_config or QualityConfig() self.stats = PipelineStats() self._chunks: List[Dict[str, Any]] = [] self._seen_hashes: set = set() def add_chunks(self, chunks: List[Dict[str, Any]], source_filename: str) -> None: """Add raw chunks from a processed file into the pipeline. Each chunk gets its source_filename attached and is tracked for stats. """ for chunk in chunks: chunk_with_source = {**chunk, "source_filename": source_filename} self._chunks.append(chunk_with_source) self.stats.source_file_counts[source_filename] += len(chunks) def _content_hash(self, chunk: Dict[str, Any]) -> str: """Compute a stable hash of chunk content for deduplication.""" text = chunk.get("text_clean", chunk.get("text", "")) source = chunk.get("source_filename", "") return hashlib.sha256(f"{source}::{text}".encode("utf-8")).hexdigest() # --- Stage 1: Validation --- def _validate(self, chunks: List[Dict[str, Any]]) -> List[Dict[str, Any]]: """Validate every chunk conforms to the expected schema. Drops chunks with missing required fields and logs errors. """ valid = [] required_keys = set(DATASET_SCHEMA.keys()) for i, chunk in enumerate(chunks): missing = required_keys - set(chunk.keys()) if missing: self.stats.validation_errors.append( f"Chunk {i} ({chunk.get('chunk_id', '?')}): missing fields {missing}" ) continue type_ok = True for key, expected_type in DATASET_SCHEMA.items(): val = chunk[key] if not isinstance(val, expected_type): self.stats.validation_errors.append( f"Chunk {i} ({chunk.get('chunk_id', '?')}): " f"field '{key}' expected {expected_type.__name__}, " f"got {type(val).__name__}" ) type_ok = False break if type_ok: valid.append(chunk) return valid # --- Stage 2: Deduplication --- def _deduplicate(self, chunks: List[Dict[str, Any]]) -> List[Dict[str, Any]]: """Remove chunks with identical content hashes.""" unique = [] for chunk in chunks: h = self._content_hash(chunk) if h not in self._seen_hashes: self._seen_hashes.add(h) unique.append(chunk) return unique # --- Stage 3: Quality Filter --- def _quality_filter(self, chunks: List[Dict[str, Any]]) -> List[Dict[str, Any]]: """Filter out low-quality chunks based on configurable thresholds.""" cfg = self.quality_config passed = [] for chunk in chunks: text = chunk.get("text", "") text_clean = chunk.get("text_clean", "") char_count = chunk.get("char_count", len(text)) # Empty text if cfg.remove_empty_text and not text.strip(): self.stats.quality_reasons["empty_text"] += 1 continue # Whitespace only if cfg.remove_whitespace_only and not text_clean.strip(): self.stats.quality_reasons["whitespace_only"] += 1 continue # Too short if char_count < cfg.min_char_count: self.stats.quality_reasons[ f"below_min_chars({cfg.min_char_count})" ] += 1 continue # Clean text too short if len(text_clean.strip()) < cfg.min_clean_char_count: self.stats.quality_reasons[ f"clean_text_too_short({cfg.min_clean_char_count})" ] += 1 continue # Too long (likely malformed) if char_count > cfg.max_char_count: self.stats.quality_reasons[ f"above_max_chars({cfg.max_char_count})" ] += 1 continue # Too few words word_count = len(text_clean.split()) if word_count < cfg.min_word_count: self.stats.quality_reasons[ f"below_min_words({cfg.min_word_count})" ] += 1 continue # Image-only chunk with no text if ( cfg.max_image_refs_without_text == 0 and chunk.get("num_images", 0) > 0 and word_count == 0 ): self.stats.quality_reasons["image_only_no_text"] += 1 continue passed.append(chunk) return passed # --- Stage 4: Compute Stats --- def _compute_stats(self, chunks: List[Dict[str, Any]]) -> None: """Compute summary statistics on the final dataset.""" if not chunks: return total_chars = sum(c.get("char_count", 0) for c in chunks) total_images = sum(c.get("num_images", 0) for c in chunks) self.stats.avg_char_count = total_chars / len(chunks) self.stats.avg_images_per_chunk = total_images / len(chunks) self.stats.chapters_found = [c.get("chapter", "") for c in chunks] # Update source file counts to reflect final dataset final_counts: Counter = Counter() for c in chunks: final_counts[c.get("source_filename", "unknown")] += 1 self.stats.source_file_counts = final_counts # --- Run Full Pipeline --- def build(self) -> tuple[Dict[str, list], PipelineStats]: """Run the full pipeline and return columnar data + stats. Returns: Tuple of (columnar_data_dict, pipeline_stats). """ chunks = list(self._chunks) self.stats.total_input_chunks = len(chunks) logger.info(f"Pipeline: {len(chunks)} input chunks") # Stage 1: Validate chunks = self._validate(chunks) self.stats.chunks_after_validation = len(chunks) logger.info(f"Pipeline: {len(chunks)} after validation") # Stage 2: Deduplicate before_dedup = len(chunks) chunks = self._deduplicate(chunks) self.stats.duplicates_removed = before_dedup - len(chunks) self.stats.chunks_after_dedup = len(chunks) logger.info( f"Pipeline: {len(chunks)} after dedup ({self.stats.duplicates_removed} removed)" ) # Stage 3: Quality filter before_quality = len(chunks) chunks = self._quality_filter(chunks) self.stats.quality_filtered = before_quality - len(chunks) self.stats.chunks_after_quality = len(chunks) logger.info( f"Pipeline: {len(chunks)} after quality filter ({self.stats.quality_filtered} removed)" ) # Stage 4: Stats self._compute_stats(chunks) # Convert to columnar format all_data: Dict[str, list] = {key: [] for key in DATASET_SCHEMA.keys()} for chunk in chunks: for key in DATASET_SCHEMA.keys(): all_data[key].append(chunk[key]) return all_data, self.stats # --- Push to Hub --- @staticmethod def push( all_data: Dict[str, list], repo_name: str, hf_token: str, stats: PipelineStats, append: bool = False, ) -> str: """Push the built dataset to Hugging Face Hub. Args: all_data: Columnar data dict from build(). repo_name: HF repo in 'username/dataset-name' format. hf_token: Hugging Face API token. stats: Pipeline stats for the dataset card. append: If True, append to existing dataset instead of overwriting. Returns: Status message string. """ if not all_data or not all_data.get("chunk_id"): return "Error: No data to push after pipeline." api = HfApi(token=hf_token) try: user_info = api.whoami() logger.info(f"Authenticated as: {user_info['name']}") except Exception as auth_err: return f"Error: Invalid HF token - authentication failed: {auth_err}" # Create repo if needed try: api.repo_info(repo_id=repo_name, repo_type="dataset") logger.info(f"Repository '{repo_name}' exists.") except huggingface_hub.utils.RepositoryNotFoundError: api.create_repo(repo_id=repo_name, repo_type="dataset", private=False) logger.info(f"Created repository '{repo_name}'.") if append: # Incremental append: load existing, concatenate, push try: existing_ds = load_dataset(repo_name, token=hf_token, split="train") existing_data = existing_ds.to_dict() for key in all_data: if key in existing_data: existing_data[key].extend(all_data[key]) else: existing_data[key] = all_data[key] # Deduplicate by chunk_id across old + new seen_ids = set() deduped: Dict[str, list] = {key: [] for key in existing_data} for i, cid in enumerate(existing_data["chunk_id"]): if cid not in seen_ids: seen_ids.add(cid) for key in existing_data: deduped[key].append(existing_data[key][i]) merged_dataset = Dataset.from_dict(deduped) total_chunks = len(deduped["chunk_id"]) new_chunks = total_chunks - len(existing_ds) commit_msg = f"Append {new_chunks} new chunks (total: {total_chunks})" except Exception as e: logger.warning( f"Could not load existing dataset for append, doing full push: {e}" ) merged_dataset = Dataset.from_dict(all_data) total_chunks = len(all_data["chunk_id"]) commit_msg = f"Add {total_chunks} chunks" else: merged_dataset = Dataset.from_dict(all_data) total_chunks = len(all_data["chunk_id"]) commit_msg = f"Add OCR data: {total_chunks} chunks" # Cast the images column so the HF Dataset Viewer renders actual images # instead of showing raw base64 strings try: merged_dataset = merged_dataset.cast_column("images", Sequence(HFImage())) logger.info( "Cast 'images' column to Sequence(Image()) for viewer rendering." ) except Exception as e: logger.warning(f"Could not cast images column to Image feature: {e}") merged_dataset.push_to_hub( repo_name, token=hf_token, commit_message=commit_msg, ) # Generate and upload dataset card card_content = DatasetBuilder._generate_dataset_card(repo_name, stats, all_data) try: api.upload_file( path_or_fileobj=card_content.encode("utf-8"), path_in_repo="README.md", repo_id=repo_name, repo_type="dataset", commit_message="Update dataset card with pipeline stats", ) except Exception as e: logger.warning(f"Failed to update dataset card: {e}") repo_url = f"https://huggingface.co/datasets/{repo_name}" return f"Success! {total_chunks} chunks pushed to: {repo_url}" @staticmethod def _generate_dataset_card( repo_name: str, stats: PipelineStats, all_data: Dict[str, list], ) -> str: """Generate a dataset card (README.md) with schema and stats.""" total = stats.chunks_after_quality sources = sorted(stats.source_file_counts.items()) unique_chapters = sorted(set(c for c in stats.chapters_found if c)) card = f"""--- license: mit task_categories: - text-generation - question-answering language: - en tags: - pdf2dataset - ocr - chunked size_categories: - {"1K= 1000 else "n<1K"} --- # {repo_name.split("/")[-1]} Dataset created with [PDF2Dataset](https://github.com/svngoku/PDF2Dataset) -- OCR + structure-aware chunking pipeline. ## Dataset Summary | Metric | Value | |---|---| | Total chunks | {total} | | Avg chars/chunk | {stats.avg_char_count:.0f} | | Avg images/chunk | {stats.avg_images_per_chunk:.2f} | | Source files | {len(sources)} | | Duplicates removed | {stats.duplicates_removed} | | Quality filtered | {stats.quality_filtered} | ## Schema | Column | Type | Description | |---|---|---| | `chunk_id` | `string` | Unique identifier: `filename_chunk_N` | | `text` | `string` | Raw markdown chunk with image refs | | `text_clean` | `string` | Cleaned text without markdown formatting | | `chapter` | `string` | H1 header active at chunk position | | `section` | `string` | H2 header active at chunk position | | `subsection` | `string` | H3+ header active at chunk position | | `images` | `list[Image]` | Rendered images extracted from chunk (viewable in Dataset Viewer) | | `image_refs` | `list[string]` | Image reference IDs in chunk text | | `num_images` | `int` | Number of images in chunk | | `has_images` | `bool` | Whether chunk contains images | | `source_filename` | `string` | Original source file name | | `start_index` | `int` | Character offset in source document | | `char_count` | `int` | Character count of chunk text | ## Source Files | File | Chunks | |---|---| """ for fname, count in sources: card += f"| `{fname}` | {count} |\n" if unique_chapters: card += "\n## Document Structure\n\n" card += "Chapters found in the source documents:\n\n" for ch in unique_chapters[:30]: card += f"- {ch}\n" if len(unique_chapters) > 30: card += f"- ... and {len(unique_chapters) - 30} more\n" card += """ ## Pipeline This dataset was processed through the PDF2Dataset pipeline: 1. **OCR** -- Mistral OCR extracts text and images from PDF/image files 2. **Chunking** -- Structure-aware recursive splitting preserves document hierarchy 3. **Validation** -- Schema validation ensures every chunk has required fields 4. **Deduplication** -- Content-hash based dedup removes identical chunks 5. **Quality Filtering** -- Removes empty, too-short, or malformed chunks """ return card def get_hf_token(explicit_token: str | None = None) -> str | None: """Retrieve Hugging Face token with fallback mechanisms.""" if explicit_token and explicit_token.strip() and explicit_token.startswith("hf_"): return explicit_token.strip() env_token = os.environ.get("HF_TOKEN") if env_token and env_token.startswith("hf_"): return env_token try: stored_token = huggingface_hub.get_token() if stored_token: return stored_token except Exception as e: logger.warning(f"Could not retrieve token from Hugging Face config: {e}") return None def process_files( file_paths: list[str], chunk_size: int, hf_token: str | None, repo_name: str, append_mode: bool = False, min_chunk_chars: int = 20, min_words: int = 3, ) -> str: """Orchestrates OCR, chunking, pipeline processing, and push to HF Hub. Pipeline stages: 1. OCR each file with Mistral 2. Chunk markdown with structure-aware splitting 3. Validate schema on every chunk 4. Deduplicate by content hash 5. Quality-filter (min chars, min words, empty, etc.) 6. Compute statistics and generate report 7. Push to HF Hub (overwrite or append) Args: file_paths: List of file paths to process. chunk_size: Maximum character count per chunk. hf_token: Explicit HF token (optional). repo_name: HF dataset repository in 'username/dataset-name' format. append_mode: If True, append to existing dataset instead of replacing. min_chunk_chars: Minimum characters per chunk for quality filter. min_words: Minimum words per chunk for quality filter. Returns: Status message string with pipeline report. """ if not file_paths: return "Error: No files uploaded." if not repo_name or "/" not in repo_name: return "Error: Invalid repository name (use 'username/dataset-name')." chunk_size = max(0, chunk_size) effective_hf_token = get_hf_token(hf_token) if not effective_hf_token: return ( "Error: No valid Hugging Face token found.\n" "Please either:\n" "1. Sign in with Hugging Face when using the Space\n" "2. Set HF_TOKEN environment variable\n" "3. Run `huggingface-cli login` in your terminal" ) try: # Initialize pipeline with quality config quality_cfg = QualityConfig( min_char_count=min_chunk_chars, min_word_count=min_words, ) builder = DatasetBuilder(quality_config=quality_cfg) files_processed = 0 error_messages = [] for file_idx, file_path in enumerate(file_paths, 1): source_filename = os.path.basename(file_path) logger.info( f"--- Processing file {file_idx}/{len(file_paths)}: {source_filename} ---" ) try: processed_markdown, raw_markdown, img_map = perform_ocr(file_path) except OCRError as e: error_messages.append(f"File '{source_filename}': {e}") logger.error(f"Failed to process file {source_filename}: {e}") continue chunks = chunk_markdown(processed_markdown, chunk_size) if not chunks: error_messages.append( f"File '{source_filename}': Failed to chunk the document." ) logger.error(f"Failed to chunk file {source_filename}") continue # Assign chunk_id before adding to pipeline for i, chunk in enumerate(chunks): chunk["chunk_id"] = f"{source_filename}_chunk_{i}" builder.add_chunks(chunks, source_filename) files_processed += 1 logger.info( f"File {source_filename}: queued {len(chunks)} chunks for pipeline" ) if files_processed == 0: return "Error: No files were processed successfully.\n" + "\n".join( error_messages ) # Run the pipeline all_data, stats = builder.build() if not all_data or not all_data.get("chunk_id"): return ( "Error: All chunks were filtered out by the pipeline.\n" + stats.summary() + ( "\n\nOCR Errors:\n" + "\n".join(error_messages) if error_messages else "" ) ) # Push to Hub push_result = DatasetBuilder.push( all_data=all_data, repo_name=repo_name, hf_token=effective_hf_token, stats=stats, append=append_mode, ) # Build final report report_parts = [push_result, "", stats.summary()] if error_messages: report_parts.append(f"\nOCR Errors ({len(error_messages)}):") report_parts.extend(f" - {e}" for e in error_messages) return "\n".join(report_parts) except huggingface_hub.utils.HfHubHTTPError as hf_http_err: status = getattr(hf_http_err.response, "status_code", "Unknown") if status == 401: return "Error: Invalid or unauthorized Hugging Face token." elif status == 403: return "Error: Token lacks write permission." return f"Error: Hugging Face Hub Error (Status {status}): {hf_http_err}" except Exception as e: logger.error(f"Unexpected error: {e}", exc_info=True) return f"Unexpected error: {e}" # --- Preview --- def render_preview(file_objs) -> list[Image.Image]: """Render uploaded files as preview images. PDFs are rendered page-by-page using PyMuPDF. Images are returned directly. """ if not file_objs: return [] if not isinstance(file_objs, list): file_objs = [file_objs] images = [] for file_obj in file_objs: file_path = file_obj.name if hasattr(file_obj, "name") else str(file_obj) ext = os.path.splitext(file_path)[1].lower() if ext == ".pdf": try: doc = fitz.open(file_path) for page in doc: pix = page.get_pixmap(dpi=150) img = Image.open(io.BytesIO(pix.tobytes("png"))) images.append(img) doc.close() except Exception as e: logger.error(f"Failed to render PDF preview: {e}") elif ext in {".png", ".jpg", ".jpeg", ".webp", ".bmp"}: try: images.append(Image.open(file_path)) except Exception as e: logger.error(f"Failed to open image preview: {e}") return images # --- Gradio Interface --- MISTRAL_CSS = """ @import url('https://fonts.googleapis.com/css2?family=JetBrains+Mono:wght@400;500;600&family=Inter:wght@400;500;600;700;800&display=swap'); :root { --mistral-bg: #FFFAEB; --mistral-bg-grid: #E9E2CB; --mistral-panel: #FFFAEB; --mistral-panel-warm: #FFF0C3; --mistral-border: #E9E2CB; --mistral-text: #1E1E1E; --mistral-muted: #444444; --mistral-soft-muted: #766B54; --mistral-accent: #FF8205; --mistral-accent-hover: #E67200; --mistral-shadow: rgba(30, 30, 30, 0.08); --mistral-grid-opacity: 0.05; } html, body, gradio-app, .gradio-container, .app, main { background-color: var(--mistral-bg) !important; background-image: linear-gradient(var(--mistral-bg-grid) 1px, transparent 1px), linear-gradient(90deg, var(--mistral-bg-grid) 1px, transparent 1px) !important; background-size: 40px 40px !important; color: var(--mistral-text) !important; font-family: 'Inter', sans-serif !important; } html, body, gradio-app { min-height: 100% !important; width: 100% !important; } .gradio-container { box-sizing: border-box !important; margin: 0 auto !important; max-width: 1440px !important; padding: clamp(1rem, 2.5vw, 2rem) !important; width: 100% !important; } .app, main { min-height: 100vh !important; max-width: 100% !important; width: 100% !important; } footer { display: none !important; } #app-shell { gap: 1rem !important; } #brand-hero { background: linear-gradient(135deg, var(--mistral-panel) 0%, var(--mistral-panel-warm) 100%) !important; border: 2px solid var(--mistral-border) !important; border-top: 5px solid var(--mistral-accent) !important; box-shadow: 0 8px 32px var(--mistral-shadow) !important; padding: clamp(1.25rem, 3vw, 2rem) !important; } #brand-hero h1 { color: var(--mistral-text) !important; font-size: clamp(2rem, 4vw, 4.5rem) !important; font-weight: 800 !important; line-height: 0.95 !important; letter-spacing: 0 !important; margin: 0 0 0.75rem !important; } #brand-hero p { color: var(--mistral-muted) !important; font-size: clamp(1rem, 1.6vw, 1.25rem) !important; line-height: 1.55 !important; margin: 0 !important; max-width: 62rem !important; } #brand-hero strong { color: var(--mistral-text) !important; font-weight: 700 !important; } .mistral-panel { background: var(--mistral-panel) !important; border: 2px solid var(--mistral-border) !important; box-shadow: 0 8px 32px var(--mistral-shadow) !important; padding: clamp(1rem, 2vw, 1.5rem) !important; } .mistral-panel .markdown h3, .mistral-section h3 { color: var(--mistral-text) !important; font-size: 0.82rem !important; font-weight: 800 !important; letter-spacing: 0.08em !important; margin: 0.35rem 0 0.85rem !important; text-transform: uppercase !important; } .mistral-section { border-top: 1px solid var(--mistral-border) !important; margin-top: 1rem !important; padding-top: 1rem !important; } .mistral-panel label, .mistral-panel .wrap label, .mistral-panel span { color: var(--mistral-text) !important; font-family: 'Inter', sans-serif !important; } .mistral-panel input, .mistral-panel textarea, .mistral-panel select { background: #FFF8E3 !important; border-color: var(--mistral-border) !important; color: var(--mistral-text) !important; font-family: 'Inter', sans-serif !important; } .mistral-panel input:focus, .mistral-panel textarea:focus { border-color: var(--mistral-accent) !important; box-shadow: 0 0 0 2px rgba(255, 130, 5, 0.18) !important; } #process-button { background: var(--mistral-accent) !important; border: 0 !important; border-radius: 0 !important; color: #FFFFFF !important; font-weight: 800 !important; letter-spacing: 0.06em !important; min-height: 3rem !important; text-transform: uppercase !important; } #process-button:hover { background: var(--mistral-accent-hover) !important; } .mistral-panel .gr-group, .mistral-panel .styler { background: transparent !important; border-color: var(--mistral-border) !important; } .mistral-panel button:not(#process-button):not(.reset-button):not(.center) { background: #FFF8E3 !important; border: 1px solid var(--mistral-border) !important; color: var(--mistral-text) !important; font-weight: 700 !important; letter-spacing: 0 !important; text-transform: none !important; } .mistral-panel button:not(#process-button):not(.reset-button):not(.center):hover { background: var(--mistral-panel-warm) !important; border-color: var(--mistral-accent) !important; } .mistral-panel button.center.boundedheight { border: 2px dashed var(--mistral-border) !important; color: var(--mistral-muted) !important; min-height: 12rem !important; } .mistral-panel button.center.boundedheight svg { color: var(--mistral-accent) !important; } .mistral-panel button.reset-button { background: transparent !important; border: 0 !important; color: var(--mistral-soft-muted) !important; min-height: auto !important; } #pipeline-report textarea { background-color: var(--mistral-panel) !important; background-image: linear-gradient(rgba(0, 0, 0, var(--mistral-grid-opacity)) 1px, transparent 1px), linear-gradient(90deg, rgba(0, 0, 0, var(--mistral-grid-opacity)) 1px, transparent 1px) !important; background-size: 20px 20px !important; color: var(--mistral-text) !important; font-family: 'JetBrains Mono', monospace !important; font-size: 0.95rem !important; line-height: 1.7 !important; } #document-preview { background: var(--mistral-panel) !important; border: 2px solid var(--mistral-border) !important; box-shadow: 0 8px 32px var(--mistral-shadow) !important; padding: clamp(1rem, 2vw, 1.5rem) !important; } #document-preview h3 { color: var(--mistral-text) !important; font-size: 0.82rem !important; font-weight: 800 !important; letter-spacing: 0.08em !important; text-transform: uppercase !important; } #document-preview .grid-wrap, #document-preview .thumbnail-item { background: #FFF8E3 !important; } .mistral-note { color: var(--mistral-soft-muted) !important; font-size: 0.9rem !important; margin-top: 0.5rem !important; } """ MISTRAL_THEME = gr.themes.Soft(primary_hue="orange", secondary_hue="yellow") GRADIO_BLOCKS_KWARGS = {"title": "PDF2Dataset -- Mistral OCR Pipeline"} GRADIO_LAUNCH_KWARGS = {} try: _GRADIO_MAJOR_VERSION = int(gr.__version__.split(".", 1)[0]) except (AttributeError, ValueError): _GRADIO_MAJOR_VERSION = 5 if _GRADIO_MAJOR_VERSION >= 6: GRADIO_LAUNCH_KWARGS.update(theme=MISTRAL_THEME, css=MISTRAL_CSS) else: GRADIO_BLOCKS_KWARGS.update(theme=MISTRAL_THEME, css=MISTRAL_CSS) def _gradio_process( file_objs: list[str] | str | None, chunk_size: int, repo_name: str, append_mode: bool, min_chars: int, min_words: int, oauth_token: gr.OAuthToken | None, ) -> str: """Bridge between Gradio inputs, the signed-in user, and core processing.""" if oauth_token is None: return "Error: Sign in with Hugging Face before processing files." if not file_objs: return "Error: No files uploaded." if not isinstance(file_objs, list): file_objs = [file_objs] file_paths = [f.name if hasattr(f, "name") else str(f) for f in file_objs] return process_files( file_paths, chunk_size, oauth_token.token, repo_name, append_mode=append_mode, min_chunk_chars=min_chars, min_words=min_words, ) with gr.Blocks(**GRADIO_BLOCKS_KWARGS) as demo: gr.Markdown( """ # PDF2Dataset Convert PDFs and images into clean Hugging Face datasets with **Mistral OCR**, structure-aware chunking, validation, deduplication, and quality filtering. """, elem_id="brand-hero", ) with gr.Column(elem_id="app-shell"): with gr.Row(): with gr.Column(scale=1, elem_classes=["mistral-panel"]): file_input = gr.File( label="Source documents", file_types=[".pdf", ".png", ".jpg", ".jpeg", ".webp", ".bmp"], type="filepath", file_count="multiple", ) with gr.Group(elem_classes=["mistral-section"]): gr.Markdown("### Chunking") chunk_size = gr.Slider( minimum=0, maximum=4096, value=512, step=64, label="Max chunk size", info="Character budget per chunk. Use 0 to keep each document as one chunk.", ) with gr.Group(elem_classes=["mistral-section"]): gr.Markdown("### Quality filters") with gr.Row(): min_chars = gr.Slider( minimum=0, maximum=500, value=20, step=5, label="Minimum characters", ) min_words = gr.Slider( minimum=0, maximum=50, value=3, step=1, label="Minimum words", ) with gr.Group(elem_classes=["mistral-section"]): gr.Markdown("### Hugging Face output") gr.LoginButton() repo_name = gr.Textbox( label="Dataset repository", placeholder="your-username/your-dataset-name", ) append_mode = gr.Checkbox( label="Append to existing dataset", value=False, ) submit_btn = gr.Button( "Process and push", variant="primary", elem_id="process-button", ) with gr.Column(scale=1, elem_classes=["mistral-panel"]): output = gr.Textbox( label="Pipeline report", lines=30, interactive=False, elem_id="pipeline-report", ) with gr.Group(elem_id="document-preview"): gr.Markdown("### Document preview") preview_gallery = gr.Gallery( label="Uploaded documents", columns=2, height="auto", object_fit="contain", ) gr.Markdown( "*Requires `MISTRAL_API_KEY`. Sign in with Hugging Face before pushing a dataset.*", elem_classes=["mistral-note"], ) file_input.change( fn=render_preview, inputs=[file_input], outputs=[preview_gallery], ) submit_btn.click( fn=_gradio_process, inputs=[ file_input, chunk_size, repo_name, append_mode, min_chars, min_words, ], outputs=output, ) gr.Examples( examples=[ [None, 512, "hf-username/my-first-ocr-dataset", False, 20, 3], [None, 1024, "hf-username/large-chunk-ocr-data", True, 50, 5], [None, 0, "hf-username/no-split-ocr-data", False, 0, 0], ], inputs=[ file_input, chunk_size, repo_name, append_mode, min_chars, min_words, ], outputs=output, fn=_gradio_process, cache_examples=False, ) def main(): """Entry point for the application.""" demo.launch( share=os.getenv("GRADIO_SHARE", "False").lower() == "true", debug=True, **GRADIO_LAUNCH_KWARGS, ) if __name__ == "__main__": main()