Update app.py

app.py

@@ -8,7 +8,6 @@ from pathlib import Path
 from transformers import AutoProcessor, Qwen2VLForConditionalGeneration
 
 from olmocr.data.renderpdf import render_pdf_to_base64png
-from olmocr.prompts import build_finetuning_prompt
 from olmocr.prompts.anchor import get_anchor_text
 
 from ebooklib import epub
@@ -36,14 +35,23 @@ def process_pdf_to_epub(pdf_file, title, author):
 
     for i in range(num_pages):
         page_num = i + 1
-        print(f"Processing page {page_num}...")  # Debugging line
+        print(f"Processing page {page_num}...")
 
         try:
             # Render page to base64 image
             image_base64 = render_pdf_to_base64png(pdf_path, page_num, target_longest_image_dim=1024)
             anchor_text = get_anchor_text(pdf_path, page_num, pdf_engine="pdfreport", target_length=4000)
-            print(f"Anchor text for page {page_num}: {anchor_text}")  # Debugging line
-            prompt = build_finetuning_prompt(anchor_text)
+            print(f"Anchor text for page {page_num}: {anchor_text}")
+
+            # New prompt format
+            prompt = (
+                "Below is the image of one page of a document, as well as some raw textual content that was previously "
+                "extracted for it. Just return the plain text representation of this document as if you were reading it naturally.\n"
+                "Do not hallucinate.\n"
+                "RAW_TEXT_START\n"
+                f"{anchor_text}\n"
+                "RAW_TEXT_END"
+            )
 
             messages = [
                 {
@@ -72,11 +80,12 @@ def process_pdf_to_epub(pdf_file, title, author):
                 num_return_sequences=1,
                 do_sample=True,
             )
+
             prompt_length = inputs["input_ids"].shape[1]
             new_tokens = output[:, prompt_length:].detach().cpu()
-            
+
             decoded = "[No output generated]"
-            if new_tokens.shape[1] > 0:
+            if new_tokens is not None and new_tokens.shape[1] > 0:
                 try:
                     decoded_list = processor.tokenizer.batch_decode(new_tokens, skip_special_tokens=True)
                     decoded = decoded_list[0].strip() if decoded_list else "[No output generated]"
@@ -87,18 +96,8 @@ def process_pdf_to_epub(pdf_file, title, author):
 
         except Exception as processing_error:
             decoded = f"[Processing error on page {page_num}: {str(processing_error)}]"
-        else:
-            try:
-                # Check if the tokens are empty
-                if not new_tokens:
-                    decoded = f"[No tokens generated for page {page_num}]"
-                else:
-                    decoded_list = processor.tokenizer.batch_decode(new_tokens, skip_special_tokens=True)
-                    decoded = decoded_list[0].strip() if decoded_list else "[No output generated]"
-            except Exception as decode_error:
-                decoded = f"[Decoding error on page {page_num}: {str(decode_error)}]"
 
-        print(f"Decoded content for page {page_num}: {decoded}")  # Debugging line
+        print(f"Decoded content for page {page_num}: {decoded}")
 
         # Create chapter
         chapter = epub.EpubHtml(title=f"Page {page_num}", file_name=f"page_{page_num}.xhtml", lang="en")
@@ -134,14 +133,14 @@ iface = gr.Interface(
     outputs=gr.File(label="Download EPUB"),
     title="PDF to EPUB Converter (with olmOCR)",
     description="Uploads a PDF, extracts text from each page with vision + prompt, and builds an EPUB using the outputs. Sets the first page as cover.",
-    allow_flagging="never"  # Add this line to avoid the flagged directory issue
+    allow_flagging="never"
 )
 
 if __name__ == "__main__":
     iface.launch(
-        server_name="0.0.0.0",  # Required to make app publicly accessible
-        server_port=7860,       # Can be changed if needed
-        share=True,             # Optional: creates a public Gradio link if supported
-        debug=True,             # Optional: helpful if you're troubleshooting
-        allowed_paths=["/tmp"]  # Optional: makes it explicit that Gradio can write here
+        server_name="0.0.0.0",
+        server_port=7860,
+        share=True,
+        debug=True,
+        allowed_paths=["/tmp"]
     )