Create app.py

app.py

@@ -0,0 +1,371 @@
+import os
+import re
+import sys
+import json
+import time
+import copy
+import base64
+import asyncio
+import tempfile
+import subprocess
+from pathlib import Path
+from datetime import datetime
+import zipfile
+
+import numpy as np
+import gradio as gr
+from PIL import Image                 
+from pdf2image import convert_from_path 
+from loguru import logger            
+from openai import OpenAI, AsyncOpenAI   
+from gradio_pdf import PDF                
+
+import uuid
+import tqdm
+
+import requests
+
+
+preset_prompts = [
+    "Please convert the document into Markdown format.",
+    "Generate a clean and structured Markdown version of the document.",
+    "Transform this content into Markdown with proper headings and bullet points.",
+    "Convert the text to Markdown, preserving structure and formatting.",
+    "Reformat this document as Markdown with clear sections and lists.",
+]
+
+
+def send_pdf_to_parse(file_path, server_ip="47.117.17.202", port=9999, route="/upload", api_key=None):
+    url = f"http://{server_ip}:{port}{route}"
+    headers = {}
+    if api_key:
+        headers["Authorization"] = f"Bearer {api_key}"
+
+    with open(file_path, "rb") as f:
+        files = {"file": (os.path.basename(file_path), f, "application/pdf")}
+        response = requests.post(url, files=files, headers=headers)
+    return response
+
+
+def extract_makrdown(text):
+    m = re.search(r'```markdown\s*([\s\S]*?)```', text, re.MULTILINE)
+    if m:
+        return m.group(1).strip()
+    else:
+        return text
+
+openai_api_key = "EMPTY"
+openai_api_base = "http://47.117.17.202:9999/v1"
+
+
+client = AsyncOpenAI(
+    api_key=openai_api_key,
+    base_url=openai_api_base,
+)
+
+
+async def request(messages):
+    
+    chat_completion_from_base64 = await client.chat.completions.create(
+        messages=messages,
+        model="Qwen2_5VL",
+        max_completion_tokens=4096,
+        stream=True,                                    
+        temperature=0.0,
+        top_p=0.95
+    )
+    
+    page = ""
+    async for chunk in chat_completion_from_base64:
+        if chunk.choices[0].delta.content:
+            content = chunk.choices[0].delta.content
+            
+            choice = chunk.choices[0]
+            if choice.finish_reason is not None:
+                print(f"end reason = {choice.finish_reason}")
+                break
+            page += content
+            
+            yield content
+            
+
+def images_to_pdf(img_paths, pdf_path):
+
+    if isinstance(img_paths, (str, Path)):
+        img_paths = [img_paths]
+
+    if not img_paths:
+        raise ValueError("img_paths is empty")
+    images = []
+    for p in img_paths:
+        p = Path(p)
+        if not p.is_file():
+            raise FileNotFoundError(p)
+
+        img = Image.open(p)
+        if img.mode in ("RGBA", "P"):     
+            img = img.convert("RGB")
+        images.append(img)
+
+    pdf_path = Path(pdf_path)
+    pdf_path.parent.mkdir(parents=True, exist_ok=True)
+    images[0].save(pdf_path,
+                   save_all=True,   
+                   append_images=images[1:],
+                   resolution=300.0)   
+    return pdf_path
+
+
+def encode_image(image_path):
+    with open(image_path, "rb") as image_file:
+        return base64.b64encode(image_file.read()).decode("utf-8")
+
+def build_message(image_path, prompt):
+    
+    content = [
+        {
+            "type": "image_url",
+            "image_url": {
+                "url": f"data:image/jpeg;base64,{encode_image(image_path)}"
+            }
+        },
+        {"type": "text", 'text': prompt}  
+    ]
+    
+    
+    messages = [
+        {"role": "system", "content": "You are a helpful assistant."},
+        {'role': 'user', 'content': content}
+        
+    ]
+    
+    return messages
+
+
+
+def download_markdown_file(md_text):
+    filename = f"markdown_{uuid.uuid4().hex[:8]}.md"
+    filepath = Path("downloads") / filename
+    filepath.parent.mkdir(exist_ok=True)
+    with open(filepath, "w", encoding="utf-8") as f:
+        f.write(md_text)
+    return str(filepath)
+
+
+async def doc_parser(doc_path, prompt):
+
+    doc_path = Path(doc_path)
+    if not doc_path.is_file():
+        raise FileNotFoundError(doc_path)
+
+    with tempfile.TemporaryDirectory() as tmpdir:
+        tmpdir = Path(tmpdir)
+
+        queries = []
+        if doc_path.suffix.lower() == ".pdf":
+            pages: List[Image.Image] = convert_from_path(doc_path, dpi=300)
+            for idx, page in enumerate(pages, start=1):
+                img_path = tmpdir / f"page_{idx}.png"
+                page.save(img_path, "PNG")
+                
+                messages = build_message(img_path, prompt)
+                queries.append(messages)
+                
+        else:
+            messages = build_message(doc_path, prompt)
+            queries.append(messages)
+            
+    all_pages = []
+    all_pages_raw = []
+    for query in queries:
+        pages = ""
+        async for chunk in request(query):
+            pages += chunk
+            yield extract_makrdown(pages), pages
+        all_pages.append(extract_makrdown(pages))
+        all_pages_raw.append(pages)
+        print(all_pages)
+        yield "\n---\n".join(all_pages), "\n\n".join(all_pages_raw)
+
+        
+def compress_directory_to_zip(directory_path, output_zip_path):
+    try:
+        with zipfile.ZipFile(output_zip_path, 'w', zipfile.ZIP_DEFLATED) as zipf:
+
+   
+            for root, dirs, files in os.walk(directory_path):
+                for file in files:
+            
+                    file_path = os.path.join(root, file)
+
+                    arcname = os.path.relpath(file_path, directory_path)
+            
+                    zipf.write(file_path, arcname)
+        return 0
+    except Exception as e:
+        logger.exception(e)
+        return -1
+
+latex_delimiters = [
+    {'left': '$$', 'right': '$$', 'display': True},
+    {'left': '$', 'right': '$', 'display': False},
+    {'left': '\\(', 'right': '\\)', 'display': False},
+    {'left': '\\[', 'right': '\\]', 'display': True},
+]
+
+def check_prompt(prompt):
+    if not prompt or prompt.strip() == "":
+        raise gr.Error("Please select or enter a prompt before parsing.")
+    return prompt
+
+def to_file(image_path):
+    
+    if image_path.endswith("Academic_Papers.png"):
+        image_path = image_path.replace("Academic_Papers.png", "Academic_Papers.pdf")
+
+    return image_path
+
+def process_file(file_path):
+    if file_path is None:
+        return None
+    if not file_path.endswith(".pdf"):
+        
+        tmp_file_path = Path(file_path)          
+        tmp_file_path = tmp_file_path.with_suffix(".pdf")
+        images_to_pdf(file_path, tmp_file_path)
+    else:
+        send_pdf_to_parse(file_path)
+        tmp_file_path = file_path
+
+    return str(tmp_file_path)
+
+
+# def render_pdf_viewer(pdf_file_path):
+#     print(pdf_file_path)
+#     # 假设 pdf_file_path 
+#     return f"""
+#     <iframe 
+#         src="https://mozilla.github.io/pdf.js/web/viewer.html?file=file://{pdf_file_path}" 
+#         width="100%" height="1000px" 
+#         style="border: none;">
+#     </iframe>
+#     """
+def render_pdf_viewer(pdf_file_path):
+    filename = os.path.basename(pdf_file_path)
+    print(pdf_file_path)
+    return f"""
+    <iframe 
+        src="https://mozilla.github.io/pdf.js/web/viewer.html?file=/file={pdf_file_path}" 
+        width="100%" height="1000px" 
+        style="border: none;">
+    </iframe>
+    """
+
+
+if __name__ == '__main__':
+    with gr.Blocks() as demo:
+        with gr.Row():
+            with gr.Column(variant='panel', scale=5):
+                
+                file = gr.File(label='Please upload a PDF or image', file_types=['.pdf', '.png', '.jpeg', '.jpg'], type="filepath")
+                prompts = gr.Dropdown(
+                    choices=preset_prompts,
+                    label="Prompt",
+                    info="Enter or select prompts...",
+                    value=preset_prompts[0],              
+                    multiselect=False,
+                    interactive=True,
+                    allow_custom_value=True, 
+                )
+
+                with gr.Row():
+                    change_bu = gr.Button('Parse')
+                    clear_bu = gr.ClearButton(value='Clear')
+                #pdf_show = PDF(label='Preview', interactive=False, visible=True, height=800)
+                pdf_show = gr.HTML()
+
+         
+                example_root = os.path.join(os.path.dirname(__file__), 'examples')
+                images = [
+                    os.path.join(example_root, f)
+                    for f in os.listdir(example_root)
+                    if f.lower().endswith(('png', 'jpg', 'jpeg'))
+                ]
+
+            with gr.Column(variant='panel', scale=5):
+                with gr.Accordion("Examples", open=True):
+                    example_root = "examples"
+                    file_path = [
+                        os.path.join(example_root, f)
+                        for f in ["Financial_Reports.png", "Books.png", "Magazines.png", "Academic_Papers.png"]
+                       
+                    ]
+                    
+                    with gr.Row():
+                        for i, label in enumerate(["Financial Reports(IMG)", "Books(IMG)", "Magazines(IMG)", "Academic Papers(PDF)"]):
+                            with gr.Column(scale=1, min_width=120):
+                                gr.Image(
+                                    value=file_path[i],
+                                    width=120,
+                                    height=90,
+                                    show_label=False,
+                                    show_download_button=False
+                                )
+                                gr.Button(label).click(fn=to_file, inputs=gr.State(file_path[i]), outputs=file)
+                            
+                            
+                download_btn = gr.Button("⬇️ Generate download link", size="sm")
+                output_file = gr.File(label='Parse result', interactive=False, elem_id="down-file-box",visible=False)
+               
+                gr.HTML("""
+                <style>
+                #down-file-box {
+                    max-height: 300px;
+                }
+                </style>
+                """)
+                with gr.Tabs():
+                    with gr.Tab('Markdown rendering'):
+                        md = gr.Markdown(label='Markdown rendering', height=1100, show_copy_button=True,
+                             latex_delimiters=latex_delimiters,
+                             line_breaks=True)
+                    with gr.Tab('Markdown text'):
+                        md_text = gr.TextArea(lines=45, show_copy_button=True)
+                            
+                        
+               
+        # file.change(fn=process_file, inputs=file, outputs=pdf_show)
+        file.change(
+            lambda f: render_pdf_viewer(f), 
+            inputs=file, 
+            outputs=pdf_show
+        )
+
+        change_bu.click(
+            fn=check_prompt,
+            inputs=prompts,
+            outputs=prompts
+        ).then(
+            lambda f: gr.update(visible=False),
+            inputs=output_file,
+            outputs=output_file
+        ).then(
+            fn=doc_parser,
+            inputs=[file, prompts],
+            outputs=[md, md_text]
+        )
+        
+        clear_bu.add([file, md, pdf_show, md_text])
+        
+        download_btn.click(
+            fn=download_markdown_file,
+            inputs=md_text,
+            outputs=output_file
+        ).then(
+            lambda f: gr.update(visible=True),
+            inputs=output_file,
+            outputs=output_file
+        )
+
+
+    demo.launch(server_name='0.0.0.0',share=True)