Update app.py

app.py

@@ -1,111 +1,108 @@
+from dotenv import load_dotenv
+
+load_dotenv()
+
 import gradio as gr
-from langchain_text_splitters import MarkdownHeaderTextSplitter, RecursiveCharacterTextSplitter
-from langchain.schema import Document
-from typing import List, Dict, Any, Tuple
+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
-from datasets import Dataset
+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 gradio_client.utils as client_utils
+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')
+logging.basicConfig(
+    level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s"
+)
 logger = logging.getLogger(__name__)
 
-# --- Patch Gradio's get_type function to handle boolean schemas ---
-def patched_get_type(schema: Any) -> str:
-    """Patched version of get_type to handle boolean schemas."""
-    if isinstance(schema, bool):
-        return "bool"
-    if "const" in schema:
-        return f"Literal[{repr(schema['const'])}]"
-    if "enum" in schema:
-        return f"Literal[{', '.join(repr(v) for v in schema['enum'])}]"
-    if "type" not in schema:
-        return "Any"
-    type_ = schema["type"]
-    if isinstance(type_, list):
-        return f"Union[{', '.join(t for t in type_ if t != 'null')}]"
-    if type_ == "array":
-        items = schema.get("items", {})
-        return f"List[{patched_json_schema_to_python_type(items, schema.get('$defs'))}]"
-    if type_ == "object":
-        return "Dict[str, Any]"
-    if type_ == "null":
-        return "None"
-    if type_ == "integer":
-        return "int"
-    if type_ == "number":
-        return "float"
-    if type_ == "boolean":
-        return "bool"
-    return type_
-
-def patched_json_schema_to_python_type(schema: Any, defs: Dict[str, Any] = None) -> str:
-    """Patched version of json_schema_to_python_type to use patched_get_type."""
-    defs = defs or {}
-    if not schema:
-        return "Any"
-    if "$ref" in schema:
-        ref = schema["$ref"].split("/")[-1]
-        return patched_json_schema_to_python_type(defs.get(ref, {}), defs)
-    if "anyOf" in schema:
-        types = [
-            patched_json_schema_to_python_type(s, defs) for s in schema["anyOf"]
-        ]
-        return f"Union[{', '.join(t for t in types if t != 'None')}]"
-    if "type" in schema and schema["type"] == "array":
-        items = schema.get("items", {})
-        elements = patched_json_schema_to_python_type(items, defs)
-        return f"List[{elements}]"
-    if "type" in schema and schema["type"] == "object":
-        if "properties" in schema:
-            des = [
-                f"{n}: {patched_json_schema_to_python_type(v, defs)}{client_utils.get_desc(v)}"
-                for n, v in schema["properties"].items()
-            ]
-            return f"Dict[str, Union[{', '.join(des)}]]"
-        if "additionalProperties" in schema:
-            return f"Dict[str, {patched_json_schema_to_python_type(schema['additionalProperties'], defs)}]"
-        return "Dict[str, Any]"
-    return patched_get_type(schema)
-
-# Override Gradio's json_schema_to_python_type
-client_utils.json_schema_to_python_type = patched_json_schema_to_python_type
-
-# --- Mistral OCR Setup ---
-api_key = os.environ.get("MISTRAL_API_KEY")
-hf_token_global = None
-client = None
-
-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.")
-    else:
-        logger.warning("No API key found.")
 
