Javedalam commited on
Commit
305c374
Β·
verified Β·
1 Parent(s): 7a48026

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +117 -0
app.py ADDED
@@ -0,0 +1,117 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import os, tempfile, traceback
2
+ import gradio as gr
3
+ import spaces
4
+ import requests
5
+
6
+ # Persist caches in /data so future runs are fast
7
+ os.environ.setdefault("HF_HOME", "/data/.cache/huggingface")
8
+ os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")
9
+
10
+ from docling.datamodel.base_models import InputFormat
11
+ from docling.document_converter import DocumentConverter, PdfFormatOption
12
+ from docling.pipeline.vlm_pipeline import VlmPipeline
13
+
14
+ # (Optional) reduce thread pressure on ZeroGPU
15
+ try:
16
+ import torch
17
+ torch.set_num_threads(max(1, int(os.environ.get("OMP_NUM_THREADS", "2"))))
18
+ HAS_CUDA = torch.cuda.is_available()
19
+ except Exception:
20
+ HAS_CUDA = False
21
+
22
+ # Build converters once
23
+ std_converter = DocumentConverter(
24
+ format_options={InputFormat.PDF: PdfFormatOption()}
25
+ )
26
+ vlm_converter = DocumentConverter(
27
+ format_options={InputFormat.PDF: PdfFormatOption(pipeline_cls=VlmPipeline)}
28
+ )
29
+
30
+ def _success(md: str, html: str):
31
+ tmpdir = tempfile.gettempdir()
32
+ md_path = os.path.join(tmpdir, "output.md")
33
+ html_path = os.path.join(tmpdir, "output.html")
34
+ with open(md_path, "w", encoding="utf-8") as f: f.write(md)
35
+ with open(html_path, "w", encoding="utf-8") as f: f.write(html)
36
+ return md, md_path, html_path
37
+
38
+ def _fail(msg: str):
39
+ # Show the full error text in the Markdown preview
40
+ return f"**Conversion failed**:\n```\n{msg}\n```", None, None
41
+
42
+ def _convert_source(path: str, use_vlm: bool):
43
+ try:
44
+ converter = vlm_converter if use_vlm else std_converter
45
+ doc = converter.convert(source=path).document
46
+ md = doc.export_to_markdown()
47
+ html = doc.export_to_html()
48
+ return _success(md, html)
49
+ except Exception as e:
50
+ return _fail(f"{e}\n\n{traceback.format_exc()}")
51
+
52
+ @spaces.GPU(duration=600)
53
+ def run_convert_file(file, mode):
54
+ if file is None:
55
+ return _fail("No file provided.")
56
+ use_vlm = mode.startswith("VLM")
57
+ return _convert_source(file.name, use_vlm)
58
+
59
+ @spaces.GPU(duration=600)
60
+ def run_convert_url(url, mode):
61
+ if not url:
62
+ return _fail("No URL provided.")
63
+ # Download to a temp file so Docling always sees a local path
64
+ try:
65
+ r = requests.get(url, stream=True, timeout=60)
66
+ r.raise_for_status()
67
+ fd, tmp_path = tempfile.mkstemp(suffix=".pdf")
68
+ with os.fdopen(fd, "wb") as tmp:
69
+ for chunk in r.iter_content(chunk_size=1 << 20):
70
+ if chunk:
71
+ tmp.write(chunk)
72
+ except Exception as e:
73
+ return _fail(f"Failed to download URL: {e}")
74
+ use_vlm = mode.startswith("VLM")
75
+ out = _convert_source(tmp_path, use_vlm)
76
+ try:
77
+ os.remove(tmp_path)
78
+ except Exception:
79
+ pass
80
+ return out
81
+
82
+ subtitle = "Device: **CUDA (ZeroGPU)**" if HAS_CUDA else "Device: **CPU** (GPU warms on first call)"
83
+
84
+ with gr.Blocks(title="Granite-Docling 258M β€” PDF β†’ Markdown/HTML") as demo:
85
+ gr.Markdown(
86
+ f"""# Granite-Docling 258M β€” PDF β†’ Markdown / HTML
87
+ {subtitle}
88
+
89
+ **Modes**
90
+ - **Standard (faster)** β†’ PDFs with a text layer
91
+ - **VLM (Granite – better for complex/scanned)** β†’ scans / heavy tables / formulas
92
+
93
+ _First call may be slow while models download and ZeroGPU warms. Cache lives in `/data`._
94
+ """
95
+ )
96
+
97
+ mode = gr.Radio(
98
+ ["Standard (faster)", "VLM (Granite – better for complex/scanned)"],
99
+ value="Standard (faster)", label="Mode"
100
+ )
101
+
102
+ with gr.Tab("Upload PDF"):
103
+ fi = gr.File(file_types=[".pdf"], label="PDF")
104
+ out_md = gr.Markdown(label="Markdown Preview")
105
+ dl_md = gr.File(label="Download Markdown")
106
+ dl_html = gr.File(label="Download HTML")
107
+ gr.Button("Convert").click(run_convert_file, [fi, mode], [out_md, dl_md, dl_html])
108
+
109
+ with gr.Tab("Convert from URL"):
110
+ url = gr.Textbox(label="Public PDF URL", placeholder="https://.../file.pdf")
111
+ out2_md = gr.Markdown(label="Markdown Preview")
112
+ dl2_md = gr.File(label="Download Markdown")
113
+ dl2_html = gr.File(label="Download HTML")
114
+ gr.Button("Convert").click(run_convert_url, [url, mode], [out2_md, dl2_md, dl2_html])
115
+
116
+ # Explicit bind/queue so Spaces picks up the interface
117
+ demo.queue().launch(server_name="0.0.0.0", server_port=int(os.environ.get("PORT", 7860)))