import gradio as gr from langchain_text_splitters import MarkdownHeaderTextSplitter, RecursiveCharacterTextSplitter from langchain.schema import Document from typing import List, Dict import logging import re from pathlib import Path import requests import base64 import io from PIL import Image from datasets import Dataset from huggingface_hub import HfApi import os from mistralai import Mistral # Configure logging logging.basicConfig(level=logging.INFO) logger = logging.getLogger(__name__) # Mistral OCR setup api_key = os.environ.get("MISTRAL_API_KEY") if not api_key: raise ValueError("MISTRAL_API_KEY environment variable not set") client = Mistral(api_key=api_key) # Function to encode image to base64 def encode_image(image_path): try: with open(image_path, "rb") as image_file: return base64.b64encode(image_file.read()).decode('utf-8') except FileNotFoundError: return "Error: The file was not found." except Exception as e: return f"Error: {e}" # Function to replace images in markdown with base64 strings def replace_images_in_markdown(markdown_str: str, images_dict: Dict[str, str]) -> str: for img_name, base64_str in images_dict.items(): markdown_str = markdown_str.replace(f"![{img_name}]({img_name})", f"![{img_name}]({base64_str})") return markdown_str # Function to combine markdown from OCR response def get_combined_markdown(ocr_response) -> tuple[str, str]: markdowns = [] raw_markdowns = [] image_data = {} # Collect all image data for page in ocr_response.pages: for img in page.images: image_data[img.id] = img.image_base64 markdowns.append(replace_images_in_markdown(page.markdown, image_data)) raw_markdowns.append(page.markdown) return "\n\n".join(markdowns), "\n\n".join(raw_markdowns), image_data # Perform OCR on uploaded file def perform_ocr_file(file): try: if file.name.lower().endswith('.pdf'): uploaded_pdf = client.files.upload( file={ "file_name": file.name, "content": open(file.name, "rb"), }, purpose="ocr" ) signed_url = client.files.get_signed_url(file_id=uploaded_pdf.id) ocr_response = client.ocr.process( model="mistral-ocr-latest", document={ "type": "document_url", "document_url": signed_url.url, }, include_image_base64=True ) client.files.delete(file_id=uploaded_pdf.id) elif file.name.lower().endswith(('.png', '.jpg', '.jpeg')): base64_image = encode_image(file.name) ocr_response = client.ocr.process( model="mistral-ocr-latest", document={ "type": "image_url", "image_url": f"data:image/jpeg;base64,{base64_image}" }, include_image_base64=True ) else: return "Unsupported file type. Please provide a PDF or an image (png, jpeg, jpg).", "", {} combined_markdown, raw_markdown, image_data = get_combined_markdown(ocr_response) return combined_markdown, raw_markdown, image_data except Exception as e: return f"Error during OCR: {str(e)}", "", {} # Function to extract image names from markdown content def extract_image_names_from_markdown(markdown_text: str) -> List[str]: # Regex to match markdown image syntax pattern = r"!\[(.*?)\]\(" return [match.replace("![","").replace("](","") for match in re.findall(pattern, markdown_text)] # Function to chunk markdown text with image handling def chunk_markdown( markdown_text: str, image_data: Dict[str, str], chunk_size: int = 1000, chunk_overlap: int = 200, strip_headers: bool = True ) -> List[Document]: try: # Define headers to split on headers_to_split_on = [ ("#", "Header 1"), ("##", "Header 2"), ("###", "Header 3"), ] # Initialize MarkdownHeaderTextSplitter markdown_splitter = MarkdownHeaderTextSplitter( headers_to_split_on=headers_to_split_on, strip_headers=strip_headers ) # Split markdown by headers logger.info("Splitting markdown by headers") chunks = markdown_splitter.split_text(markdown_text) # If chunk_size is specified, further split large chunks if chunk_size > 0: text_splitter = RecursiveCharacterTextSplitter( chunk_size=chunk_size, chunk_overlap=chunk_overlap, length_function=len, separators=["\n\n", "\n", ".", " ", ""], keep_separator=True, add_start_index=True ) logger.info(f"Applying character-level splitting with chunk_size={chunk_size}") final_chunks = [] for chunk in chunks: if len(chunk.page_content) > chunk_size: sub_chunks = text_splitter.split_documents([chunk]) final_chunks.extend(sub_chunks) else: final_chunks.append(chunk) chunks = final_chunks # Add images to metadata for chunk in chunks: image_names = extract_image_names_from_markdown(chunk.page_content) chunk.metadata["images"] = {name: image_data.get(name, None) for name in image_names} logger.info(f"Created {len(chunks)} chunks") return chunks except Exception as e: logger.error(f"Error processing markdown: {str(e)}") raise # Placeholder image generation def text_to_base64_dummy(text: str, chunk_index: int): img = Image.new('RGB', (200, 200), color='white') buffer = io.BytesIO() img.save(buffer, format="PNG") return base64.b64encode(buffer.getvalue()).decode("utf-8") # Process file: OCR -> Chunk -> Save def process_file_and_save(file, chunk_size, chunk_overlap, strip_headers, hf_token, repo_name): try: # Step 1: Perform OCR combined_markdown, raw_markdown, image_data = perform_ocr_file(file) if "Error" in combined_markdown: return combined_markdown # Step 2: Chunk the markdown chunks = chunk_markdown(combined_markdown, image_data, chunk_size, chunk_overlap, strip_headers) # Step 3: Prepare dataset data = { "chunk_id": [], "content": [], "metadata": [], } for i, chunk in enumerate(chunks): data["chunk_id"].append(i) data["content"].append(chunk.page_content) data["metadata"].append(chunk.metadata) # Step 4: Create and push dataset to Hugging Face dataset = Dataset.from_dict(data) api = HfApi() api.create_repo(repo_id=repo_name, token=hf_token, repo_type="dataset", exist_ok=True) dataset.push_to_hub(repo_name, token=hf_token) return f"Dataset created with {len(chunks)} chunks and saved to Hugging Face at {repo_name}" except Exception as e: return f"Error: {str(e)}" # Gradio Interface with gr.Blocks(title="PDF/Image OCR, Markdown Chunking, and Dataset Creator") as demo: gr.Markdown("# PDF/Image OCR, Markdown Chunking, and Dataset Creator") gr.Markdown("Upload a PDF or image, extract text/images with Mistral OCR, chunk the markdown by headers, and save to Hugging Face.") with gr.Row(): with gr.Column(): file_input = gr.File(label="Upload PDF or Image") chunk_size = gr.Slider(0, 2000, value=1000, step=100, label="Max Chunk Size (0 to disable)") chunk_overlap = gr.Slider(0, 500, value=200, step=50, label="Chunk Overlap") strip_headers = gr.Checkbox(label="Strip Headers from Content", value=True) hf_token = gr.Textbox(label="Hugging Face Token", type="password") repo_name = gr.Textbox(label="Hugging Face Repository Name (e.g., username/dataset-name)") submit_btn = gr.Button("Process and Save") with gr.Column(): output = gr.Textbox(label="Result") submit_btn.click( fn=process_file_and_save, inputs=[file_input, chunk_size, chunk_overlap, strip_headers, hf_token, repo_name], outputs=output ) demo.launch(share=True)