-if api_key:
-    try:
-        client = Mistral(api_key=api_key)
-        logger.info("Mistral client initialized successfully.")
-    except Exception as e:
-        logger.error(f"Failed to initialize Mistral client: {e}", exc_info=True)
-        raise RuntimeError(f"Failed to initialize Mistral client: {e}")
-else:
-    logger.error("Mistral API key not available. OCR will fail.")
+# --- 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')
+    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": <raw 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]:
     """
@@ -113,12 +110,15 @@ def extract_images_from_markdown(markdown_text: str) -> Dict[str, str]:
     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)
+    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}"
+        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).
@@ -130,145 +130,203 @@ def replace_image_references(markdown_text: str, image_map: Dict[str, str]) -> s
         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]]:
+
+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' attribute or pages list is empty.")
+    if not hasattr(ocr_response, "pages") or not ocr_response.pages:
+        logger.warning("OCR response has no pages.")
         return "", "", {}
 
-    try:
-        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
-                        logger.debug(f"Page {page_idx}: Image ID {img.id} added to image_data_map.")
-                    else:
-                        logger.warning(f"Page {page_idx}: Image object lacks 'id' or valid 'image_base64'. Image: {img}")
-            else:
-                logger.info(f"Page {page_idx}: No images found.")
+    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
+        if not hasattr(page, "markdown"):
+            logger.warning(f"Page {page_idx} lacks 'markdown' attribute. Skipping.")
+            continue
 
-            current_raw_markdown = page.markdown if page.markdown else ""
-            raw_markdowns.append(current_raw_markdown)
-            current_processed_markdown = current_raw_markdown
-
-            img_refs = re.findall(r"!\[.*?\]\((.*?)\)", current_processed_markdown)
-            logger.debug(f"Page {page_idx}: Found {len(img_refs)} image references in 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"(\))"
-                    if re.search(pattern, current_processed_markdown):
-                        current_processed_markdown = re.sub(
-                            pattern,
-                            r"\1" + base64_data_uri + r"\2",
-                            current_processed_markdown
-                        )
-                        logger.debug(f"Page {page_idx}: Replaced image ID {img_id} with base64 data URI.")
-                elif not img_id.startswith(('http:', 'https:', 'data:')):
-                    logger.warning(f"Page {page_idx}: Image ID '{img_id}' not in image data.")
+        current_raw_markdown = page.markdown or ""
+        raw_markdowns.append(current_raw_markdown)
+        current_processed_markdown = current_raw_markdown
 
-            processed_markdowns.append(current_processed_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."
+                )
 
-        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
+        processed_markdowns.append(current_processed_markdown)
 
-    except Exception as e:
-        logger.error(f"Error processing OCR response markdown: {e}", exc_info=True)
-        raise
+    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(file_obj: Any) -> Tuple[str, str, Dict[str, str]]:
-    """Performs OCR on an uploaded file using Mistral API."""
-    if not client:
-        return "Error: Mistral client not initialized.", "", {}
-    if not file_obj:
-        return "Error: No file provided.", "", {}
 
-    try:
-        file_path = file_obj.name
-        file_name = getattr(file_obj, 'orig_name', os.path.basename(file_path))
-        logger.info(f"Performing OCR on file: {file_name}")
-        file_ext = os.path.splitext(file_name)[1].lower()
+def perform_ocr(file_path: str) -> tuple[str, str, Dict[str, str]]:
+    """Performs OCR on a file using Mistral API.
 
-        ocr_response = None
-        uploaded_file_id = None
+    Args:
+        file_path: Path to the file on disk.
 
-        if file_ext == '.pdf':
-            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 successfully. 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
-                )
-                logger.info(f"OCR response received: {ocr_response}")
-            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 ['.png', '.jpg', '.jpeg', '.webp', '.bmp']:
+    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:
-                image_bytes = f.read()
-                if not image_bytes:
-                    return f"Error: Uploaded image file '{file_name}' is empty.", "", {}
-                base64_encoded_image = 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_image}"
-                ocr_response = client.ocr.process(
-                    model="mistral-ocr-latest",
-                    document={"type": "image_url", "image_url": data_uri},
-                    include_image_base64=True
-                )
-                logger.info(f"OCR response received: {ocr_response}")
+                file_content = f.read()
 
-        else:
-            return f"Unsupported file type: '{file_name}'.", "", {}
+            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}")
 
-        if ocr_response:
-            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
-        return f"Error: OCR failed for '{file_name}'.", "", {}
+            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()
 
