Files changed (4) hide show
  1. .gitignore +4 -0
  2. README.md +94 -7
  3. app.py +444 -54
  4. requirements.txt +11 -0
.gitignore ADDED
@@ -0,0 +1,4 @@
 
 
 
 
 
1
+ __pycache__/
2
+ *.pyc
3
+ .gradio/
4
+ /tmp/
README.md CHANGED
@@ -1,16 +1,103 @@
1
  ---
2
  title: Chandra OCR 2
3
- emoji: πŸ’¬
4
- colorFrom: yellow
5
  colorTo: purple
6
  sdk: gradio
7
- sdk_version: 6.5.1
8
  app_file: app.py
9
  pinned: false
10
- hf_oauth: true
11
- hf_oauth_scopes:
12
- - inference-api
13
  license: apache-2.0
 
 
 
 
 
 
 
 
14
  ---
15
 
16
- An example chatbot using [Gradio](https://gradio.app), [`huggingface_hub`](https://huggingface.co/docs/huggingface_hub/v0.22.2/en/index), and the [Hugging Face Inference API](https://huggingface.co/docs/api-inference/index).
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
  ---
2
  title: Chandra OCR 2
3
+ emoji: πŸ“„
4
+ colorFrom: indigo
5
  colorTo: purple
6
  sdk: gradio
 
7
  app_file: app.py
8
  pinned: false
 
 
 
9
  license: apache-2.0
10
+ short_description: Layout-aware document OCR to markdown, HTML, or JSON
11
+ models:
12
+ - datalab-to/chandra-ocr-2
13
+ tags:
14
+ - ocr
15
+ - document-ai
16
+ - vision-language
17
+ suggested_hardware: zero-a10g
18
  ---
19
 
20
+ # Chandra OCR 2 β€” Space demo
21
+
22
+ Gradio demo for [`datalab-to/chandra-ocr-2`](https://huggingface.co/datalab-to/chandra-ocr-2),
23
+ Datalab's layout-aware document OCR model. Upload PDFs or images and get back
24
+ markdown with tables, math, forms, and reading order preserved.
25
+
26
+ ## Hardware
27
+
28
+ **This will not run on the free CPU tier.** The model is ~10B parameters,
29
+ roughly 20 GB in bf16.
30
+
31
+ | Hardware | Works? |
32
+ |---|---|
33
+ | CPU basic (free) | No β€” OOM at load |
34
+ | ZeroGPU (H200 slice) | Yes β€” recommended |
35
+ | L40S / A100 | Yes |
36
+ | T4 (16 GB) | No, unless you add 4-bit quantisation |
37
+
38
+ Set this under **Settings β†’ Hardware** after creating the Space. `suggested_hardware`
39
+ in the frontmatter is only a hint to visitors; it does not provision anything.
40
+
41
+ ## How it works
42
+
43
+ - PDFs are rasterised page by page with PyMuPDF (no poppler/apt needed).
44
+ - Each page is capped on its long edge before inference β€” visual token count
45
+ scales with area, so this is the main lever on latency and memory.
46
+ - Pages are batched `PAGES_PER_GPU_CALL` at a time so each ZeroGPU allocation
47
+ finishes inside its duration budget.
48
+ - Output is offered three ways: rendered markdown, markdown source, and the raw
49
+ model string (useful when a prompt type returns HTML or JSON instead).
50
+
51
+ The app prefers the official `chandra` package (`generate_hf` + `BatchInputItem`
52
+ + `parse_markdown`). If that import fails, it falls back to driving the chat
53
+ template through plain `transformers` so the Space still boots.
54
+
55
+ ## Configuration
56
+
57
+ Edit the constants at the top of `app.py`:
58
+
59
+ | Constant | Default | Purpose |
60
+ |---|---|---|
61
+ | `PROMPT_TYPES` | `ocr_layout`, … | Prompt types offered in the dropdown |
62
+ | `MAX_PAGES` | 20 | Per-run page cap |
63
+ | `PAGES_PER_GPU_CALL` | 3 | Pages per ZeroGPU allocation |
64
+ | `GPU_DURATION` | 180 | Seconds requested per allocation |
65
+
66
+ Only `ocr_layout` is confirmed from the model card Quickstart. The others are
67
+ exposed on the assumption that the card's markdown/HTML/JSON output modes map to
68
+ prompt types; if one errors, remove it or check the `chandra` package docs.
69
+
70
+ ## Licence β€” read before making this public
71
+
72
+ The Space **code** here is Apache-2.0. The **model weights** are not:
73
+
74
+ > Code is Apache 2.0. Model weights are under a modified OpenRAIL-M license.
75
+ > Free for research, personal use, and startups under $2M funding/revenue.
76
+ > Cannot be used competitively with our API.
77
+
78
+ A public, free, hosted OCR endpoint is plausibly "competitive with our API."
79
+ If you are past the revenue threshold, or intend this as a product rather than a
80
+ demo, check with Datalab first. Setting the Space to **private** avoids the
81
+ question entirely.
82
+
83
+ ## Troubleshooting: "upload isn't working"
84
+
85
+ Nine times out of ten the upload widget is fine and the **container crashed at
86
+ startup**, so every interaction fails. Check **Logs β†’ Container** (not Build)
87
+ and read the last ~30 lines.
88
+
89
+ | Log line | Cause | Fix |
90
+ |---|---|---|
91
+ | `MODEL FAILED TO LOAD` + `CUDA out of memory` / `Killed` | Running on CPU basic or too small a GPU | Settings β†’ Hardware β†’ **ZeroGPU** |
92
+ | `No module named 'chandra'` | `chandra-ocr` didn't install | Check Build logs for the pip failure; app falls back to plain transformers and still works |
93
+ | `Unrecognized configuration class` / `KeyError` | `transformers` too old for this architecture | Raise the pin in `requirements.txt` |
94
+ | `Application startup failed` / port timeout | Model download exceeded the startup window | Factory reboot (Settings); weights are cached after the first success |
95
+ | Nothing obviously wrong, UI loads, upload button does nothing | Browser/Gradio `file_types` mismatch | Already handled β€” this app sets no `file_types` filter |
96
+
97
+ The app now opens a **Diagnostics** panel automatically when the model failed to
98
+ load, showing torch/transformers/gradio versions, whether CUDA is visible, and
99
+ the exact load error. Start there.
100
+
101
+ Unsupported or corrupt files no longer fail the whole run β€” they are listed in
102
+ the status box as `Skipped: filename (reason)` and the readable pages still
103
+ process.
app.py CHANGED
@@ -1,69 +1,459 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
  import gradio as gr
2
- from huggingface_hub import InferenceClient
3
-
4
-
5
- def respond(
6
- message,
7
- history: list[dict[str, str]],
8
- system_message,
9
- max_tokens,
10
- temperature,
11
- top_p,
12
- hf_token: gr.OAuthToken,
13
- ):
14
- """
15
- For more information on `huggingface_hub` Inference API support, please check the docs: https://huggingface.co/docs/huggingface_hub/v0.22.2/en/guides/inference
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
16
  """
17
- client = InferenceClient(token=hf_token.token, model="openai/gpt-oss-20b")
 
 
 
18
 
19
- messages = [{"role": "system", "content": system_message}]
 
 
 
 
20
 
21
- messages.extend(history)
 
 
 
 
 
 
 
22
 
23
- messages.append({"role": "user", "content": message})
24
 
25
- response = ""
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
26
 
27
- for message in client.chat_completion(
28
- messages,
29
- max_tokens=max_tokens,
30
- stream=True,
31
- temperature=temperature,
32
- top_p=top_p,
33
- ):
34
- choices = message.choices
35
- token = ""
36
- if len(choices) and choices[0].delta.content:
37
- token = choices[0].delta.content
38
 
39
- response += token
40
- yield response
 
 
 
 
 
 
 
41
 
42
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
43
  """
44
- For information on how to customize the ChatInterface, peruse the gradio docs: https://www.gradio.app/docs/chatinterface
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
45
  """
46
- chatbot = gr.ChatInterface(
47
- respond,
48
- additional_inputs=[
49
- gr.Textbox(value="You are a friendly Chatbot.", label="System message"),
50
- gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max new tokens"),
51
- gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature"),
52
- gr.Slider(
53
- minimum=0.1,
54
- maximum=1.0,
55
- value=0.95,
56
- step=0.05,
57
- label="Top-p (nucleus sampling)",
58
- ),
59
- ],
60
- )
61
-
62
- with gr.Blocks() as demo:
63
- with gr.Sidebar():
64
- gr.LoginButton()
65
- chatbot.render()
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
66
 
67
 
68
  if __name__ == "__main__":
69
- demo.launch()
 
1
+ """
2
+ Chandra OCR 2 β€” Hugging Face Space demo.
3
+
4
+ Model: datalab-to/chandra-ocr-2 (~10B, bf16)
5
+ Docs: https://huggingface.co/datalab-to/chandra-ocr-2
6
+
7
+ Hardware: needs ZeroGPU (H200 slice) or a paid A100/L40S.
8
+ The model is ~20 GB in bf16 and will NOT run on the free CPU tier.
9
+ """
10
+
11
+ import inspect
12
+ import json
13
+ import os
14
+ import time
15
+ import zipfile
16
+ from pathlib import Path
17
+
18
  import gradio as gr
19
+ import torch
20
+ from PIL import Image
21
+
22
+ # ---------------------------------------------------------------------------
23
+ # Config
24
+ # ---------------------------------------------------------------------------
25
+
26
+ MODEL_ID = "datalab-to/chandra-ocr-2"
27
+
28
+ # Chandra is prompt-type driven rather than free-text prompted. 'ocr_layout' is
29
+ # the one shown in the model card Quickstart; the others are exposed because the
30
+ # card advertises markdown / HTML / JSON output. If one errors, the package
31
+ # doesn't support that name in your installed version β€” check `chandra` docs.
32
+ PROMPT_TYPES = ["ocr_layout", "ocr", "ocr_html", "ocr_json", "layout"]
33
+ DEFAULT_PROMPT_TYPE = "ocr_layout"
34
+
35
+ PAGES_PER_GPU_CALL = 3 # keep each ZeroGPU allocation inside its duration budget
36
+ GPU_DURATION = 180 # seconds requested per allocation
37
+ MAX_PAGES = 20 # guard against someone uploading a 500-page PDF
38
+
39
+ OUT_DIR = Path(os.environ.get("CHANDRA_OUT_DIR", "/tmp/chandra_out"))
40
+ OUT_DIR.mkdir(parents=True, exist_ok=True)
41
+
42
+ IMAGE_EXTS = {".png", ".jpg", ".jpeg", ".webp", ".bmp", ".tif", ".tiff"}
43
+
44
+ # ---------------------------------------------------------------------------
45
+ # ZeroGPU shim β€” lets the same file run locally without the `spaces` package
46
+ # ---------------------------------------------------------------------------
47
+
48
+ # Set by the platform on ZeroGPU hardware.
49
+ ON_ZERO = os.environ.get("SPACES_ZERO_GPU") == "true"
50
+
51
+ try:
52
+ import spaces
53
+
54
+ gpu = spaces.GPU
55
+ except ImportError: # local / non-ZeroGPU deploy
56
+
57
+ def gpu(*args, **kwargs):
58
+ if args and callable(args[0]):
59
+ return args[0]
60
+
61
+ def deco(fn):
62
+ return fn
63
+
64
+ return deco
65
+
66
+
67
+ # ---------------------------------------------------------------------------
68
+ # Chandra package β€” preferred path. Falls back to plain transformers.
69
+ # ---------------------------------------------------------------------------
70
+
71
+ CHANDRA_ERR = None
72
+ try:
73
+ from chandra.model import generate_hf
74
+ from chandra.model.schema import BatchInputItem
75
+
76
+ try:
77
+ from chandra.output import parse_markdown
78
+ except ImportError:
79
+ from chandra.model.output import parse_markdown
80
+
81
+ USE_CHANDRA = True
82
+ except Exception as e: # noqa: BLE001
83
+ USE_CHANDRA = False
84
+ CHANDRA_ERR = f"{type(e).__name__}: {e}"
85
+
86
+
87
+ # ---------------------------------------------------------------------------
88
+ # Model
89
+ # ---------------------------------------------------------------------------
90
+
91
+ def _load():
92
+ from transformers import AutoProcessor
93
+
94
+ try:
95
+ from transformers import AutoModelForImageTextToText as VLM
96
+ except ImportError:
97
+ from transformers import AutoModelForVision2Seq as VLM
98
+
99
+ # On ZeroGPU there is no GPU visible at import time, so accelerate's
100
+ # device_map="auto" would strand the model on CPU. ZeroGPU instead
101
+ # intercepts .to("cuda") at global scope. Elsewhere, device_map is fine.
102
+ kw = dict(low_cpu_mem_usage=True)
103
+ if not ON_ZERO:
104
+ kw["device_map"] = "auto"
105
+
106
+ try:
107
+ m = VLM.from_pretrained(MODEL_ID, dtype=torch.bfloat16, **kw)
108
+ except TypeError: # transformers < 4.56 spelled it torch_dtype
109
+ m = VLM.from_pretrained(MODEL_ID, torch_dtype=torch.bfloat16, **kw)
110
+
111
+ if ON_ZERO:
112
+ m = m.to("cuda")
113
+
114
+ m.eval()
115
+ p = AutoProcessor.from_pretrained(MODEL_ID)
116
+ p.tokenizer.padding_side = "left" # required by chandra's batched generate
117
+ m.processor = p
118
+ return m, p
119
+
120
+
121
+ model = processor = None
122
+ MODEL_ERR = None
123
+
124
+ print(f"Loading {MODEL_ID} ...")
125
+ _t0 = time.time()
126
+ try:
127
+ model, processor = _load()
128
+ print(f"Loaded in {time.time() - _t0:.0f}s")
129
+ except Exception as e: # noqa: BLE001
130
+ # Do NOT re-raise. A crash here kills the container and every interaction
131
+ # afterwards fails with an opaque error, including the file picker.
132
+ import traceback
133
+ MODEL_ERR = f"{type(e).__name__}: {e}"
134
+ traceback.print_exc()
135
+ print(f"!! MODEL FAILED TO LOAD: {MODEL_ERR}")
136
+
137
+ print(f"chandra pkg: {USE_CHANDRA} ({CHANDRA_ERR or 'ok'})")
138
+
139
+
140
+ def diagnostics() -> str:
141
+ import platform
142
+ try:
143
+ import transformers as _tf
144
+ tfv = _tf.__version__
145
+ except Exception: # noqa: BLE001
146
+ tfv = "?"
147
+ gpu_name = torch.cuda.get_device_name(0) if torch.cuda.is_available() else "none visible"
148
+ return (
149
+ f"python {platform.python_version()} | torch {torch.__version__} | "
150
+ f"transformers {tfv} | gradio {gr.__version__}\n"
151
+ f"ZeroGPU: {ON_ZERO} | cuda now: {torch.cuda.is_available()} ({gpu_name})\n"
152
+ f"chandra package: {'ok' if USE_CHANDRA else CHANDRA_ERR}\n"
153
+ f"model: {'loaded' if model is not None else 'FAILED β€” ' + str(MODEL_ERR)}"
154
+ )
155
+
156
+
157
+ # ---------------------------------------------------------------------------
158
+ # Page extraction
159
+ # ---------------------------------------------------------------------------
160
+
161
+ def as_path(f) -> Path:
162
+ """Normalise str / Path / Gradio file object into a real Path.
163
+
164
+ Note pathlib.Path also has .name, but there it's the basename β€” checking
165
+ hasattr first would silently drop the directory.
166
  """
167
+ if isinstance(f, (str, os.PathLike)):
168
+ return Path(f)
169
+ return Path(getattr(f, "name", str(f)))
170
+
171
 
172
+ def pdf_to_images(path: Path, dpi: int):
173
+ try:
174
+ import pymupdf
175
+ except ImportError:
176
+ import fitz as pymupdf
177
 
178
+ doc = pymupdf.open(str(path))
179
+ pages = []
180
+ for i, page in enumerate(doc):
181
+ pix = page.get_pixmap(dpi=dpi)
182
+ pages.append((f"{path.stem}_p{i + 1:03d}",
183
+ Image.frombytes("RGB", (pix.width, pix.height), pix.samples)))
184
+ doc.close()
185
+ return pages
186
 
 
187
 
188
+ def collect_pages(files, dpi: int):
189
+ """Returns (pages, skipped). Never raises on a single bad file."""
190
+ pages, skipped = [], []
191
+ for f in files:
192
+ p = as_path(f)
193
+ ext = p.suffix.lower()
194
+ try:
195
+ if ext == ".pdf":
196
+ pages.extend(pdf_to_images(p, dpi))
197
+ elif ext in IMAGE_EXTS:
198
+ pages.append((p.stem, Image.open(p).convert("RGB")))
199
+ else:
200
+ skipped.append(f"{p.name} (unsupported type '{ext or 'none'}')")
201
+ except Exception as e: # noqa: BLE001
202
+ skipped.append(f"{p.name} ({type(e).__name__}: {e})")
203
+ return pages, skipped
204
 
 
 
 
 
 
 
 
 
 
 
 
205
 
206
+ def fit(img: Image.Image, max_side: int) -> Image.Image:
207
+ """Cap the long edge β€” visual token count scales with area, so this is the
208
+ single biggest lever on VRAM and latency."""
209
+ img = img.convert("RGB")
210
+ if max(img.size) > max_side:
211
+ s = max_side / max(img.size)
212
+ img = img.resize((max(1, int(img.width * s)), max(1, int(img.height * s))),
213
+ Image.LANCZOS)
214
+ return img
215
 
216
 
217
+ # ---------------------------------------------------------------------------
218
+ # Inference
219
+ # ---------------------------------------------------------------------------
220
+
221
+ @gpu(duration=GPU_DURATION)
222
+ @torch.inference_mode()
223
+ def _infer_chunk(images, prompt_type: str, max_new_tokens: int):
224
+ """OCR a small batch of PIL images. Returns a list of raw model strings."""
225
+ if model is None:
226
+ raise RuntimeError(
227
+ f"Model never loaded: {MODEL_ERR}. Check the Container logs and "
228
+ f"confirm Settings > Hardware is ZeroGPU or a GPU tier."
229
+ )
230
+
231
+ if USE_CHANDRA:
232
+ batch = [BatchInputItem(image=im, prompt_type=prompt_type) for im in images]
233
+ try:
234
+ results = generate_hf(batch, model, max_tokens=max_new_tokens)
235
+ except TypeError:
236
+ results = generate_hf(batch, model)
237
+ return [getattr(r, "raw", None) or getattr(r, "markdown", "") or str(r)
238
+ for r in results]
239
+
240
+ # ---- fallback: drive the chat template directly ----
241
+ outs = []
242
+ for im in images:
243
+ msgs = [{"role": "user", "content": [
244
+ {"type": "image", "image": im},
245
+ {"type": "text", "text": prompt_type},
246
+ ]}]
247
+ inputs = processor.apply_chat_template(
248
+ msgs, tokenize=True, add_generation_prompt=True,
249
+ return_dict=True, return_tensors="pt",
250
+ ).to(model.device)
251
+ if "pixel_values" in inputs:
252
+ inputs["pixel_values"] = inputs["pixel_values"].to(model.dtype)
253
+ gen = model.generate(**inputs, max_new_tokens=int(max_new_tokens), do_sample=False)
254
+ trimmed = [o[len(i):] for i, o in zip(inputs["input_ids"], gen)]
255
+ outs.append(processor.batch_decode(trimmed, skip_special_tokens=True)[0].strip())
256
+ return outs
257
+
258
+
259
+ def to_markdown(raw: str) -> str:
260
+ if USE_CHANDRA:
261
+ try:
262
+ return parse_markdown(raw)
263
+ except Exception: # noqa: BLE001
264
+ pass
265
+ return raw
266
+
267
+
268
+ # ---------------------------------------------------------------------------
269
+ # Orchestration
270
+ # ---------------------------------------------------------------------------
271
+
272
+ def run(files, prompt_type, dpi, max_side, max_new_tokens,
273
+ progress=gr.Progress()):
274
+
275
+ def msg(m):
276
+ return "", "", "", m, None, None
277
+
278
+ if not files:
279
+ yield msg("Upload a PDF or some images first.")
280
+ return
281
+
282
+ pages, skipped = collect_pages(files, int(dpi))
283
+ note = ("\nSkipped: " + "; ".join(skipped)) if skipped else ""
284
+
285
+ if not pages:
286
+ yield msg("Nothing readable in that upload." + note +
287
+ "\n\nSupported: .pdf, .png, .jpg, .jpeg, .webp, .bmp, .tif, .tiff")
288
+ return
289
+
290
+ truncated = note
291
+ if len(pages) > MAX_PAGES:
292
+ truncated += f" (truncated from {len(pages)} pages)"
293
+ pages = pages[:MAX_PAGES]
294
+
295
+ md_parts, raw_parts, per_page = [], [], []
296
+ t_start = time.time()
297
+
298
+ for start in range(0, len(pages), PAGES_PER_GPU_CALL):
299
+ chunk = pages[start:start + PAGES_PER_GPU_CALL]
300
+ names = [n for n, _ in chunk]
301
+ imgs = [fit(im, int(max_side)) for _, im in chunk]
302
+
303
+ progress(start / len(pages),
304
+ desc=f"{names[0]} … ({start + 1}-{start + len(chunk)}/{len(pages)})")
305
+
306
+ try:
307
+ raws = _infer_chunk(imgs, prompt_type, int(max_new_tokens))
308
+ except torch.cuda.OutOfMemoryError:
309
+ torch.cuda.empty_cache()
310
+ raws = ["[OUT OF MEMORY β€” lower 'Max image side']"] * len(chunk)
311
+ except Exception as e: # noqa: BLE001
312
+ raws = [f"[FAILED: {type(e).__name__}: {e}]"] * len(chunk)
313
+
314
+ for name, raw in zip(names, raws):
315
+ md = to_markdown(raw)
316
+ md_parts.append(f"\n\n---\n\n## {name}\n\n{md}")
317
+ raw_parts.append(f"===== {name} =====\n{raw}")
318
+ per_page.append((name, md, raw))
319
+
320
+ elapsed = time.time() - t_start
321
+ joined = "\n".join(md_parts)
322
+ yield (joined, joined,
323
+ "\n\n".join(raw_parts),
324
+ f"{len(per_page)}/{len(pages)} pages{truncated} Β· {elapsed:.0f}s "
325
+ f"({elapsed / max(1, len(per_page)):.1f}s/page)",
326
+ None, None)
327
+
328
+ # ---- artefacts ----
329
+ stamp = time.strftime("%Y%m%d_%H%M%S")
330
+ md_path = OUT_DIR / f"chandra_{stamp}.md"
331
+ md_path.write_text("\n".join(md_parts), encoding="utf-8")
332
+
333
+ zip_path = OUT_DIR / f"chandra_{stamp}.zip"
334
+ with zipfile.ZipFile(zip_path, "w", zipfile.ZIP_DEFLATED) as z:
335
+ for name, md, raw in per_page:
336
+ safe = "".join(c if c.isalnum() or c in "-_." else "_" for c in name)
337
+ z.writestr(f"markdown/{safe}.md", md)
338
+ z.writestr(f"raw/{safe}.txt", raw)
339
+ z.writestr("manifest.json", json.dumps({
340
+ "model": MODEL_ID,
341
+ "prompt_type": prompt_type,
342
+ "dpi": dpi,
343
+ "max_side": max_side,
344
+ "pages": [n for n, _, _ in per_page],
345
+ }, indent=2))
346
+
347
+ total = time.time() - t_start
348
+ joined = "\n".join(md_parts)
349
+ yield (joined, joined,
350
+ "\n\n".join(raw_parts),
351
+ f"Done β€” {len(per_page)} page(s){truncated} in {total:.0f}s "
352
+ f"({total / max(1, len(per_page)):.1f}s/page)",
353
+ str(md_path), str(zip_path))
354
+
355
+
356
+ # ---------------------------------------------------------------------------
357
+ # UI
358
+ # ---------------------------------------------------------------------------
359
+
360
+ GR_MAJOR = int(gr.__version__.split(".")[0])
361
+
362
+
363
+ def C(cls, **kw):
364
+ """Build a component, dropping kwargs this Gradio version rejects.
365
+
366
+ Gradio 6 removed Textbox.show_copy_button and moved theme/css from
367
+ Blocks() to launch(). This keeps one file working across 4/5/6.
368
+ """
369
+ try:
370
+ allowed = set(inspect.signature(cls.__init__).parameters)
371
+ if "kwargs" not in allowed:
372
+ kw = {k: v for k, v in kw.items() if k in allowed}
373
+ except (TypeError, ValueError):
374
+ pass
375
+ return cls(**kw)
376
+
377
+
378
+ CSS = """
379
+ #raw_out textarea { font-family: ui-monospace, monospace; font-size: 12px; }
380
+ #status textarea { font-family: ui-monospace, monospace; font-size: 12px; }
381
+ .md_pane { max-height: 640px; overflow-y: auto; }
382
  """
383
+
384
+ _STYLE = dict(theme=gr.themes.Soft(), css=CSS)
385
+ _BLOCKS_KW = {} if GR_MAJOR >= 6 else _STYLE
386
+ _LAUNCH_KW = _STYLE if GR_MAJOR >= 6 else {}
387
+
388
+ with gr.Blocks(title="Chandra OCR 2", **_BLOCKS_KW) as demo:
389
+ gr.Markdown(
390
+ f"""
391
+ # Chandra OCR 2 β€” document β†’ markdown / HTML / JSON
392
+
393
+ Layout-aware OCR from [Datalab](https://datalab.to). Handles tables, math,
394
+ forms, handwriting and 90+ languages, preserving reading order and structure.
395
+
396
+ Upload **PDFs and/or images**; each page is processed separately and results
397
+ stream in below. Capped at **{MAX_PAGES} pages** per run in this demo.
398
+
399
+ Model: [`{MODEL_ID}`](https://huggingface.co/{MODEL_ID}) Β· weights are under a
400
+ modified OpenRAIL-M licence (free for research, personal use, and companies
401
+ under $2M funding/revenue β€” **not** for building a competitor to Datalab's API).
402
  """
403
+ )
404
+
405
+ with gr.Row():
406
+ with gr.Column(scale=1):
407
+ # No file_types filter on purpose: it silently rejects valid files
408
+ # on some Gradio/browser combinations. Validation happens in
409
+ # collect_pages() instead, which can give a useful message.
410
+ files = C(gr.Files, label="PDFs / images", file_count="multiple")
411
+ go = C(gr.Button, value="Run OCR", variant="primary")
412
+
413
+ prompt_type = C(gr.Dropdown, label="Prompt type", choices=PROMPT_TYPES,
414
+ value=DEFAULT_PROMPT_TYPE, allow_custom_value=True,
415
+ info="'ocr_layout' is the documented default.")
416
+
417
+ with gr.Accordion("Advanced", open=False):
418
+ dpi = C(gr.Slider, minimum=100, maximum=400, value=200, step=25,
419
+ label="PDF render DPI",
420
+ info="200-300 suits most scans.")
421
+ max_side = C(gr.Slider, minimum=768, maximum=2560, value=1540, step=64,
422
+ label="Max image side (px)",
423
+ info="Biggest lever on speed and VRAM.")
424
+ max_new = C(gr.Slider, minimum=512, maximum=8192, value=4096, step=256,
425
+ label="Max new tokens",
426
+ info="Layout output is verbose β€” keep this high.")
427
+
428
+ status = C(gr.Textbox, label="Status", lines=3, elem_id="status",
429
+ value=("" if model is not None else
430
+ f"MODEL NOT LOADED β€” {MODEL_ERR}\n"
431
+ f"Check Settings > Hardware (needs ZeroGPU or a GPU tier)."))
432
+
433
+ with gr.Accordion("Diagnostics", open=model is None):
434
+ C(gr.Textbox, label=None, lines=5, value=diagnostics(),
435
+ interactive=False)
436
+
437
+ with gr.Column(scale=2):
438
+ with gr.Tabs():
439
+ with gr.Tab("Rendered"):
440
+ md_view = C(gr.Markdown, value="", elem_classes=["md_pane"])
441
+ with gr.Tab("Markdown source"):
442
+ md_src = C(gr.Textbox, label=None, lines=24,
443
+ show_copy_button=True)
444
+ with gr.Tab("Raw model output"):
445
+ raw_view = C(gr.Textbox, label=None, lines=24,
446
+ elem_id="raw_out", show_copy_button=True)
447
+ with gr.Row():
448
+ md_file = C(gr.File, label="Combined .md")
449
+ zip_file = C(gr.File, label="All pages .zip")
450
+
451
+ go.click(
452
+ run,
453
+ inputs=[files, prompt_type, dpi, max_side, max_new],
454
+ outputs=[md_view, md_src, raw_view, status, md_file, zip_file],
455
+ )
456
 
457
 
458
  if __name__ == "__main__":
459
+ demo.queue(max_size=12).launch(show_error=True, **_LAUNCH_KW)
requirements.txt ADDED
@@ -0,0 +1,11 @@
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Chandra OCR 2 Space
2
+ # torch is preinstalled on ZeroGPU/GPU Spaces β€” do not pin it here.
3
+
4
+ chandra-ocr
5
+ transformers>=4.57.0
6
+ accelerate>=0.34.0
7
+ huggingface_hub>=0.26.0
8
+ pymupdf>=1.24.0
9
+ pillow>=10.0.0
10
+ gradio>=5.0.0
11
+ spaces