Update app.py

app.py

@@ -1,7 +1,7 @@
 import gradio as gr
 from langchain_text_splitters import MarkdownHeaderTextSplitter, RecursiveCharacterTextSplitter
 from langchain.schema import Document
-from typing import List
+from typing import List, Dict
 import logging
 import re
 from pathlib import Path
@@ -35,22 +35,22 @@ def encode_image(image_path):
         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:
+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:
+def get_combined_markdown(ocr_response) -> tuple[str, str]:
     markdowns = []
     raw_markdowns = []
+    image_data = {}  # Collect all image data
     for page in ocr_response.pages:
-        image_data = {}
         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)
+    return "\n\n".join(markdowns), "\n\n".join(raw_markdowns), image_data
 
 # Perform OCR on uploaded file
 def perform_ocr_file(file):
@@ -85,22 +85,23 @@ def perform_ocr_file(file):
                 include_image_base64=True
             )
         else:
-            return "Unsupported file type. Please provide a PDF or an image (png, jpeg, jpg).", ""
+            return "Unsupported file type. Please provide a PDF or an image (png, jpeg, jpg).", "", {}
 
-        combined_markdown, raw_markdown = get_combined_markdown(ocr_response)
-        return combined_markdown, raw_markdown
+        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)}", ""
+        return f"Error during OCR: {str(e)}", "", {}
 
-# Function to extract base64 images from markdown content
-def extract_images_from_markdown(markdown_text: str) -> List[str]:
-    # Regex to match markdown image syntax with base64 data
-    pattern = r"!\[.*?\]\((data:image/[a-z]+;base64,[^\)]+)\)"
-    return re.findall(pattern, markdown_text)
+# 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
@@ -111,7 +112,6 @@ def chunk_markdown(
             ("#", "Header 1"),
             ("##", "Header 2"),
             ("###", "Header 3"),
-            ("####", "Header 4")
         ]
 
         # Initialize MarkdownHeaderTextSplitter
@@ -146,8 +146,8 @@ def chunk_markdown(
 
         # Add images to metadata
         for chunk in chunks:
-            images = extract_images_from_markdown(chunk.page_content)
-            chunk.metadata["images"] = images
+            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
@@ -167,29 +167,24 @@ def text_to_base64_dummy(text: str, chunk_index: int):
 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 = perform_ocr_file(file)
+        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, chunk_size, chunk_overlap, strip_headers)
+        chunks = chunk_markdown(combined_markdown, image_data, chunk_size, chunk_overlap, strip_headers)
         
         # Step 3: Prepare dataset
         data = {
             "chunk_id": [],
             "content": [],
             "metadata": [],
-            "images": []  # Changed to store list of images
         }
         
         for i, chunk in enumerate(chunks):
             data["chunk_id"].append(i)
             data["content"].append(chunk.page_content)
-            data["metadata"].append({k: v for k, v in chunk.metadata.items() if k != "images"})  # Exclude images from metadata column
-            images = chunk.metadata.get("images", [])
-            if not images:  # If no images, add a placeholder
-                images = [text_to_base64_dummy(chunk.page_content, i)]
-            data["images"].append(images)
+            data["metadata"].append(chunk.metadata)
 
         # Step 4: Create and push dataset to Hugging Face
         dataset = Dataset.from_dict(data)