-    except Exception as e:
-        logger.error(f"Error during OCR: {e}", exc_info=True)
-        return f"Error during OCR: {str(e)}", "", {}
 
 def chunk_markdown(
     markdown_text_with_images: str,
-    chunk_size: int = 1000,
-    chunk_overlap: int = 200,
-    strip_headers: bool = True
-) -> List[Document]:
-    """Chunks markdown text, preserving headers in metadata and extracting images."""
+    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 []
@@ -278,168 +336,730 @@ def chunk_markdown(
     updated_markdown = replace_image_references(markdown_text_with_images, image_map)
     logger.info(f"Extracted {len(image_map)} images from markdown.")
 
-    headers_to_split_on = [
-        ("#", "Header 1"), ("##", "Header 2"), ("###", "Header 3"),
-        ("####", "Header 4"), ("#####", "Header 5"), ("######", "Header 6"),
-    ]
-    markdown_splitter = MarkdownHeaderTextSplitter(
-        headers_to_split_on=headers_to_split_on, strip_headers=strip_headers
+    # 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,
     )
-    header_chunks = markdown_splitter.split_text(updated_markdown)
-
-    if not header_chunks:
-        logger.warning("No header chunks created. Treating entire text as one chunk.")
-        return [Document(page_content=updated_markdown, metadata={"images_base64": list(image_map.values())})]
-
-    final_chunks = []
-    if chunk_size > 0:
-        text_splitter = RecursiveCharacterTextSplitter(
-            chunk_size=chunk_size, chunk_overlap=chunk_overlap, length_function=len,
-            separators=["\n\n", "\n", "(?<=\. )", "(?<=\? )", "(?<=! )", ", ", "; ", " ", ""],
-            add_start_index=True
+    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),
+            }
         )
-        for i, header_chunk in enumerate(header_chunks):
-            if header_chunk.page_content:
-                sub_chunks = text_splitter.split_documents([header_chunk])
-                final_chunks.extend(sub_chunks)
-                logger.debug(f"Header chunk {i}: Split into {len(sub_chunks)} sub-chunks.")
-            else:
-                logger.debug(f"Header chunk {i}: Empty, skipping.")
-    else:
-        final_chunks = [chunk for chunk in header_chunks if chunk.page_content]
-
-    # Add image references to metadata for each chunk
-    for chunk in final_chunks:
-        if not hasattr(chunk, 'metadata'):
-            chunk.metadata = {}
-        # Find image references in this chunk
-        chunk_img_refs = re.findall(r"!\[.*?\]\((img_ref_\d+)\)", chunk.page_content)
-        chunk_images = [image_map[ref_id] for ref_id in chunk_img_refs if ref_id in image_map]
-        chunk.metadata["images_base64"] = chunk_images
-        chunk.metadata["image_references"] = chunk_img_refs
-        logger.debug(f"Chunk {chunk.metadata.get('start_index', 'unknown')}: Found {len(chunk_images)} images.")
-
-    logger.info(f"Created {len(final_chunks)} final chunks.")
-    return final_chunks
-
-def get_hf_token(explicit_token: str = None) -> str:
+
+    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<n<10K" if total >= 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."""
-    global hf_token_global
-    
-    if explicit_token and explicit_token.strip() and explicit_token.startswith('hf_'):
+    if explicit_token and explicit_token.strip() and explicit_token.startswith("hf_"):
         return explicit_token.strip()
-    
-    if hf_token_global:
-        return hf_token_global
-        
+
     env_token = os.environ.get("HF_TOKEN")
-    if env_token and env_token.startswith('hf_'):
-        hf_token_global = env_token
+    if env_token and env_token.startswith("hf_"):
         return env_token
-    
+
     try:
         stored_token = huggingface_hub.get_token()
         if stored_token:
-            hf_token_global = 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_file_and_save(
-    file_objs: Any, chunk_size: int, chunk_overlap: int,
-    strip_headers: bool, hf_token: str, repo_name: str
+
+def process_files(
+    file_paths: list[str],
+    chunk_size: int,
+    hf_token: str,
+    repo_name: str,
+    append_mode: bool = False,
+    min_chunk_chars: int = 20,
+    min_words: int = 3,
 ) -> str:
-    """Orchestrates OCR, chunking, and saving to Hugging Face for multiple files."""
-    # Handle case where file_objs is a single file or None
-    if not file_objs:
+    """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 isinstance(file_objs, list):
-        file_objs = [file_objs]
-    
-    if not repo_name or '/' not in repo_name:
+
+    if not repo_name or "/" not in repo_name:
         return "Error: Invalid repository name (use 'username/dataset-name')."
 
-    if chunk_size < 0:
-        chunk_size = 0
-    if chunk_overlap < 0:
-        chunk_overlap = 0
-    if chunk_size > 0 and chunk_overlap >= chunk_size:
-        chunk_overlap = min(200, chunk_size // 2)
+    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.
-        Please either:
-        1. Provide a token in the input field (starts with 'hf_')
-        2. Set HF_TOKEN environment variable
-        3. Run `huggingface-cli login` in your terminal"""
+        return (
+            "Error: No valid Hugging Face token found.\n"
+            "Please either:\n"
+            "1. Provide a token in the input field (starts with 'hf_')\n"
+            "2. Set HF_TOKEN environment variable\n"
+            "3. Run `huggingface-cli login` in your terminal"
+        )
 
     try:
-        all_data = {
-            "chunk_id": [],
-            "text": [],
-            "metadata": [],
-            "source_filename": []
-        }
-        total_chunks = 0
+        # 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_obj in enumerate(file_objs, 1):
-            source_filename = getattr(file_obj, 'orig_name', os.path.basename(file_obj.name))
-            logger.info(f"--- Processing file {file_idx}/{len(file_objs)}: {source_filename} ---")
+        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} ---"
+            )
 
-            processed_markdown, raw_markdown, img_map = perform_ocr_file(file_obj)
-            if processed_markdown.startswith("Error:"):
-                error_messages.append(f"File '{source_filename}': {processed_markdown}")
-                logger.error(f"Failed to process file {source_filename}: {processed_markdown}")
+            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, chunk_overlap, strip_headers)
+            chunks = chunk_markdown(processed_markdown, chunk_size)
             if not chunks:
-                error_messages.append(f"File '{source_filename}': Failed to chunk the document.")
+                error_messages.append(
+                    f"File '{source_filename}': Failed to chunk the document."
+                )
                 logger.error(f"Failed to chunk file {source_filename}")
                 continue
 
-            all_data["chunk_id"].extend([f"{source_filename}_chunk_{i}" for i in range(len(chunks))])
-            all_data["text"].extend([chunk.page_content or "" for chunk in chunks])
-            all_data["metadata"].extend([chunk.metadata for chunk in chunks])
-            all_data["source_filename"].extend([source_filename] * len(chunks))
-            total_chunks += len(chunks)
+            # 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}: Added {len(chunks)} chunks. Total chunks: {total_chunks}")
+            logger.info(
+                f"File {source_filename}: queued {len(chunks)} chunks for pipeline"
+            )
 
-        if not all_data["chunk_id"]:
-            return "Error: No valid data processed from any files.\n" + "\n".join(error_messages)
+        if files_processed == 0:
+            return "Error: No files were processed successfully.\n" + "\n".join(
+                error_messages
+            )
 
-        dataset = Dataset.from_dict(all_data)
-        api = HfApi(token=effective_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}"
+        # Run the pipeline
+        all_data, stats = builder.build()
 
-        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 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 ""
+                )
+            )
 
-        dataset.push_to_hub(repo_name, token=effective_hf_token, 
-                          commit_message=f"Add OCR data from {files_processed} files")
-        repo_url = f"https://huggingface.co/datasets/{repo_name}"
-        result = f"Success! Dataset with {total_chunks} chunks from {files_processed}/{len(file_objs)} files saved to: {repo_url}"
+        # 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:
-            result += "\n\nErrors encountered:\n" + "\n".join(error_messages)
-        return result
+            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')
+        status = getattr(hf_http_err.response, "status_code", "Unknown")
         if status == 401:
             return "Error: Invalid or unauthorized Hugging Face token."
         elif status == 403:
@@ -447,81 +1067,457 @@ def process_file_and_save(
         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: {str(e)}\n" + "\n".join(error_messages)
+        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 ---
-with gr.Blocks(title="Mistral OCR & Dataset Creator", 
-              theme=gr.themes.Soft(primary_hue="blue", secondary_hue="cyan")) as demo:
-    gr.Markdown("# Mistral OCR, Markdown Chunking, and Hugging Face Dataset Creator")
+
+
+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, chunk_size, hf_token, repo_name, append_mode, min_chars, min_words
+):
+    """Bridge between Gradio file objects and core processing logic."""
+    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,
+        hf_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(
         """
-        Upload one or more PDF or image files. The application will:
-        1. Extract text and images using Mistral OCR for each file
-        2. Embed images as base64 data URIs in markdown
-        3. Chunk markdown by headers and optionally character count
-        4. Store embedded images in chunk metadata
-        5. Create/update a Hugging Face Dataset with all processed data
-        """
+        # 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.Row():
-        with gr.Column(scale=1):
-            file_input = gr.File(
-                label="Upload PDF or Image Files",
-                file_types=['.pdf', '.png', '.jpg', '.jpeg', '.webp', '.bmp'],
-                type="filepath",
-                file_count="multiple"  # Allow multiple file uploads
+    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")
+                    repo_name = gr.Textbox(
+                        label="Dataset repository",
+                        placeholder="your-username/your-dataset-name",
+                    )
+                    hf_token = gr.Textbox(
+                        label="Hugging Face token",
+                        type="password",
+                        placeholder="hf_... or set HF_TOKEN",
+                    )
+                    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("## Chunking Options")
-            chunk_size = gr.Slider(minimum=0, maximum=8000, value=1000, step=100,
-                                 label="Max Chunk Size (Characters)")
-            chunk_overlap = gr.Slider(minimum=0, maximum=1000, value=200, step=50,
-                                    label="Chunk Overlap (Characters)")
-            strip_headers = gr.Checkbox(label="Strip Headers from Content", value=True)
-            gr.Markdown("## Hugging Face Output Options")
-            repo_name = gr.Textbox(label="HF Dataset Repository",
-                                 placeholder="your-username/your-dataset-name")
-            hf_token = gr.Textbox(label="Hugging Face Token", type="password",
-                                placeholder="hf_...")
-            submit_btn = gr.Button("Process and Save", variant="primary")
-
-        with gr.Column(scale=1):
-            output = gr.Textbox(label="Result Status", lines=20, interactive=False)
+
+        gr.Markdown(
+            "*Requires `MISTRAL_API_KEY`, a Hugging Face token, or an active Hugging Face CLI login.*",
+            elem_classes=["mistral-note"],
+        )
+
+    file_input.change(
+        fn=render_preview,
+        inputs=[file_input],
+        outputs=[preview_gallery],
+    )
 
     submit_btn.click(
-        fn=process_file_and_save,
-        inputs=[file_input, chunk_size, chunk_overlap, strip_headers, hf_token, repo_name],
-        outputs=output
+        fn=_gradio_process,
+        inputs=[
+            file_input,
+            chunk_size,
+            hf_token,
+            repo_name,
+            append_mode,
+            min_chars,
+            min_words,
+        ],
+        outputs=output,
     )
 
     gr.Examples(
         examples=[
-            [None, 1000, 200, True, "", "hf-username/my-first-ocr-dataset"],
-            [None, 2000, 400, True, "", "hf-username/large-chunk-ocr-data"],
-            [None, 0, 0, False, "", "hf-username/header-only-ocr-data"],
+            [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,
+            hf_token,
+            repo_name,
+            append_mode,
+            min_chars,
+            min_words,
         ],
-        inputs=[file_input, chunk_size, chunk_overlap, strip_headers, hf_token, repo_name],
         outputs=output,
-        fn=process_file_and_save,
-        cache_examples=False
+        fn=_gradio_process,
+        cache_examples=False,
     )
-    
-    gr.Markdown("*Requires MISTRAL_API_KEY or HF token*")
 
-if __name__ == "__main__":
-    import gradio
-    logger.info(f"Using Gradio version: {gradio.__version__}")
-    if not gradio.__version__.startswith("4."):
-        logger.warning("Gradio version is not 4.x. Updating to the latest version is recommended.")
-        print("Consider running: pip install --upgrade gradio")
-
-    initial_token = get_hf_token()
-    if not initial_token and not client:
-        print("\nWARNING: Neither Mistral API key nor HF token found.")
-        print("Set MISTRAL_API_KEY and/or HF_TOKEN, or use `huggingface-cli login`")
-    
+
+def main():
+    """Entry point for the application."""
     demo.launch(
-        share=os.getenv('GRADIO_SHARE', 'False').lower() == 'true',
+        share=os.getenv("GRADIO_SHARE", "False").lower() == "true",
         debug=True,
-        auth_message="Provide a valid Hugging Face token if prompted"
+        **GRADIO_LAUNCH_KWARGS,
     )
+
+
+if __name__ == "__main__":
+    main()