Image-Text-to-Text
Transformers
Joblib
Safetensors
nemotron_parse
text-generation
document-parsing
ocr
tables
charts
ground-truth-sft
grpo
verifiable-rewards
routed-refinement
thrisul
conversational
custom_code
Instructions to use sandeeprdy1729/thrisul-parser with libraries, inference providers, notebooks, and local apps. Follow these links to get started.
- Libraries
- Transformers
How to use sandeeprdy1729/thrisul-parser with Transformers:
# Use a pipeline as a high-level helper from transformers import pipeline pipe = pipeline("image-text-to-text", model="sandeeprdy1729/thrisul-parser", trust_remote_code=True) messages = [ { "role": "user", "content": [ {"type": "image", "url": "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/p-blog/candy.JPG"}, {"type": "text", "text": "What animal is on the candy?"} ] }, ] pipe(text=messages)# pip install -U transformers accelerate # Load model directly from transformers import AutoModelForCausalLM model = AutoModelForCausalLM.from_pretrained("sandeeprdy1729/thrisul-parser", trust_remote_code=True, device_map="auto") - Notebooks
- Google Colab
- Kaggle
- Local Apps Settings
- vLLM
How to use sandeeprdy1729/thrisul-parser with vLLM:
Install from pip and serve model
# Install vLLM from pip: pip install vllm # Start the vLLM server: vllm serve "sandeeprdy1729/thrisul-parser" # Call the server using curl (OpenAI-compatible API): curl -X POST "http://localhost:8000/v1/chat/completions" \ -H "Content-Type: application/json" \ --data '{ "model": "sandeeprdy1729/thrisul-parser", "messages": [ { "role": "user", "content": [ { "type": "text", "text": "Describe this image in one sentence." }, { "type": "image_url", "image_url": { "url": "https://cdn.britannica.com/61/93061-050-99147DCE/Statue-of-Liberty-Island-New-York-Bay.jpg" } } ] } ] }'Use Docker
docker model run hf.co/sandeeprdy1729/thrisul-parser
- SGLang
How to use sandeeprdy1729/thrisul-parser with SGLang:
Install from pip and serve model
# Install SGLang from pip: pip install sglang # Start the SGLang server: python3 -m sglang.launch_server \ --model-path "sandeeprdy1729/thrisul-parser" \ --host 0.0.0.0 \ --port 30000 # Call the server using curl (OpenAI-compatible API): curl -X POST "http://localhost:30000/v1/chat/completions" \ -H "Content-Type: application/json" \ --data '{ "model": "sandeeprdy1729/thrisul-parser", "messages": [ { "role": "user", "content": [ { "type": "text", "text": "Describe this image in one sentence." }, { "type": "image_url", "image_url": { "url": "https://cdn.britannica.com/61/93061-050-99147DCE/Statue-of-Liberty-Island-New-York-Bay.jpg" } } ] } ] }'Use Docker images
docker run --gpus all \ --shm-size 32g \ -p 30000:30000 \ -v ~/.cache/huggingface:/root/.cache/huggingface \ --env "HF_TOKEN=<secret>" \ --ipc=host \ lmsysorg/sglang:latest \ python3 -m sglang.launch_server \ --model-path "sandeeprdy1729/thrisul-parser" \ --host 0.0.0.0 \ --port 30000 # Call the server using curl (OpenAI-compatible API): curl -X POST "http://localhost:30000/v1/chat/completions" \ -H "Content-Type: application/json" \ --data '{ "model": "sandeeprdy1729/thrisul-parser", "messages": [ { "role": "user", "content": [ { "type": "text", "text": "Describe this image in one sentence." }, { "type": "image_url", "image_url": { "url": "https://cdn.britannica.com/61/93061-050-99147DCE/Statue-of-Liberty-Island-New-York-Bay.jpg" } } ] } ] }' - Docker Model Runner
How to use sandeeprdy1729/thrisul-parser with Docker Model Runner:
docker model run hf.co/sandeeprdy1729/thrisul-parser
Thrisul-Parser v0.1
Browse filesThis view is limited to 50 files because it contains too many changes. See raw diff
- CONTRIBUTING.md +48 -0
- Dockerfile +19 -0
- LICENSE +57 -0
- README.md +102 -0
- THIRD_PARTY_NOTICES.md +0 -0
- TOKENIZER_ATTRIBUTION.md +41 -0
- bias.md +7 -0
- chat_template.jinja +3 -0
- config.json +207 -0
- docker-compose.yaml +14 -0
- docs/postprocessing.md +446 -0
- docs/serving.md +106 -0
- eval/router_ablation.png +0 -0
- eval/sample_chart.png +0 -0
- eval/sample_table.png +0 -0
- eval/synthetic_eval.png +0 -0
- explainability.md +15 -0
- generation_config.json +13 -0
- hf_logits_processor.py +352 -0
- hf_nemotron_parse_config.py +133 -0
- hf_nemotron_parse_modeling.py +275 -0
- hf_nemotron_parse_processor.py +501 -0
- hf_prefix_geometry.py +149 -0
- latex2html.py +11 -0
- logitsprocs/nemotron_parse_vllm_logitprocs.py +350 -0
- model-cards/nemotron-3-5-parse-subcards/bias.md +7 -0
- model-cards/nemotron-3-5-parse-subcards/explainability.md +15 -0
- model-cards/nemotron-3-5-parse-subcards/privacy.md +2 -0
- model-cards/nemotron-3-5-parse-subcards/safety_security.md +12 -0
- model-cards/nemotron-3-5-parse.md +158 -0
- model.safetensors +3 -0
- nemotron_parse/SOURCE_MANIFEST.json +110 -0
- nemotron_parse/__init__.py +6 -0
- nemotron_parse/batch.py +571 -0
- nemotron_parse/engine/__init__.py +2 -0
- nemotron_parse/engine/caption_relocation.py +89 -0
- nemotron_parse/engine/chart_caption.py +71 -0
- nemotron_parse/engine/chart_context.py +40 -0
- nemotron_parse/engine/convert/__init__.py +2 -0
- nemotron_parse/engine/convert/hybrid_cleanup.py +67 -0
- nemotron_parse/engine/convert/hybrid_inline.py +164 -0
- nemotron_parse/engine/convert/latex2html.py +653 -0
- nemotron_parse/engine/convert/rich_table_fallback.py +105 -0
- nemotron_parse/engine/convert/rich_table_reader.py +164 -0
- nemotron_parse/engine/convert/semantic_styles.py +68 -0
- nemotron_parse/engine/geometry.py +438 -0
- nemotron_parse/engine/header_inference_probe.py +84 -0
- nemotron_parse/engine/markdown.py +61 -0
- nemotron_parse/engine/merge.py +630 -0
- nemotron_parse/engine/policies.py +338 -0
CONTRIBUTING.md
ADDED
|
@@ -0,0 +1,48 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# Contributing
|
| 2 |
+
|
| 3 |
+
## Developer Certificate of Origin (DCO)
|
| 4 |
+
|
| 5 |
+
We require that all contributors "sign-off" on their commits. This certifies that the contribution is your original
|
| 6 |
+
work, or you have rights to submit it under the same license, or a compatible license.
|
| 7 |
+
|
| 8 |
+
Any contribution which contains commits that are not signed off are not accepted.
|
| 9 |
+
|
| 10 |
+
To sign off on a commit, use the `--signoff` (or `-s`) option when you commit your changes:
|
| 11 |
+
|
| 12 |
+
```bash
|
| 13 |
+
git commit --signoff --message "Add cool feature."
|
| 14 |
+
```
|
| 15 |
+
|
| 16 |
+
This appends the following text to your commit message:
|
| 17 |
+
|
| 18 |
+
```text
|
| 19 |
+
Signed-off-by: Your Name <your@email.com>
|
| 20 |
+
```
|
| 21 |
+
|
| 22 |
+
### Full text of the Developer Certificate of Origin
|
| 23 |
+
|
| 24 |
+
```text
|
| 25 |
+
Developer Certificate of Origin
|
| 26 |
+
Version 1.1
|
| 27 |
+
|
| 28 |
+
Copyright (C) 2004, 2006 The Linux Foundation and its contributors.
|
| 29 |
+
|
| 30 |
+
Everyone is permitted to copy and distribute verbatim copies of this license document, but changing it is not allowed.
|
| 31 |
+
|
| 32 |
+
|
| 33 |
+
Developer's Certificate of Origin 1.1
|
| 34 |
+
|
| 35 |
+
By making a contribution to this project, I certify that:
|
| 36 |
+
|
| 37 |
+
(a) The contribution was created in whole or in part by me and I have the right to submit it under the open source license indicated in the file; or
|
| 38 |
+
|
| 39 |
+
(b) The contribution is based upon previous work that, to the best of my knowledge, is covered under an appropriate open source license and I have the right under that license to submit that work with modifications, whether created in whole or in part by me, under the same open source license (unless I am permitted to submit under a different license), as indicated in the file; or
|
| 40 |
+
|
| 41 |
+
(c) The contribution was provided directly to me by some other person who certified (a), (b) or (c) and I have not modified it.
|
| 42 |
+
|
| 43 |
+
(d) I understand and agree that this project and the contribution are public and that a record of the contribution (including all personal information I submit with it, including my sign-off) is maintained indefinitely and may be redistributed consistent with this project or the open source license(s) involved.
|
| 44 |
+
```
|
| 45 |
+
|
| 46 |
+
## License
|
| 47 |
+
|
| 48 |
+
By contributing to this project, you agree that your contributions will be licensed under the [OpenMDW-1.1 License](LICENSE). All contributions must comply with the terms of this license.
|
Dockerfile
ADDED
|
@@ -0,0 +1,19 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
| 2 |
+
# SPDX-License-Identifier: OpenMDW-1.1
|
| 3 |
+
|
| 4 |
+
FROM vllm/vllm-openai:v0.26.0
|
| 5 |
+
LABEL org.opencontainers.image.title="Nemotron 3.5 Parse" \
|
| 6 |
+
org.opencontainers.image.version="3.5.0.post7" \
|
| 7 |
+
org.opencontainers.image.licenses="OpenMDW-1.1"
|
| 8 |
+
WORKDIR /workspace
|
| 9 |
+
RUN python3 -m venv /opt/nemotron-parse-venv
|
| 10 |
+
ENV PATH="/opt/nemotron-parse-venv/bin:$PATH" \
|
| 11 |
+
VIRTUAL_ENV="/opt/nemotron-parse-venv" \
|
| 12 |
+
PYTHONPATH=""
|
| 13 |
+
COPY pyproject.toml LICENSE THIRD_PARTY_NOTICES.md TOKENIZER_ATTRIBUTION.md postprocessing.py latex2html.py /tmp/nemotron-package/
|
| 14 |
+
COPY nemotron_parse /tmp/nemotron-package/nemotron_parse
|
| 15 |
+
COPY vllm_tied_patch /tmp/nemotron-package/vllm_tied_patch
|
| 16 |
+
RUN python3 -m pip install "/tmp/nemotron-package[vllm,client]" && python3 -m pip check
|
| 17 |
+
COPY . /workspace
|
| 18 |
+
ENTRYPOINT ["python3", "-m", "vllm.entrypoints.openai.api_server"]
|
| 19 |
+
CMD ["--model", "/workspace", "--served-model-name", "nemotron-parse", "--trust-remote-code", "--dtype", "bfloat16", "--distributed-executor-backend", "uni", "--tensor-parallel-size", "1", "--max-model-len", "16384", "--max-num-seqs", "192", "--max-num-batched-tokens", "16384", "--no-enable-prefix-caching", "--attention-config", "{\"flash_attn_version\":3}", "--limit-mm-per-prompt", "{\"image\":1}", "--chat-template", "/workspace/chat_template.jinja"]
|
LICENSE
ADDED
|
@@ -0,0 +1,57 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
Unless separately indicated, the model weights and NVIDIA-provided source code
|
| 2 |
+
and documentation are licensed under the OpenMDW-1.1 License.
|
| 3 |
+
The tokenizer included with this model is separately licensed under the
|
| 4 |
+
MIT license. See TOKENIZER_ATTRIBUTION.md for the license text,
|
| 5 |
+
original copyright notice, source, and NVIDIA modification history.
|
| 6 |
+
|
| 7 |
+
------------
|
| 8 |
+
Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
| 9 |
+
|
| 10 |
+
|
| 11 |
+
OpenMDW License Agreement, version 1.1 (OpenMDW-1.1)
|
| 12 |
+
|
| 13 |
+
By exercising rights granted to you under this agreement, you accept and agree
|
| 14 |
+
to its terms.
|
| 15 |
+
|
| 16 |
+
As used in this agreement, "Model Materials" means the materials provided to
|
| 17 |
+
you under this agreement, consisting of: (1) one or more machine learning
|
| 18 |
+
models (including architecture and parameters); and (2) all related artifacts
|
| 19 |
+
(including associated data, documentation and software) that are provided to
|
| 20 |
+
you hereunder.
|
| 21 |
+
|
| 22 |
+
Subject to your compliance with this agreement, permission is hereby granted,
|
| 23 |
+
free of charge, to deal in the Model Materials without restriction, including
|
| 24 |
+
under all copyright, patent, database, and trade secret rights included or
|
| 25 |
+
embodied therein.
|
| 26 |
+
|
| 27 |
+
If you distribute any portion of the Model Materials, you shall retain in your
|
| 28 |
+
distribution (1) a copy of this agreement, and (2) all copyright notices and
|
| 29 |
+
other notices of origin included in the Model Materials that are applicable to
|
| 30 |
+
your distribution.
|
| 31 |
+
If you file, maintain, or voluntarily participate in a lawsuit against any
|
| 32 |
+
person or entity asserting that the Model Materials directly or indirectly
|
| 33 |
+
infringe any patent or copyright, then all rights and grants made to you
|
| 34 |
+
hereunder are terminated, unless that lawsuit was in response to a
|
| 35 |
+
corresponding lawsuit first brought against you.
|
| 36 |
+
|
| 37 |
+
This agreement does not impose any restrictions or obligations with respect to
|
| 38 |
+
any use, modification, or sharing of any outputs generated by using the Model
|
| 39 |
+
Materials.
|
| 40 |
+
|
| 41 |
+
THE MODEL MATERIALS ARE PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS
|
| 42 |
+
OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
| 43 |
+
FITNESS FOR A PARTICULAR PURPOSE, TITLE, NONINFRINGEMENT, ACCURACY, OR THE
|
| 44 |
+
ABSENCE OF LATENT OR OTHER DEFECTS OR ERRORS, WHETHER OR NOT DISCOVERABLE, ALL
|
| 45 |
+
TO THE GREATEST EXTENT PERMISSIBLE UNDER APPLICABLE LAW.
|
| 46 |
+
|
| 47 |
+
YOU ARE SOLELY RESPONSIBLE FOR (1) CLEARING RIGHTS OF OTHER PERSONS THAT MAY
|
| 48 |
+
APPLY TO THE MODEL MATERIALS OR ANY USE THEREOF, INCLUDING WITHOUT LIMITATION
|
| 49 |
+
ANY PERSON'S COPYRIGHTS OR OTHER RIGHTS INCLUDED OR EMBODIED IN THE MODEL
|
| 50 |
+
MATERIALS; (2) OBTAINING ANY NECESSARY CONSENTS, PERMISSIONS OR OTHER RIGHTS
|
| 51 |
+
REQUIRED FOR ANY USE OF THE MODEL MATERIALS; OR (3) PERFORMING ANY DUE
|
| 52 |
+
DILIGENCE OR UNDERTAKING ANY OTHER INVESTIGATIONS INTO THE MODEL MATERIALS OR
|
| 53 |
+
ANYTHING INCORPORATED OR EMBODIED THEREIN.
|
| 54 |
+
IN NO EVENT SHALL THE PROVIDERS OF THE MODEL MATERIALS BE LIABLE FOR ANY CLAIM,
|
| 55 |
+
DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR
|
| 56 |
+
OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE MODEL MATERIALS, THE
|
| 57 |
+
USE THEREOF OR OTHER DEALINGS THEREIN.
|
README.md
ADDED
|
@@ -0,0 +1,102 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
---
|
| 2 |
+
license: other
|
| 3 |
+
license_name: openmdw-1.1
|
| 4 |
+
base_model: nvidia/Nemotron-3.5-Parse
|
| 5 |
+
library_name: transformers
|
| 6 |
+
pipeline_tag: image-text-to-text
|
| 7 |
+
tags:
|
| 8 |
+
- document-parsing
|
| 9 |
+
- ocr
|
| 10 |
+
- tables
|
| 11 |
+
- charts
|
| 12 |
+
- thrisul
|
| 13 |
+
---
|
| 14 |
+
# Thrisul-Parser
|
| 15 |
+
|
| 16 |
+
A derivative of [`nvidia/Nemotron-3.5-Parse`](https://huggingface.co/nvidia/Nemotron-3.5-Parse) (631M parameters) built by TIMPS.
|
| 17 |
+
It adds two things on top of the base parser: **(A) verifier-routed refinement** and **(B) a lightweight resolution self-distillation post-training pass**.
|
| 18 |
+
|
| 19 |
+
## What is new
|
| 20 |
+
1. **Verifier-routed refinement (inference).** The base model's double/triple pass re-decodes *every* detected table/chart at native resolution.
|
| 21 |
+
Thrisul runs the page once, scores each table/chart region with a cheap verifier (decoder token uncertainty, structure sanity checks,
|
| 22 |
+
and a render-vs-crop similarity feature from the vision encoder), and re-decodes **only** the regions a small gradient-boosted router predicts will gain.
|
| 23 |
+
Code: `thrisul_pipeline.py` (`parse_page(..., mode="routed")`), router: `thrisul_router.joblib`.
|
| 24 |
+
2. **Resolution self-distillation (weights).** LoRA (r=16, decoder linears) trained so that the *page-scale view* of a table/chart produces the output the model
|
| 25 |
+
itself gives at *native resolution*, filtered by ground-truth quality (teacher score >= 0.8 and not worse than the page-scale output), with native-resolution anchors and text-page replay to avoid forgetting.
|
| 26 |
+
The LoRA was merged losslessly into the original weights (untouched tensors, including the MTP head, are bit-identical to the base).
|
| 27 |
+
|
| 28 |
+
## Training data (all exact-ground-truth or model-generated; no benchmark data was used for training)
|
| 29 |
+
- Procedurally rendered tables (merged cells, multi-level headers, numeric/currency/percent columns) and charts (bar/line/pie/stacked, labelled and unlabelled) with exact ground truth.
|
| 30 |
+
- 300 real charts from ChartNet (if available in the build session).
|
| 31 |
+
- Targets are the base model's own native-resolution outputs (so the output format is unchanged), quality-filtered against ground truth.
|
| 32 |
+
- Compute: Kaggle free tier (T4), LoRA, 166 optimizer steps.
|
| 33 |
+
|
| 34 |
+
## Results (read the caveats)
|
| 35 |
+
All numbers below were measured by us on **held-out synthetic data and a small ParseBench-proxy sample**, on a **T4** GPU with Transformers (no vLLM).
|
| 36 |
+
They are **not** the official ParseBench score and are **not comparable** to the H100 numbers in the base model's card.
|
| 37 |
+
|
| 38 |
+
Held-out synthetic regions (page-scale input; `routed` refines only regions the router selects):
|
| 39 |
+
|
| 40 |
+
| system | score | refine_frac | sec_per_region | n |
|
| 41 |
+
|---|---|---|---|---|
|
| 42 |
+
| base_single | 0.800 | 0.000 | 0.723 | 150 |
|
| 43 |
+
| base_double(all) | 0.954 | 1.000 | 1.612 | 150 |
|
| 44 |
+
| base_routed | 0.951 | 0.500 | 1.165 | 150 |
|
| 45 |
+
| thr_single | 0.816 | 0.000 | 0.939 | 150 |
|
| 46 |
+
| thr_double(all) | 0.957 | 1.000 | 1.818 | 150 |
|
| 47 |
+
| thr_routed | 0.956 | 0.680 | 1.517 | 150 |
|
| 48 |
+
|
| 49 |
+
Full synthetic pages through the real pipeline (table score):
|
| 50 |
+
|
| 51 |
+
| system | table_score | sec_per_page | pages_per_s | map_fail_rate | regions_refined | n |
|
| 52 |
+
|---|---|---|---|---|---|---|
|
| 53 |
+
| base_single | 0.996 | 5.272 | 0.190 | 0.000 | 0 | 24 |
|
| 54 |
+
| base_double(all) | 0.988 | 7.664 | 0.130 | 0.000 | 24 | 24 |
|
| 55 |
+
| base_routed | 0.988 | 7.305 | 0.137 | 0.000 | 20 | 24 |
|
| 56 |
+
| thr_single | 0.996 | 5.294 | 0.189 | 0.000 | 0 | 24 |
|
| 57 |
+
| thr_routed | 0.997 | 7.504 | 0.133 | 0.000 | 22 | 24 |
|
| 58 |
+
|
| 59 |
+
ParseBench-proxy on real PDFs (first table/chart of the page vs `expected_markdown`; crude proxy, small n):
|
| 60 |
+
|
| 61 |
+
not run (dataset load failed in this session)
|
| 62 |
+
|
| 63 |
+
Verifier ablation (out-of-fold, which signals predict the benefit of refinement):
|
| 64 |
+
|
| 65 |
+
| features | spearman | auc | refine_frac_for_90pct_gain | gain_captured_at_30pct |
|
| 66 |
+
|---|---|---|---|---|
|
| 67 |
+
| entropy+structure | 0.653 | 0.850 | 0.548 | 0.791 |
|
| 68 |
+
| all (entropy+structure+render) | 0.659 | 0.852 | 0.538 | 0.788 |
|
| 69 |
+
| entropy | 0.646 | 0.844 | 0.552 | 0.773 |
|
| 70 |
+
| render only | 0.417 | 0.732 | 0.767 | 0.606 |
|
| 71 |
+
| structure | 0.287 | 0.663 | 0.917 | 0.541 |
|
| 72 |
+
|
| 73 |
+
Text regression check (Thrisul-single vs base-single page markdown similarity): 97.0 %
|
| 74 |
+
|
| 75 |
+
Published reference numbers (self-reported by their authors, different protocol/hardware): Nemotron 3.5 Parse ParseBench 72.74 / 74.92 / 75.20 (single / double / triple pass); Infinity-Parser2-Pro 74.3.
|
| 76 |
+
|
| 77 |
+
## Usage
|
| 78 |
+
```python
|
| 79 |
+
import torch, sys
|
| 80 |
+
from huggingface_hub import snapshot_download
|
| 81 |
+
from transformers import AutoModel, AutoProcessor, GenerationConfig
|
| 82 |
+
path = snapshot_download("sandeeprdy1729/thrisul-parser")
|
| 83 |
+
sys.path.insert(0, path)
|
| 84 |
+
from postprocessing import postprocess_document # shipped from the base model
|
| 85 |
+
import thrisul_pipeline as tp
|
| 86 |
+
model = AutoModel.from_pretrained(path, trust_remote_code=True, dtype=torch.bfloat16).cuda().eval() # use float16 on T4
|
| 87 |
+
processor = AutoProcessor.from_pretrained(path, trust_remote_code=True)
|
| 88 |
+
gen_cfg = GenerationConfig.from_pretrained(path)
|
| 89 |
+
router = tp.Router.load(path + "/thrisul_router.joblib")
|
| 90 |
+
post = lambda text, image_size: postprocess_document(text, image_size=image_size)
|
| 91 |
+
out = tp.parse_page(model, processor, gen_cfg, image, post, mode="routed", router=router) # or mode="single" / "all"
|
| 92 |
+
print(out["md"])
|
| 93 |
+
```
|
| 94 |
+
|
| 95 |
+
## Limitations
|
| 96 |
+
- Trained on synthetic + a small number of real charts; real-document gains are unproven beyond the proxy above.
|
| 97 |
+
- Router features depend on this tokenizer's `<class_*>` tokens and on region<->table mapping; pages where mapping fails fall back to single pass.
|
| 98 |
+
- The router was fitted on crop-level labels and applied to page-level generations (distribution shift).
|
| 99 |
+
- Evaluated on a T4 with Transformers; throughput claims do not transfer to H100/vLLM.
|
| 100 |
+
|
| 101 |
+
## License and attribution
|
| 102 |
+
Derivative work of NVIDIA Nemotron 3.5 Parse; use is governed by the base model's license (OpenMDW 1.1) and the third-party notices in this repository. See `LICENSE` and the notice files copied from the base repo.
|
THIRD_PARTY_NOTICES.md
ADDED
|
The diff for this file is too large to render.
See raw diff
|
|
|
TOKENIZER_ATTRIBUTION.md
ADDED
|
@@ -0,0 +1,41 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# Tokenizer Attribution
|
| 2 |
+
|
| 3 |
+
## Meta Platforms, Inc. and affiliates — Nougat tokenizer — MIT License
|
| 4 |
+
|
| 5 |
+
Component: `nougat/dataset/tokenizer.json`
|
| 6 |
+
|
| 7 |
+
Source: https://github.com/facebookresearch/nougat/blob/5a92920d342fb6acf05fc9b594ccb4053dbe8e7a/nougat/dataset/tokenizer.json
|
| 8 |
+
|
| 9 |
+
Source revision: `5a92920d342fb6acf05fc9b594ccb4053dbe8e7a`
|
| 10 |
+
|
| 11 |
+
Attribution Statements: Modified by NVIDIA. The base 50,000-token vocabulary and token IDs, BPE merge rules, decoder, and original special-token IDs are retained. NVIDIA added multilingual and document-parsing tokens, including chart-related tokens, expanding the tokenizer to 72,505 entries, and modified pre-tokenization.
|
| 12 |
+
|
| 13 |
+
License Text (https://github.com/facebookresearch/nougat/blob/5a92920d342fb6acf05fc9b594ccb4053dbe8e7a/LICENSE)
|
| 14 |
+
|
| 15 |
+
MIT License
|
| 16 |
+
|
| 17 |
+
Copyright (c) Meta Platforms, Inc. and affiliates.
|
| 18 |
+
|
| 19 |
+
Permission is hereby granted, free of charge, to any person obtaining a copy
|
| 20 |
+
of this software and associated documentation files (the "Software"), to deal
|
| 21 |
+
in the Software without restriction, including without limitation the rights
|
| 22 |
+
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
| 23 |
+
copies of the Software, and to permit persons to whom the Software is
|
| 24 |
+
furnished to do so, subject to the following conditions:
|
| 25 |
+
|
| 26 |
+
The above copyright notice and this permission notice shall be included in all
|
| 27 |
+
copies or substantial portions of the Software.
|
| 28 |
+
|
| 29 |
+
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
| 30 |
+
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
| 31 |
+
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
| 32 |
+
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
| 33 |
+
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
| 34 |
+
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
| 35 |
+
SOFTWARE.
|
| 36 |
+
|
| 37 |
+
## NVIDIA tokenizer modifications
|
| 38 |
+
|
| 39 |
+
Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES.
|
| 40 |
+
|
| 41 |
+
License: MIT
|
bias.md
ADDED
|
@@ -0,0 +1,7 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
## Bias
|
| 2 |
+
|
| 3 |
+
Field | Response
|
| 4 |
+
:---- | :----
|
| 5 |
+
Participation considerations from adversely impacted groups [protected classes](https://calcivilrights.ca.gov/disputeresolution/protected-characteristics/) in model design and testing: | None
|
| 6 |
+
Measures taken to mitigate against unwanted bias: | Dataset curation, filtering, and evaluation include review across document types, layouts, languages, and scripts where available. Users should evaluate performance on their target document domains before deployment.
|
| 7 |
+
Bias Metric (If Measured): | Not measured as a demographic fairness metric. Model quality is evaluated through document parsing, OCR, layout, table, chart, and reading-order metrics.
|
chat_template.jinja
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. #}
|
| 2 |
+
{# SPDX-License-Identifier: OpenMDW-1.1 #}
|
| 3 |
+
{% for message in messages %}{{ message['content'] }}{% endfor %}
|
config.json
ADDED
|
@@ -0,0 +1,207 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"architectures": [
|
| 3 |
+
"RadioMBartPrefixForCausalLM"
|
| 4 |
+
],
|
| 5 |
+
"auto_map": {
|
| 6 |
+
"AutoConfig": "hf_nemotron_parse_config.NemotronParseConfig",
|
| 7 |
+
"AutoModel": "hf_nemotron_parse_modeling.NemotronParseForConditionalGeneration",
|
| 8 |
+
"AutoImageProcessor": "hf_nemotron_parse_processor.NemotronParseImageProcessor",
|
| 9 |
+
"AutoProcessor": "hf_nemotron_parse_processor.NemotronParseProcessor"
|
| 10 |
+
},
|
| 11 |
+
"bos_token_id": 0,
|
| 12 |
+
"decoder": {
|
| 13 |
+
"_attn_implementation": "sdpa",
|
| 14 |
+
"_name_or_path": "",
|
| 15 |
+
"activation_dropout": 0.0,
|
| 16 |
+
"activation_function": "gelu",
|
| 17 |
+
"add_cross_attention": false,
|
| 18 |
+
"add_final_layer_norm": true,
|
| 19 |
+
"architectures": null,
|
| 20 |
+
"attention_dropout": 0.0,
|
| 21 |
+
"bad_words_ids": null,
|
| 22 |
+
"begin_suppress_tokens": null,
|
| 23 |
+
"bos_token_id": 0,
|
| 24 |
+
"chunk_size_feed_forward": 0,
|
| 25 |
+
"classifier_dropout": 0.0,
|
| 26 |
+
"cross_attention_hidden_size": null,
|
| 27 |
+
"d_model": 1024,
|
| 28 |
+
"decoder_attention_heads": 16,
|
| 29 |
+
"decoder_ffn_dim": 4096,
|
| 30 |
+
"decoder_layerdrop": 0.0,
|
| 31 |
+
"decoder_layers": 10,
|
| 32 |
+
"decoder_start_token_id": null,
|
| 33 |
+
"diversity_penalty": 0.0,
|
| 34 |
+
"do_sample": false,
|
| 35 |
+
"dropout": 0.1,
|
| 36 |
+
"early_stopping": false,
|
| 37 |
+
"encoder_attention_heads": 16,
|
| 38 |
+
"encoder_ffn_dim": 4096,
|
| 39 |
+
"encoder_layerdrop": 0.0,
|
| 40 |
+
"encoder_layers": 12,
|
| 41 |
+
"encoder_no_repeat_ngram_size": 0,
|
| 42 |
+
"eos_token_id": 2,
|
| 43 |
+
"exponential_decay_length_penalty": null,
|
| 44 |
+
"finetuning_task": null,
|
| 45 |
+
"forced_bos_token_id": null,
|
| 46 |
+
"forced_eos_token_id": 2,
|
| 47 |
+
"hidden_size": 1024,
|
| 48 |
+
"id2label": {
|
| 49 |
+
"0": "LABEL_0",
|
| 50 |
+
"1": "LABEL_1",
|
| 51 |
+
"2": "LABEL_2"
|
| 52 |
+
},
|
| 53 |
+
"init_std": 0.02,
|
| 54 |
+
"is_decoder": true,
|
| 55 |
+
"is_encoder_decoder": false,
|
| 56 |
+
"label2id": {
|
| 57 |
+
"LABEL_0": 0,
|
| 58 |
+
"LABEL_1": 1,
|
| 59 |
+
"LABEL_2": 2
|
| 60 |
+
},
|
| 61 |
+
"length_penalty": 1.0,
|
| 62 |
+
"max_length": 20,
|
| 63 |
+
"min_length": 0,
|
| 64 |
+
"model_type": "nemotron_parse_text",
|
| 65 |
+
"no_repeat_ngram_size": 0,
|
| 66 |
+
"num_beam_groups": 1,
|
| 67 |
+
"num_beams": 1,
|
| 68 |
+
"num_hidden_layers": 10,
|
| 69 |
+
"num_return_sequences": 1,
|
| 70 |
+
"output_attentions": false,
|
| 71 |
+
"output_hidden_states": false,
|
| 72 |
+
"output_scores": false,
|
| 73 |
+
"pad_token_id": 1,
|
| 74 |
+
"prefix": null,
|
| 75 |
+
"problem_type": null,
|
| 76 |
+
"pruned_heads": {},
|
| 77 |
+
"remove_invalid_values": false,
|
| 78 |
+
"repetition_penalty": 1.0,
|
| 79 |
+
"return_dict": true,
|
| 80 |
+
"return_dict_in_generate": false,
|
| 81 |
+
"scale_embedding": true,
|
| 82 |
+
"sep_token_id": null,
|
| 83 |
+
"suppress_tokens": null,
|
| 84 |
+
"task_specific_params": null,
|
| 85 |
+
"temperature": 1.0,
|
| 86 |
+
"tf_legacy_loss": false,
|
| 87 |
+
"tie_encoder_decoder": false,
|
| 88 |
+
"tie_word_embeddings": true,
|
| 89 |
+
"tokenizer_class": null,
|
| 90 |
+
"top_k": 50,
|
| 91 |
+
"top_p": 1.0,
|
| 92 |
+
"torch_dtype": "bfloat16",
|
| 93 |
+
"torchscript": false,
|
| 94 |
+
"transformers_version": "5.6.1",
|
| 95 |
+
"typical_p": 1.0,
|
| 96 |
+
"use_bfloat16": true,
|
| 97 |
+
"use_cache": true,
|
| 98 |
+
"vocab_size": 72512,
|
| 99 |
+
"self_attention_kv_heads": 8,
|
| 100 |
+
"cross_attention_kv_heads": 0,
|
| 101 |
+
"causal_vision_prefix": true
|
| 102 |
+
},
|
| 103 |
+
"decoder_start_token_id": 2,
|
| 104 |
+
"encoder": {
|
| 105 |
+
"_name_or_path": "nvidia/C-RADIOv4-SO400M",
|
| 106 |
+
"adaptor_configs": {},
|
| 107 |
+
"adaptor_names": null,
|
| 108 |
+
"architectures": [
|
| 109 |
+
"RADIOModel"
|
| 110 |
+
],
|
| 111 |
+
"args": {
|
| 112 |
+
"cls_token_per_teacher": true,
|
| 113 |
+
"cpe_max_size": 2048,
|
| 114 |
+
"model": "vit_so400m_patch16_224",
|
| 115 |
+
"register_multiple": 10,
|
| 116 |
+
"teachers": [
|
| 117 |
+
{
|
| 118 |
+
"name": "siglip2-g",
|
| 119 |
+
"use_summary": true
|
| 120 |
+
},
|
| 121 |
+
{
|
| 122 |
+
"name": "dino_v3_7b",
|
| 123 |
+
"use_summary": true
|
| 124 |
+
},
|
| 125 |
+
{
|
| 126 |
+
"name": "sam3",
|
| 127 |
+
"use_summary": false
|
| 128 |
+
}
|
| 129 |
+
]
|
| 130 |
+
},
|
| 131 |
+
"auto_map": {
|
| 132 |
+
"AutoConfig": "nvidia/C-RADIOv4-SO400M--hf_model.RADIOConfig",
|
| 133 |
+
"AutoModel": "nvidia/C-RADIOv4-SO400M--hf_model.RADIOModel"
|
| 134 |
+
},
|
| 135 |
+
"chunk_size_feed_forward": 0,
|
| 136 |
+
"dtype": "float32",
|
| 137 |
+
"feature_normalizer_config": null,
|
| 138 |
+
"id2label": {
|
| 139 |
+
"0": "LABEL_0",
|
| 140 |
+
"1": "LABEL_1"
|
| 141 |
+
},
|
| 142 |
+
"inter_feature_normalizer_config": null,
|
| 143 |
+
"is_encoder_decoder": false,
|
| 144 |
+
"label2id": {
|
| 145 |
+
"LABEL_0": 0,
|
| 146 |
+
"LABEL_1": 1
|
| 147 |
+
},
|
| 148 |
+
"max_resolution": 2048,
|
| 149 |
+
"model_type": "",
|
| 150 |
+
"output_attentions": false,
|
| 151 |
+
"output_hidden_states": false,
|
| 152 |
+
"patch_size": 16,
|
| 153 |
+
"preferred_resolution": [
|
| 154 |
+
512,
|
| 155 |
+
512
|
| 156 |
+
],
|
| 157 |
+
"problem_type": null,
|
| 158 |
+
"return_dict": true,
|
| 159 |
+
"version": "c-radio_v4-so400m",
|
| 160 |
+
"vitdet_window_size": null,
|
| 161 |
+
"code_revision": "c0457f5dc26ca145f954cd4fc5bb6114e5705ad8"
|
| 162 |
+
},
|
| 163 |
+
"eos_token_id": 2,
|
| 164 |
+
"image_size": [
|
| 165 |
+
2048,
|
| 166 |
+
1664
|
| 167 |
+
],
|
| 168 |
+
"is_encoder_decoder": false,
|
| 169 |
+
"max_sequence_length": 16384,
|
| 170 |
+
"model_type": "nemotron_parse",
|
| 171 |
+
"pad_token_id": 1,
|
| 172 |
+
"tie_word_embeddings": true,
|
| 173 |
+
"torch_dtype": "bfloat16",
|
| 174 |
+
"transformers_version": "5.6.1",
|
| 175 |
+
"vocab_size": 72512,
|
| 176 |
+
"class_token_start_idx": 52315,
|
| 177 |
+
"legacy_use_hf_vision_model": true,
|
| 178 |
+
"vision_backbone_name": "nvidia/C-RADIOv4-SO400M",
|
| 179 |
+
"vision_spatial_dim": 1152,
|
| 180 |
+
"vision_summary_dim": 2304,
|
| 181 |
+
"architecture_benchmark": {
|
| 182 |
+
"backbone": "nvidia/C-RADIOv4-SO400M",
|
| 183 |
+
"attention_mode": "global",
|
| 184 |
+
"spatial_dim": 1152,
|
| 185 |
+
"summary_dim": 2304
|
| 186 |
+
},
|
| 187 |
+
"navit_disable_cpe_augmentation": true,
|
| 188 |
+
"self_attention_kv_heads": 8,
|
| 189 |
+
"cross_attention_kv_heads": 0,
|
| 190 |
+
"legacy_causal_vision_prefix": true,
|
| 191 |
+
"legacy_causal_prefix_context_limit": 16384,
|
| 192 |
+
"image_token_index": 1,
|
| 193 |
+
"num_hidden_layers": 10,
|
| 194 |
+
"decoder_layers": 10,
|
| 195 |
+
"hidden_size": 1024,
|
| 196 |
+
"intermediate_size": 4096,
|
| 197 |
+
"num_attention_heads": 16,
|
| 198 |
+
"num_key_value_heads": 8,
|
| 199 |
+
"max_position_embeddings": 16384,
|
| 200 |
+
"navit_position_embedding_mode": "content_normalized_2d_interpolate",
|
| 201 |
+
"navit_packing": true,
|
| 202 |
+
"navit_coordinate_frame": "content_normalized",
|
| 203 |
+
"virtual_canvas_width": 1664,
|
| 204 |
+
"virtual_canvas_height": 2048,
|
| 205 |
+
"virtual_canvas_placement": "aligned_center",
|
| 206 |
+
"num_extra_heads": 1
|
| 207 |
+
}
|
docker-compose.yaml
ADDED
|
@@ -0,0 +1,14 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
services:
|
| 2 |
+
nemotron-parse:
|
| 3 |
+
build: .
|
| 4 |
+
image: nemotron-parse:3.5.0.post7
|
| 5 |
+
deploy:
|
| 6 |
+
resources:
|
| 7 |
+
reservations:
|
| 8 |
+
devices:
|
| 9 |
+
- driver: nvidia
|
| 10 |
+
count: 1
|
| 11 |
+
capabilities: [gpu]
|
| 12 |
+
shm_size: 8gb
|
| 13 |
+
ports:
|
| 14 |
+
- "8000:8000"
|
docs/postprocessing.md
ADDED
|
@@ -0,0 +1,446 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# Document postprocessing and multipass parsing
|
| 2 |
+
|
| 3 |
+
The public helpers operate on model predictions and ordinary page images. They do
|
| 4 |
+
not load benchmark labels, scorer code, category IDs, reference answers, or
|
| 5 |
+
page-specific rules. They are adapted from the general conversion and crop-merge
|
| 6 |
+
components of the bundled ParseBench pipeline. The pinned reproduction bundle
|
| 7 |
+
remains unchanged so its recorded hashes and results stay reproducible.
|
| 8 |
+
|
| 9 |
+
## Installation
|
| 10 |
+
|
| 11 |
+
For vLLM, use the model card's repository-source installation:
|
| 12 |
+
|
| 13 |
+
```bash
|
| 14 |
+
hf download nvidia/Nemotron-3.5-Parse --local-dir nemotron-parse
|
| 15 |
+
cd nemotron-parse
|
| 16 |
+
python -m pip install '.[vllm,client]'
|
| 17 |
+
```
|
| 18 |
+
|
| 19 |
+
This installs the plugin, public `nemotron_parse` Python package, compatibility
|
| 20 |
+
modules, and `nemotron-parse` command. No prebuilt wheel is required.
|
| 21 |
+
The repository source installs the vLLM plugin, parsing helpers and CLI together.
|
| 22 |
+
Archived release artifacts and their digests are listed in the
|
| 23 |
+
[build evidence](../reference/package/README.md).
|
| 24 |
+
Use version 3.5.0.post5 or later for the restored merge default and post7 or later
|
| 25 |
+
for the corrected third-party notices.
|
| 26 |
+
|
| 27 |
+
For Transformers, keep the model card's inference dependencies. To download the
|
| 28 |
+
helpers without installing the package:
|
| 29 |
+
|
| 30 |
+
```bash
|
| 31 |
+
python -m pip install 'beautifulsoup4==4.13.3' 'Pillow==12.2.0' 'numpy==2.2.6' 'httpx==0.28.1'
|
| 32 |
+
hf download nvidia/Nemotron-3.5-Parse \
|
| 33 |
+
--include 'postprocessing.py' 'latex2html.py' 'nemotron_parse/**' --local-dir .
|
| 34 |
+
```
|
| 35 |
+
|
| 36 |
+
Run Python in that directory, or put it on your Python import path. Basic text/table
|
| 37 |
+
postprocessing only needs Beautiful Soup; the image wrappers also need Pillow and
|
| 38 |
+
NumPy. The HTTP backend needs httpx. Downloaded source and installed packages
|
| 39 |
+
provide the same public API.
|
| 40 |
+
|
| 41 |
+
## Render an existing prediction
|
| 42 |
+
|
| 43 |
+
```python
|
| 44 |
+
from nemotron_parse import postprocess_document, PostprocessOptions
|
| 45 |
+
|
| 46 |
+
options = PostprocessOptions(
|
| 47 |
+
infer_headers=True,
|
| 48 |
+
render_charts=True,
|
| 49 |
+
split_chart_value_percent=True,
|
| 50 |
+
relocate_chart_captions=True,
|
| 51 |
+
)
|
| 52 |
+
document = postprocess_document(raw_model_output, image_size=(width, height), options=options)
|
| 53 |
+
print(document['markdown'])
|
| 54 |
+
print(document['blocks'])
|
| 55 |
+
```
|
| 56 |
+
|
| 57 |
+
`raw_out` contains the cleaned model serialization. Each block has its original
|
| 58 |
+
class, normalized `bbox`, literal `text`, and rendered `content`. Supplying
|
| 59 |
+
`image_size` adds clipped, ordered `bbox_pixels`. `markdown` renders the complete
|
| 60 |
+
page in model order, with section headings and HTML tables for span preservation.
|
| 61 |
+
Raw per-block text remains available independently of presentation changes.
|
| 62 |
+
Malformed table markup that cannot be converted is preserved rather than passed
|
| 63 |
+
through a lossy row splitter; document `warnings` identify these blocks.
|
| 64 |
+
|
| 65 |
+
Default cleanup preserves empty cells, spans, line breaks, nested inline styles,
|
| 66 |
+
and escaped literals. It removes exact repeated blocks at identical coordinates,
|
| 67 |
+
repairs empty incomplete serialization prefixes, removes internal continuation
|
| 68 |
+
markers, and normalizes existing bullets and soft hyphens in textual roles.
|
| 69 |
+
Code/math roles and marked code/math text are protected from soft-hyphen cleanup.
|
| 70 |
+
Table conversion is applied to Table/Chart blocks; literal LaTeX in prose, code,
|
| 71 |
+
formulas, or Picture blocks retains its source representation.
|
| 72 |
+
Header inference uses predicted span topology and existing bold styling, never an
|
| 73 |
+
English vocabulary list. It can still make mistakes; disable `infer_headers` to
|
| 74 |
+
use only explicit predicted headers. Unknown classes and page numbers are retained.
|
| 75 |
+
|
| 76 |
+
Chart titles, subtitles, and explicit panel labels can be rendered with their
|
| 77 |
+
associated tables. Two guarded presentation choices are enabled by default and can be disabled:
|
| 78 |
+
|
| 79 |
+
- `split_chart_value_percent=True` splits consistently predicted numeric pairs
|
| 80 |
+
such as `32 (16.6%)` into Value and Percent subcolumns. It requires at least two
|
| 81 |
+
matching nonempty cells, adds those column labels, and does not calculate values.
|
| 82 |
+
- `relocate_chart_captions=True` associates a plain caption immediately above a
|
| 83 |
+
chart when geometry and text-preservation checks pass. Notes below remain
|
| 84 |
+
standalone. Set `relocate_chart_captions=False` to retain the standalone caption in its source order.
|
| 85 |
+
|
| 86 |
+
`normalize_text=False`, `deduplicate=False`, and `render_charts=False` disable the
|
| 87 |
+
corresponding cleanup/rendering choices. The existing root imports continue to
|
| 88 |
+
work:
|
| 89 |
+
|
| 90 |
+
```python
|
| 91 |
+
from postprocessing import extract_classes_bboxes, postprocess_text
|
| 92 |
+
from latex2html import convert_html_table_to_json_hierarchical
|
| 93 |
+
```
|
| 94 |
+
|
| 95 |
+
`postprocess_text(..., cls='Table', table_format=...)` supports LaTeX, HTML,
|
| 96 |
+
Markdown, record JSON, hierarchical JSON, and CSV. HTML retains merged-cell
|
| 97 |
+
structure. Flat exports expand spans and use neutral column names for headerless
|
| 98 |
+
tables without consuming the first data row as a header. Duplicate column names
|
| 99 |
+
are disambiguated without overwriting values. Hierarchical JSON retains the
|
| 100 |
+
existing `headers`, `columns`, and flat `data` fields and additionally exposes
|
| 101 |
+
`body_cells` with explicit text, span, and header-role metadata.
|
| 102 |
+
|
| 103 |
+
## Parse a page with a running vLLM server
|
| 104 |
+
|
| 105 |
+
Start the server using the model card. Then:
|
| 106 |
+
|
| 107 |
+
```python
|
| 108 |
+
from nemotron_parse.inference import DocumentParser
|
| 109 |
+
|
| 110 |
+
with DocumentParser(base_url='http://localhost:8000/v1', model='nemotron-parse') as parser:
|
| 111 |
+
single = parser.parse('page.png')
|
| 112 |
+
double = parser.parse_double('page.png')
|
| 113 |
+
triple = parser.parse_triple('page.png', triple_regions='union')
|
| 114 |
+
documents = parser.parse_many(['one.png', 'two.png'], mode='double', concurrency=128)
|
| 115 |
+
```
|
| 116 |
+
|
| 117 |
+
| Mode | Work performed |
|
| 118 |
+
| --- | --- |
|
| 119 |
+
| Single | Full-page parse and document rendering |
|
| 120 |
+
| Double | Full-page parse, then Table/Chart crops from its predicted regions |
|
| 121 |
+
| Triple | Full-page parse, separate detection, then Table/Chart crop parsing |
|
| 122 |
+
|
| 123 |
+
A pass denotes a stage, not one request: each retained crop needs its own model
|
| 124 |
+
call. `triple_regions='union'` retains detector regions first and adds parser
|
| 125 |
+
regions after overlap filtering; `'detector'` uses only the separate detector.
|
| 126 |
+
The region budget defaults to 80 per page and can be changed with `max_regions`.
|
| 127 |
+
Crops use 4% plus 12-pixel margins and JPEG quality 95, configurable margins, and
|
| 128 |
+
content-normalized coordinates mapped back to the original image. Tiny crops are
|
| 129 |
+
skipped. Crops are made in memory and no temporary image files remain.
|
| 130 |
+
|
| 131 |
+
The default `merge_policy="alternatives"` preserves differing overlapping
|
| 132 |
+
Table/Chart predictions from the full-page and crop passes. Near-identical blocks
|
| 133 |
+
are deduplicated using geometry and text. Guarded table refinement checks native
|
| 134 |
+
grid validity and compatible column counts; the original block order is preserved
|
| 135 |
+
and crop alternatives are inserted by their predicted geometry. This policy
|
| 136 |
+
matches the reference approach using only model predictions, without labels,
|
| 137 |
+
benchmark categories or scorer feedback.
|
| 138 |
+
|
| 139 |
+
A returned document can contain overlapping or conflicting Table/Chart hypotheses.
|
| 140 |
+
Use `DocumentParser(merge_policy="guarded")` (CLI: `--merge-policy guarded`) when
|
| 141 |
+
one selected prediction per overlapping same-class region is preferable. This
|
| 142 |
+
option applies conservative structural replacement and rejects overlapping
|
| 143 |
+
same-class additions. It can discard a correct alternative, reducing recall.
|
| 144 |
+
Raw stage responses remain available with either policy.
|
| 145 |
+
|
| 146 |
+
Returned results include `stages` with original raw responses and usage, and
|
| 147 |
+
`inference` with request counts, regions, sampling, and crop/rendering settings.
|
| 148 |
+
Empty, failed, filtered, or token-limited generations raise `GenerationError`;
|
| 149 |
+
its `stage` and `response` expose the failed response. Successful stage responses
|
| 150 |
+
from the same page are only returned when parsing completes. No automatic retry
|
| 151 |
+
policy changes the request settings.
|
| 152 |
+
|
| 153 |
+
Sampling is configurable per stage:
|
| 154 |
+
|
| 155 |
+
```python
|
| 156 |
+
from nemotron_parse.inference import DocumentParser, GenerationSettings
|
| 157 |
+
|
| 158 |
+
with DocumentParser(
|
| 159 |
+
fullpage=GenerationSettings(temperature=0, top_k=1, repetition_penalty=1.1),
|
| 160 |
+
detection=GenerationSettings(temperature=0, top_k=1, max_tokens=4096),
|
| 161 |
+
crop=GenerationSettings(temperature=0, top_k=1, repetition_penalty=1.0),
|
| 162 |
+
) as parser:
|
| 163 |
+
document = parser.parse_double('page.png')
|
| 164 |
+
```
|
| 165 |
+
|
| 166 |
+
Defaults are temperature 0.2 / repetition penalty 1.05 for full-page and detection,
|
| 167 |
+
and 0.1 / 1.0 for crops, top-k 50, top-p 1.0, seed 0. Token caps are 8,744 for
|
| 168 |
+
parsing and 4,096 for detection. These are practical starting settings rather than
|
| 169 |
+
a claim of optimal sampling for every document type.
|
| 170 |
+
|
| 171 |
+
## Use an already loaded Transformers model
|
| 172 |
+
|
| 173 |
+
Continue from the model card's Transformers setup:
|
| 174 |
+
|
| 175 |
+
```python
|
| 176 |
+
from nemotron_parse.inference import DocumentParser, TransformersBackend
|
| 177 |
+
|
| 178 |
+
backend = TransformersBackend(model, processor, generation_config=generation_config)
|
| 179 |
+
with DocumentParser(backend=backend) as parser:
|
| 180 |
+
document = parser.parse_triple(image)
|
| 181 |
+
```
|
| 182 |
+
|
| 183 |
+
This uses the same image/crop orchestration without a server or plugin. The adapter
|
| 184 |
+
serializes calls to the shared model, uses inference mode, applies stage settings,
|
| 185 |
+
and restores the CPU/current-GPU RNG after generation. Model weights and processor
|
| 186 |
+
remain caller-owned. Reaching the token cap is conservatively reported as
|
| 187 |
+
truncation, including when generation forces an EOS at the cap.
|
| 188 |
+
|
| 189 |
+
Other inference implementations can supply a callable
|
| 190 |
+
`backend(image, *, prompt, stage, settings)` returning
|
| 191 |
+
`{'raw_out': text, 'finish_reason': 'stop', 'usage': {...}}`. A string return is also
|
| 192 |
+
accepted, but then the custom backend must detect truncation itself. Custom
|
| 193 |
+
backends used with `parse_many` should support concurrent calls or use concurrency 1.
|
| 194 |
+
|
| 195 |
+
## Command line
|
| 196 |
+
|
| 197 |
+
```bash
|
| 198 |
+
nemotron-parse --image page.png --mode double --output parsed-page
|
| 199 |
+
nemotron-parse --image-dir pages --mode triple --triple-regions union --output parsed-pages
|
| 200 |
+
```
|
| 201 |
+
|
| 202 |
+
The download-only equivalent is `python -m nemotron_parse.inference` with the same
|
| 203 |
+
arguments. The default outputs are compact `documents.jsonl.gz` and `metrics.json`, containing
|
| 204 |
+
blocks, raw stages, settings and timing. Add `--write-page-files` for individual
|
| 205 |
+
Markdown and JSON exports. Use a new/empty output directory. Files are written only
|
| 206 |
+
after all pages complete. The API accepts PIL images and image paths; the CLI
|
| 207 |
+
accepts image files. Render PDF pages to images before calling it.
|
| 208 |
+
|
| 209 |
+
## Validation and scope
|
| 210 |
+
|
| 211 |
+
Run the focused CPU regression suite from the downloaded repository:
|
| 212 |
+
|
| 213 |
+
```bash
|
| 214 |
+
python -m pip install pytest
|
| 215 |
+
python -m pytest -q tests/test_postprocessing.py tests/test_inference.py tests/test_batch.py
|
| 216 |
+
```
|
| 217 |
+
|
| 218 |
+
The 86 focused tests passed from both source and the installed wheel. They cover
|
| 219 |
+
table structure, text preservation, chart defaults and opt-outs, crop geometry,
|
| 220 |
+
merge guards, HTTP handling, explicit truncation reporting, batch ordering,
|
| 221 |
+
independent-stage overlap, early crop submission, a shared concurrency bound,
|
| 222 |
+
and staged/pipelined region and output equivalence with deterministic responses.
|
| 223 |
+
The Transformers adapter is checked with a small CPU model stub. The real-model
|
| 224 |
+
H100 runs and official full-page scoring are recorded below and in
|
| 225 |
+
[the validation receipt](../reference/validation/postprocessing.json). The pinned reproduction
|
| 226 |
+
bundle and its recorded headline result remain unchanged.
|
| 227 |
+
|
| 228 |
+
## Prepared batches and timing
|
| 229 |
+
|
| 230 |
+
`parser.prepare(images, workers=8)` returns a reusable in-memory batch containing
|
| 231 |
+
decoded RGB pixels and encoded image requests. Existing RGB JPEG/PNG/WebP payloads
|
| 232 |
+
are retained without another lossy encode. Other inputs are encoded as PNG.
|
| 233 |
+
The preparation time is reported separately, and cached images can occupy several
|
| 234 |
+
GB for thousands of document pages.
|
| 235 |
+
|
| 236 |
+
`parser.parse_prepared(batch, mode="double", concurrency=128, postprocess_workers=8)`
|
| 237 |
+
returns a `BatchResult` with `documents` and `metrics`. The default
|
| 238 |
+
`schedule="pipelined"` overlaps independent full-page and detection requests,
|
| 239 |
+
submits detector crops as soon as their regions are ready, and adds parser-only
|
| 240 |
+
regions in the original union order. Double-pass crops similarly start as each
|
| 241 |
+
page finishes its full-page request. Ready crops take the next available HTTP
|
| 242 |
+
slot, and new pages fill remaining slots. All stages share one bounded
|
| 243 |
+
`concurrency` window, including requests waiting inside vLLM. Crop preparation
|
| 244 |
+
uses separate CPU threads; finished pages merge and render while other pages
|
| 245 |
+
continue inference. Large batches use spawned rendering processes.
|
| 246 |
+
|
| 247 |
+
Run standalone scripts under a main guard, or choose `postprocess_workers=1`.
|
| 248 |
+
HTTP `parse_many` uses this same pipelined batch engine and returns only the
|
| 249 |
+
documents. Its default request window is 128; a custom backend retains the
|
| 250 |
+
four-worker default. Pass `concurrency=1` for a backend that requires serial calls.
|
| 251 |
+
Use `schedule="staged"` for stage-barrier comparisons. A custom Transformers
|
| 252 |
+
backend still serializes its shared model calls.
|
| 253 |
+
|
| 254 |
+
Metrics include separately reported initial preparation, pipeline elapsed time,
|
| 255 |
+
request counts, token-limit counts, and processed-page and fully completed-page
|
| 256 |
+
throughput. Pipeline time excludes
|
| 257 |
+
initial preparation and output-file writing; model-dependent crop preparation is
|
| 258 |
+
included. A benchmark should also report serialization and warmup/server-startup
|
| 259 |
+
exclusions explicitly. The same prepared batch can be reused for different modes;
|
| 260 |
+
each call generates its own responses and provides an independent pipeline timing.
|
| 261 |
+
Pipelined stage spans overlap and must not be added together.
|
| 262 |
+
`request_service_seconds` sums concurrent request latency, including server queue
|
| 263 |
+
wait; it is not GPU compute time. `peak_inflight_requests` records the client
|
| 264 |
+
window, while actual server running/waiting requests require vLLM metrics.
|
| 265 |
+
|
| 266 |
+
Strict generation handling remains the default. Set `allow_partial=True` (CLI:
|
| 267 |
+
`--allow-partial`) only when partial results are useful: capped raw responses are
|
| 268 |
+
preserved, affected pages have `complete=False` and `generation_warnings`, and
|
| 269 |
+
truncated crops are excluded from refinement. HTTP/transport errors and empty
|
| 270 |
+
responses still raise. A successful benchmark execution does not imply that every
|
| 271 |
+
page completed within its token limit; inspect both coverage and throughput.
|
| 272 |
+
|
| 273 |
+
The CLI supports `--no-split-chart-value-percent` and
|
| 274 |
+
`--no-relocate-chart-captions` to disable the default presentation rules.
|
| 275 |
+
|
| 276 |
+
## Measured prepared-batch performance
|
| 277 |
+
|
| 278 |
+
The following matched scheduler comparison used the earlier **guarded** merge
|
| 279 |
+
policy in package 3.5.0.post4. The restored default and its fresh results are in
|
| 280 |
+
[accuracy restoration](#accuracy-restoration-with-the-faster-scheduler).
|
| 281 |
+
|
| 282 |
+
Measured on all 2,078 unique ParseBench pages with one H100 80GB, vLLM 0.26.0,
|
| 283 |
+
`--distributed-executor-backend uni`, tensor parallelism 1, one-token MTP, FA3,
|
| 284 |
+
`max_num_seqs=192`, `max_num_batched_tokens=16384`, and `max_model_len=16384`.
|
| 285 |
+
Both chart presentation rules were enabled. Each run generated all of its own
|
| 286 |
+
responses. The same-server comparison used the same prepared images and
|
| 287 |
+
per-mode concurrency limit for both schedulers.
|
| 288 |
+
|
| 289 |
+
| Mode | Concurrency | Staged pages/s | Pipelined pages/s | Pipelined elapsed | Fully completed pages |
|
| 290 |
+
| --- | ---: | ---: | ---: | ---: | ---: |
|
| 291 |
+
| Single | 128 | 14.56 | 14.91 | 139.4 s | 2047/2078 |
|
| 292 |
+
| Double | 256 | 11.16 | 11.87 | 175.0 s | 2038/2078 |
|
| 293 |
+
| Triple | 128 | 7.31 | 10.90 | 190.6 s | 2039/2078 |
|
| 294 |
+
|
| 295 |
+
Initial decoding/encoding of the full batch took 9.7 seconds and was excluded
|
| 296 |
+
from these rates. Timed work includes all model calls, predicted-region crop
|
| 297 |
+
preparation, merge and rendering; startup, warmup and output serialization are
|
| 298 |
+
excluded. The receipt separately reports serialization-inclusive rates.
|
| 299 |
+
|
| 300 |
+
These diagnostic runs explicitly enabled `allow_partial=True`. Token-limited
|
| 301 |
+
responses remain visible and affected pages count as incomplete. The table reports
|
| 302 |
+
processed-page throughput; fully completed-page throughput and cap counts are in
|
| 303 |
+
[the benchmark receipt](../reference/validation/batch_benchmarks.json). Strict failure handling remains
|
| 304 |
+
the API default. A successful scoring run does not make truncated generations complete.
|
| 305 |
+
|
| 306 |
+
Concurrency 128 and 256 were compared on a 256-page pilot; the selected settings
|
| 307 |
+
above were then checked on the full batch. The earlier full single-pass sweep reached
|
| 308 |
+
15.09 pages/s at 256. These are measured serving profiles, not a throughput ceiling.
|
| 309 |
+
Sampling is nonzero: changed batching can change output lengths and crop counts.
|
| 310 |
+
The receipt includes request and completion-token totals for each run.
|
| 311 |
+
Stages share one bounded client request window. In the pipelined runner, crops
|
| 312 |
+
receive available slots while independent full-page/detection requests continue.
|
| 313 |
+
Stage spans overlap; adding them does not produce elapsed time. Queue averages
|
| 314 |
+
include pipeline ramp-up and final drain, so mean running-request counts need not
|
| 315 |
+
equal the configured limit.
|
| 316 |
+
|
| 317 |
+
### Quality and the reference output contract
|
| 318 |
+
|
| 319 |
+
| Previous guarded merge, pipelined | Official overall (%) |
|
| 320 |
+
| --- | ---: |
|
| 321 |
+
| Single | 72.79 |
|
| 322 |
+
| Double | 73.28 |
|
| 323 |
+
| Triple | 73.48 |
|
| 324 |
+
|
| 325 |
+
Each full-run output shown here was freshly scored on all 2,078 pages with the
|
| 326 |
+
pinned official scorer
|
| 327 |
+
`9af4102e26b51f234ab5f037e19fec05c0696ec6`, with all 241 Python source hashes
|
| 328 |
+
verified and zero official scoring errors. Capped generations are included in scoring.
|
| 329 |
+
|
| 330 |
+
The pinned reference merge can retain an original Table/Chart block and an
|
| 331 |
+
overlapping crop alternative when their text differs. The chart evaluator searches
|
| 332 |
+
rendered tables for the best matching data. The previous guarded policy uses
|
| 333 |
+
replacement and rejects overlapping same-class crop additions; raw alternatives
|
| 334 |
+
remain in `stages.crops` for inspection.
|
| 335 |
+
|
| 336 |
+
A controlled replay of identical baseline triple responses scored 73.44
|
| 337 |
+
with the previous guarded merge and 75.19 with reference-style merging of
|
| 338 |
+
the same completed crops. The latter inserted 1,660 overlapping alternative blocks.
|
| 339 |
+
This identified the output contract as the cause of the score gap. Package
|
| 340 |
+
3.5.0.post5 restores alternatives as the default, keeping the faster scheduler.
|
| 341 |
+
Use the separately pinned reproduction pipeline for the exact headline 75.20
|
| 342 |
+
contract. The receipt includes category metrics and the merge-only replay checks.
|
| 343 |
+
|
| 344 |
+
## Accuracy restoration with the faster scheduler
|
| 345 |
+
|
| 346 |
+
Package 3.5.0.post5 defaults to `merge_policy="alternatives"`. The scheduler,
|
| 347 |
+
model settings, region selection and chart presentation defaults are unchanged.
|
| 348 |
+
The merge retains differing full-page/crop hypotheses as the reference does;
|
| 349 |
+
it adds no model calls. `merge_policy="guarded"` remains available explicitly.
|
| 350 |
+
|
| 351 |
+
Fresh model inference on all 2,078 pages with one H100, the same uni/FA3/MTP server
|
| 352 |
+
profile, and prepared inputs produced:
|
| 353 |
+
|
| 354 |
+
| Pipeline | Concurrency | Pages/s | Elapsed | Official overall (%) | Fully completed pages |
|
| 355 |
+
| --- | ---: | ---: | ---: | ---: | ---: |
|
| 356 |
+
| Double | 256 | 11.64 | 178.5 s | 74.92 | 2043/2078 |
|
| 357 |
+
| Triple | 128 | 10.61 | 195.8 s | 75.20 | 2038/2078 |
|
| 358 |
+
|
| 359 |
+
Timing includes crop preparation and merge/rendering. Preparation, startup, warmup
|
| 360 |
+
and serialization are excluded; the receipt also reports serialization-inclusive
|
| 361 |
+
and fully completed-page rates. These diagnostic runs explicitly used
|
| 362 |
+
`allow_partial=True`; token-limit cases remain marked incomplete.
|
| 363 |
+
|
| 364 |
+
The same saved responses from the earlier staged and pipelined runs were also
|
| 365 |
+
rescored with the restored merge. This isolates the merge from generation:
|
| 366 |
+
|
| 367 |
+
| Saved responses | Guarded overall (%) | Alternatives overall (%) |
|
| 368 |
+
| --- | ---: | ---: |
|
| 369 |
+
| Double Staged | 73.25 | 74.88 |
|
| 370 |
+
| Double Pipelined | 73.28 | 74.96 |
|
| 371 |
+
| Triple Staged | 73.51 | 75.23 |
|
| 372 |
+
| Triple Pipelined | 73.48 | 75.23 |
|
| 373 |
+
|
| 374 |
+
All six outputs were freshly scored on 2,078 unique pages with the pinned official
|
| 375 |
+
scorer and all 241 source hashes verified, with zero official scoring errors.
|
| 376 |
+
Batching at nonzero temperature can change individual responses; this does not
|
| 377 |
+
establish bit-identical generations or a guarantee for every document.
|
| 378 |
+
|
| 379 |
+
The reference-compatible score includes alternative tables/charts. The evaluator
|
| 380 |
+
can match the best alternative, so score recovery does not prove that a single
|
| 381 |
+
chosen transcription is equally accurate. The alternatives are present in the
|
| 382 |
+
actual returned document and its raw stages. Use the guarded policy for a more
|
| 383 |
+
compact document when that tradeoff fits the application.
|
| 384 |
+
|
| 385 |
+
See [the restoration receipt](../reference/validation/accuracy_restoration.json) for category
|
| 386 |
+
scores, generation coverage, exact settings, request/token counts, hashes and
|
| 387 |
+
controlled replay results.
|
| 388 |
+
|
| 389 |
+
|
| 390 |
+
## Inference throughput with prepared pages
|
| 391 |
+
|
| 392 |
+
For prediction throughput, use `result.metrics["inference_pages_per_second"]`
|
| 393 |
+
(package 3.5.0.post6 or later). Its wall clock runs from the first model request
|
| 394 |
+
start to the last model response across all passes. It includes server queueing,
|
| 395 |
+
image encoding inside the server, generation, HTTP transfer and any gaps waiting
|
| 396 |
+
for predicted-region selection, crop encoding or dispatch. It excludes initial
|
| 397 |
+
page preparation, startup/warmup, final merge/rendering after inference, and writes.
|
| 398 |
+
`pipeline_seconds` and `processed_pages_per_second` retain their earlier meanings,
|
| 399 |
+
including merge/rendering, for applications that need complete documents.
|
| 400 |
+
Concurrent rendering can still compete for CPU during inference; that observed
|
| 401 |
+
contention remains in the wall time. We do not subtract overlapping CPU time or
|
| 402 |
+
sum concurrent request latencies to claim a higher rate.
|
| 403 |
+
|
| 404 |
+
The fastest observed profiles compatible with the restored output quality are:
|
| 405 |
+
|
| 406 |
+
| Pipeline | Concurrency | Schedule | Prediction pages/s | Inference wall time | Fully completed pages |
|
| 407 |
+
| --- | ---: | --- | ---: | ---: | ---: |
|
| 408 |
+
| Single | 256 | Staged | **15.42** | 134.78 s | 2045/2078 |
|
| 409 |
+
| Double | 256 | Pipelined | **11.65** | 178.35 s | 2043/2078 |
|
| 410 |
+
| Triple | 128 | Pipelined | **10.62** | 195.70 s | 2038/2078 |
|
| 411 |
+
|
| 412 |
+
These use 2,078 unique prepared ParseBench page images on one H100 80GB and the
|
| 413 |
+
same uni/FA3/MTP serving profile. The single-pass measurement predates the merge
|
| 414 |
+
restoration; it has no crops and is unaffected by that merge change. Its saved
|
| 415 |
+
stage clock includes executor setup/teardown and is slightly conservative.
|
| 416 |
+
Double/triple times are reconstructed exactly from saved stage start/end spans,
|
| 417 |
+
including gaps between calls. These are existing measured runs, not a new GPU
|
| 418 |
+
run of the timing-only post6 release. See the [timing receipt](../reference/validation/inference_benchmarks.json).
|
| 419 |
+
|
| 420 |
+
The earlier guarded-merge runs reached 11.88/10.91 prediction pages/s for
|
| 421 |
+
double/triple, respectively, but had the lower scores documented above. Use the
|
| 422 |
+
restored profiles when comparing speed while preserving the reference output
|
| 423 |
+
contract. The restored double/triple scores are 74.92/75.20.
|
| 424 |
+
|
| 425 |
+
This is realistic batch prediction throughput when pages are prepared ahead of
|
| 426 |
+
time and enough work is queued. It is not PDF ingestion throughput or single-page
|
| 427 |
+
latency, and is not a measured hardware ceiling. Document complexity, generation
|
| 428 |
+
length, crop count, CPU capacity, network placement and batching affect speed.
|
| 429 |
+
The numerator includes token-limited pages; these diagnostic runs explicitly
|
| 430 |
+
allow partial output. Use `complete_inference_pages_per_second` for the number
|
| 431 |
+
of fully completed pages divided by the same wall time. Strict errors remain the
|
| 432 |
+
default outside these diagnostic runs.
|
| 433 |
+
|
| 434 |
+
```python
|
| 435 |
+
prepared = parser.prepare(page_images, workers=8) # outside inference clock
|
| 436 |
+
result = parser.parse_prepared(
|
| 437 |
+
prepared, mode="triple", concurrency=128, postprocess_workers=8,
|
| 438 |
+
)
|
| 439 |
+
print(result.metrics["inference_seconds"])
|
| 440 |
+
print(result.metrics["inference_pages_per_second"])
|
| 441 |
+
```
|
| 442 |
+
|
| 443 |
+
For double pass, use `mode="double", concurrency=256`. The measured fastest
|
| 444 |
+
single-pass profile uses `mode="single", concurrency=256, schedule="staged"`;
|
| 445 |
+
a single-pass pipelined run at concurrency 128 measured 14.92 prediction pages/s.
|
| 446 |
+
A pipelined single-pass run at concurrency 256 has not been measured here.
|
docs/serving.md
ADDED
|
@@ -0,0 +1,106 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# Nemotron 3.5 Parse: Transformers and vLLM
|
| 2 |
+
|
| 3 |
+
## Installation
|
| 4 |
+
|
| 5 |
+
Use Python 3.12 and a fresh environment.
|
| 6 |
+
|
| 7 |
+
The reference serving contract is vLLM 0.26.0, Transformers 5.14.1,
|
| 8 |
+
huggingface-hub 1.24.0, and BF16 on one NVIDIA H100. The package declares the exact
|
| 9 |
+
core dependencies; its `[vllm,client]` extras add serving and the HTTP evaluation
|
| 10 |
+
client. See `pyproject.toml` for pinned direct dependencies.
|
| 11 |
+
|
| 12 |
+
Download the repository and install its serving extras:
|
| 13 |
+
|
| 14 |
+
```bash
|
| 15 |
+
python3.12 -m venv .venv
|
| 16 |
+
source .venv/bin/activate
|
| 17 |
+
python -m pip install "huggingface-hub==1.24.0"
|
| 18 |
+
hf download nvidia/Nemotron-3.5-Parse --local-dir nemotron-parse
|
| 19 |
+
cd nemotron-parse
|
| 20 |
+
export MODEL_REPO="$PWD"
|
| 21 |
+
python -m pip install ".[vllm,client]"
|
| 22 |
+
python -m pip check
|
| 23 |
+
```
|
| 24 |
+
|
| 25 |
+
One-token MTP and the compatible FA3 scheduler are enabled automatically for this model.
|
| 26 |
+
|
| 27 |
+
For Transformers-only inference, no Nemotron package is required. Follow the
|
| 28 |
+
[Quick Start](../README.md#quick-start) to load the model directly from the Hub.
|
| 29 |
+
|
| 30 |
+
## Generation contract
|
| 31 |
+
|
| 32 |
+
The recommended prompt is identical in Transformers, vLLM, and evaluation:
|
| 33 |
+
|
| 34 |
+
`</s><s><predict_bbox><predict_classes><output_markdown><predict_text_in_pic>`
|
| 35 |
+
|
| 36 |
+
Use one RGB page per request, temperature 0, top-k 1, repetition penalty 1.1,
|
| 37 |
+
`max_tokens=8744`, and `skip_special_tokens=False`. Preserve the coordinate and
|
| 38 |
+
class tokens in the raw prediction used by postprocessing and scoring.
|
| 39 |
+
|
| 40 |
+
Boxes are normalized to the evaluated image content. Scale x by that image's
|
| 41 |
+
width and y by its height; no legacy center-pad inverse is needed. The image
|
| 42 |
+
processor accepts dynamically sized pages and applies the exported NaViT
|
| 43 |
+
preprocessing configuration.
|
| 44 |
+
|
| 45 |
+
## One-token MTP
|
| 46 |
+
|
| 47 |
+
`model.safetensors` contains the main model and all four trained MTP tensors.
|
| 48 |
+
No auxiliary weight sidecar is required. The model has **630,800,496 unique
|
| 49 |
+
trainable parameters**, including **2,099,200 MTP parameters**. The decoder
|
| 50 |
+
embedding and LM head share weights and are counted once. Eight stored scalar
|
| 51 |
+
buffers are excluded from the parameter count.
|
| 52 |
+
|
| 53 |
+
Transformers generation uses the normal decoder. vLLM enables the trained
|
| 54 |
+
one-token proposer automatically; it shares the decoder embedding and output
|
| 55 |
+
head.
|
| 56 |
+
|
| 57 |
+
```bash
|
| 58 |
+
vllm serve "$MODEL_REPO" \
|
| 59 |
+
--served-model-name nemotron-parse \
|
| 60 |
+
--trust-remote-code \
|
| 61 |
+
--dtype bfloat16 \
|
| 62 |
+
--distributed-executor-backend uni \
|
| 63 |
+
--tensor-parallel-size 1 \
|
| 64 |
+
--data-parallel-size 1 \
|
| 65 |
+
--max-model-len 16384 \
|
| 66 |
+
--max-num-seqs 192 \
|
| 67 |
+
--max-num-batched-tokens 16384 \
|
| 68 |
+
--no-enable-prefix-caching \
|
| 69 |
+
--limit-mm-per-prompt '{"image":1}' \
|
| 70 |
+
--attention-config '{"flash_attn_version":3}' \
|
| 71 |
+
--chat-template "$MODEL_REPO/chat_template.jinja" \
|
| 72 |
+
--host 127.0.0.1 \
|
| 73 |
+
--port 8000
|
| 74 |
+
```
|
| 75 |
+
|
| 76 |
+
This is the reference H100 throughput profile. Reduce `--max-num-seqs` when
|
| 77 |
+
using a GPU with less memory. Set `NP_VLLM_MTP=0` before startup to run the matched normal-decoding path. Clients use `model="nemotron-parse"`, matching
|
| 78 |
+
the served model name. Startup must show MTP with one speculative token; vLLM's
|
| 79 |
+
`SpecDecoding metrics` report its measured draft acceptance during generation.
|
| 80 |
+
|
| 81 |
+
## Verification and references
|
| 82 |
+
|
| 83 |
+
The current prepared-batch inference and official scoring records are in the
|
| 84 |
+
[validation receipts](../reference/validation/). The
|
| 85 |
+
[reference reproduction records](../reference/README.md) retain their pinned
|
| 86 |
+
model, runtime and page-set contracts. See
|
| 87 |
+
[runtime metadata](../reference/model/runtime_compatibility.json) and
|
| 88 |
+
[package evidence](../reference/package/README.md) for artifact identities and
|
| 89 |
+
verification scope.
|
| 90 |
+
|
| 91 |
+
Set `NP_VLLM_MTP=0` before starting vLLM to use normal decoding.
|
| 92 |
+
|
| 93 |
+
|
| 94 |
+
### Docker alternative
|
| 95 |
+
|
| 96 |
+
From the downloaded folder, run the included Compose recipe:
|
| 97 |
+
|
| 98 |
+
```bash
|
| 99 |
+
docker compose up --build
|
| 100 |
+
```
|
| 101 |
+
|
| 102 |
+
Use the OpenAI-compatible client in the [Quick Start](../README.md#quick-start). Container verification details are in [reference/README.md](https://huggingface.co/nvidia/Nemotron-3.5-Parse/blob/main/reference/README.md).
|
| 103 |
+
|
| 104 |
+
### HTTP client
|
| 105 |
+
|
| 106 |
+
See [vllm_example.py](../vllm_example.py) for an OpenAI-compatible HTTP client.
|
eval/router_ablation.png
ADDED
|
eval/sample_chart.png
ADDED
|
eval/sample_table.png
ADDED
|
eval/synthetic_eval.png
ADDED
|
explainability.md
ADDED
|
@@ -0,0 +1,15 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
## Explainability
|
| 2 |
+
|
| 3 |
+
Field | Response
|
| 4 |
+
:---- | :----
|
| 5 |
+
Intended Task/Domain: | Document understanding, OCR, layout parsing, table extraction, and text extraction with spatial grounding.
|
| 6 |
+
Model Type: | Transformer-based vision-encoder-decoder model.
|
| 7 |
+
Intended Users: | Developers and teams building document intelligence, retrieval-augmented generation (RAG), data-curation, extraction, and multimodal agent workflows.
|
| 8 |
+
Output: | Structured text containing document element classes, bounding boxes, reading-order information, chart/table annotations, and extracted text.
|
| 9 |
+
Describe how the model works: | Nemotron 3.5 Parse encodes a dynamically sized document image with C-RADIOv4-SO400M, compresses the visual features with a neck, and places them before the OCR prompt as a causal prefix to a 10-layer mBART decoder. The decoder uses grouped-query self-attention and produces structured document annotations with content-normalized boxes. The 72,505-token tokenizer includes multilingual, task-control, and chart-aware tokens such as `<class_Chart>`.
|
| 10 |
+
Name the adversely impacted groups this has been tested to deliver comparable outcomes regardless of: | Not Applicable
|
| 11 |
+
Technical Limitations & Mitigation: | Performance can vary for low-resolution scans, motion blur, handwriting, decorative fonts, dense tables, unusual layouts, unsupported languages or scripts, and documents with heavy compression artifacts. Users should evaluate outputs on representative documents and use human review for high-impact workflows.
|
| 12 |
+
Verified to have met prescribed NVIDIA quality standards: | Yes
|
| 13 |
+
Performance Metrics: | OCR accuracy, layout/class accuracy, table extraction quality, chart extraction quality, reading-order quality, grounding quality, latency, throughput, and qualitative visual inspection.
|
| 14 |
+
Potential Known Risks: | The model may miss text, hallucinate structure, assign incorrect classes, produce inaccurate bounding boxes, or incorrectly order elements. Downstream systems should treat output as model-generated extraction results and apply validation, confidence checks, and human review where appropriate.
|
| 15 |
+
Licensing: | Governing Terms: Your use of this model is governed by the [OpenMDW License Agreement, version 1.1](https://openmdw.ai/license/1-1/). Use of the tokenizer included in this model is governed by the [MIT license and tokenizer attribution](https://huggingface.co/nvidia/Nemotron-3.5-Parse/blob/main/TOKENIZER_ATTRIBUTION.md). See [third-party notices](https://huggingface.co/nvidia/Nemotron-3.5-Parse/blob/main/THIRD_PARTY_NOTICES.md) for separately installed dependencies.
|
generation_config.json
ADDED
|
@@ -0,0 +1,13 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"_from_model_config": true,
|
| 3 |
+
"bos_token_id": 0,
|
| 4 |
+
"decoder_start_token_id": 2,
|
| 5 |
+
"eos_token_id": 2,
|
| 6 |
+
"forced_eos_token_id": 2,
|
| 7 |
+
"pad_token_id": 1,
|
| 8 |
+
"transformers_version": "5.6.1",
|
| 9 |
+
"max_new_tokens": 9000,
|
| 10 |
+
"do_sample": false,
|
| 11 |
+
"num_beams": 1,
|
| 12 |
+
"repetition_penalty": 1.1
|
| 13 |
+
}
|
hf_logits_processor.py
ADDED
|
@@ -0,0 +1,352 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
| 2 |
+
# SPDX-License-Identifier: OpenMDW-1.1
|
| 3 |
+
|
| 4 |
+
"""
|
| 5 |
+
LogitsProcessors for document parsing models.
|
| 6 |
+
|
| 7 |
+
Contains:
|
| 8 |
+
- TableInsertionLogitsProcessor: Forces table structure insertion after coordinates
|
| 9 |
+
- RepetitionStopProcessor: Detects repetition/hallucination and forces coordinate tokens
|
| 10 |
+
"""
|
| 11 |
+
|
| 12 |
+
import torch
|
| 13 |
+
from collections import Counter
|
| 14 |
+
from transformers import LogitsProcessor
|
| 15 |
+
from typing import Set, List, Tuple
|
| 16 |
+
|
| 17 |
+
|
| 18 |
+
class TableInsertionLogitsProcessor(LogitsProcessor):
|
| 19 |
+
"""
|
| 20 |
+
A LogitsProcessor that inserts \begin{tabular} tokens after coordinate pairs
|
| 21 |
+
that mark the START of an object (not the END coordinates).
|
| 22 |
+
|
| 23 |
+
Pattern: <x_start><y_start>CONTENT<x_end><y_end><class_...>
|
| 24 |
+
|
| 25 |
+
This processor triggers after the first <x_...><y_...> pair (start coords),
|
| 26 |
+
but not after the second pair (end coords).
|
| 27 |
+
|
| 28 |
+
Args:
|
| 29 |
+
tokenizer: The tokenizer used for encoding/decoding
|
| 30 |
+
table_prefix: The string to insert (default: "\\begin{tabular}")
|
| 31 |
+
"""
|
| 32 |
+
|
| 33 |
+
def __init__(
|
| 34 |
+
self,
|
| 35 |
+
tokenizer,
|
| 36 |
+
table_prefix: str = "\\begin{tabular}"
|
| 37 |
+
):
|
| 38 |
+
self.tokenizer = tokenizer
|
| 39 |
+
self.table_prefix = table_prefix
|
| 40 |
+
|
| 41 |
+
# Tokenize the table prefix to get the sequence of tokens to force
|
| 42 |
+
self.table_prefix_ids = tokenizer.encode(
|
| 43 |
+
table_prefix,
|
| 44 |
+
add_special_tokens=False
|
| 45 |
+
)
|
| 46 |
+
|
| 47 |
+
# Build sets of token IDs for detection
|
| 48 |
+
self._build_token_sets()
|
| 49 |
+
|
| 50 |
+
# State tracking for forced insertion
|
| 51 |
+
self._insertion_position = {} # batch_idx -> position in table_prefix_ids
|
| 52 |
+
self._insertion_active = {} # batch_idx -> bool
|
| 53 |
+
|
| 54 |
+
# State tracking for coordinate pairs
|
| 55 |
+
# False = expecting START coordinates (should trigger)
|
| 56 |
+
# True = expecting END coordinates (should NOT trigger)
|
| 57 |
+
self._expecting_end_coords = {} # batch_idx -> bool
|
| 58 |
+
|
| 59 |
+
def _build_token_sets(self):
|
| 60 |
+
"""Build sets of token IDs for x_, y_, class_, and special tokens."""
|
| 61 |
+
vocab = self.tokenizer.get_vocab()
|
| 62 |
+
|
| 63 |
+
# Collect all x_ coordinate tokens
|
| 64 |
+
self.x_token_ids: Set[int] = set()
|
| 65 |
+
for token, token_id in vocab.items():
|
| 66 |
+
if token.startswith("<x_") and token.endswith(">"):
|
| 67 |
+
self.x_token_ids.add(token_id)
|
| 68 |
+
|
| 69 |
+
# Collect all y_ coordinate tokens
|
| 70 |
+
self.y_token_ids: Set[int] = set()
|
| 71 |
+
for token, token_id in vocab.items():
|
| 72 |
+
if token.startswith("<y_") and token.endswith(">"):
|
| 73 |
+
self.y_token_ids.add(token_id)
|
| 74 |
+
|
| 75 |
+
# Collect all class tokens
|
| 76 |
+
self.class_token_ids: Set[int] = set()
|
| 77 |
+
for token, token_id in vocab.items():
|
| 78 |
+
if token.startswith("<class_") and token.endswith(">"):
|
| 79 |
+
self.class_token_ids.add(token_id)
|
| 80 |
+
|
| 81 |
+
def _is_xy_pair(self, input_ids: torch.Tensor, batch_idx: int) -> bool:
|
| 82 |
+
"""Check if the sequence ends with <x_...><y_...>."""
|
| 83 |
+
seq = input_ids[batch_idx].tolist()
|
| 84 |
+
|
| 85 |
+
if len(seq) < 2:
|
| 86 |
+
return False
|
| 87 |
+
|
| 88 |
+
return seq[-1] in self.y_token_ids and seq[-2] in self.x_token_ids
|
| 89 |
+
|
| 90 |
+
def _last_token_is_class(self, input_ids: torch.Tensor, batch_idx: int) -> bool:
|
| 91 |
+
"""Check if the last token is a <class_...> token."""
|
| 92 |
+
seq = input_ids[batch_idx].tolist()
|
| 93 |
+
|
| 94 |
+
if len(seq) < 1:
|
| 95 |
+
return False
|
| 96 |
+
|
| 97 |
+
return seq[-1] in self.class_token_ids
|
| 98 |
+
|
| 99 |
+
def reset(self):
|
| 100 |
+
"""Reset the processor state for a new generation."""
|
| 101 |
+
self._insertion_position = {}
|
| 102 |
+
self._insertion_active = {}
|
| 103 |
+
self._expecting_end_coords = {}
|
| 104 |
+
|
| 105 |
+
def __call__(
|
| 106 |
+
self,
|
| 107 |
+
input_ids: torch.LongTensor,
|
| 108 |
+
scores: torch.FloatTensor
|
| 109 |
+
) -> torch.FloatTensor:
|
| 110 |
+
"""
|
| 111 |
+
Process logits to force table prefix insertion when appropriate.
|
| 112 |
+
|
| 113 |
+
Args:
|
| 114 |
+
input_ids: (batch_size, seq_len) - tokens generated so far
|
| 115 |
+
scores: (batch_size, vocab_size) - logits for next token
|
| 116 |
+
|
| 117 |
+
Returns:
|
| 118 |
+
Modified scores with forced tokens where appropriate
|
| 119 |
+
"""
|
| 120 |
+
batch_size = input_ids.shape[0]
|
| 121 |
+
|
| 122 |
+
for batch_idx in range(batch_size):
|
| 123 |
+
# Check if we're currently in an insertion sequence
|
| 124 |
+
if self._insertion_active.get(batch_idx, False):
|
| 125 |
+
pos = self._insertion_position[batch_idx]
|
| 126 |
+
|
| 127 |
+
if pos < len(self.table_prefix_ids):
|
| 128 |
+
# Force the next token in the sequence
|
| 129 |
+
forced_token_id = self.table_prefix_ids[pos]
|
| 130 |
+
scores[batch_idx] = torch.full_like(scores[batch_idx], float('-inf'))
|
| 131 |
+
scores[batch_idx, forced_token_id] = 0.0
|
| 132 |
+
|
| 133 |
+
# Advance position
|
| 134 |
+
self._insertion_position[batch_idx] = pos + 1
|
| 135 |
+
else:
|
| 136 |
+
# Finished inserting, deactivate
|
| 137 |
+
self._insertion_active[batch_idx] = False
|
| 138 |
+
continue
|
| 139 |
+
|
| 140 |
+
# Check if we just saw a <class_...> token - reset state to expect START coords
|
| 141 |
+
if self._last_token_is_class(input_ids, batch_idx):
|
| 142 |
+
self._expecting_end_coords[batch_idx] = False
|
| 143 |
+
|
| 144 |
+
# Check if we have an <x_...><y_...> pair
|
| 145 |
+
if self._is_xy_pair(input_ids, batch_idx):
|
| 146 |
+
expecting_end = self._expecting_end_coords.get(batch_idx, False)
|
| 147 |
+
|
| 148 |
+
if not expecting_end:
|
| 149 |
+
# This is START coordinates - trigger insertion!
|
| 150 |
+
self._insertion_active[batch_idx] = True
|
| 151 |
+
self._insertion_position[batch_idx] = 0
|
| 152 |
+
|
| 153 |
+
# After START coords, we expect END coords next
|
| 154 |
+
self._expecting_end_coords[batch_idx] = True
|
| 155 |
+
|
| 156 |
+
# Force the first token
|
| 157 |
+
if len(self.table_prefix_ids) > 0:
|
| 158 |
+
forced_token_id = self.table_prefix_ids[0]
|
| 159 |
+
scores[batch_idx] = torch.full_like(scores[batch_idx], float('-inf'))
|
| 160 |
+
scores[batch_idx, forced_token_id] = 0.0
|
| 161 |
+
self._insertion_position[batch_idx] = 1
|
| 162 |
+
else:
|
| 163 |
+
# This is END coordinates - don't trigger, but reset for next object
|
| 164 |
+
# The <class_...> token will come next, which resets expecting_end_coords
|
| 165 |
+
pass
|
| 166 |
+
|
| 167 |
+
return scores
|
| 168 |
+
|
| 169 |
+
|
| 170 |
+
class RepetitionStopProcessor(LogitsProcessor):
|
| 171 |
+
"""
|
| 172 |
+
A LogitsProcessor that detects CONSECUTIVE repetition/hallucination in generated
|
| 173 |
+
content and forces the model to output <x_...> tokens to close the current object.
|
| 174 |
+
|
| 175 |
+
This detects patterns like:
|
| 176 |
+
- <x_0.0><x_0.0><x_0.0>... (same token repeating consecutively)
|
| 177 |
+
- "ABC ABC ABC ABC..." (same phrase repeating consecutively)
|
| 178 |
+
|
| 179 |
+
But NOT patterns like:
|
| 180 |
+
- "& 1 & 2 & 3 & 4..." (delimiters with different content between them)
|
| 181 |
+
|
| 182 |
+
When consecutive repetition exceeds the threshold, the processor forces an <x_...>
|
| 183 |
+
token to start coordinates, effectively ending the current content. After triggering,
|
| 184 |
+
it enters a cooldown until a <class_...> token is seen (indicating the object was
|
| 185 |
+
closed properly).
|
| 186 |
+
|
| 187 |
+
Args:
|
| 188 |
+
tokenizer: The tokenizer used for encoding/decoding
|
| 189 |
+
max_repetitions: Max consecutive repetitions before forcing stop (default: 10)
|
| 190 |
+
ngram_sizes: List of n-gram sizes to check for repetition (default: [1, 2, 3, 4, 5])
|
| 191 |
+
window_size: Size of the sliding window to check for repetitions (default: 2 * max_ngram * (max_repetitions + 1))
|
| 192 |
+
"""
|
| 193 |
+
|
| 194 |
+
def __init__(
|
| 195 |
+
self,
|
| 196 |
+
tokenizer,
|
| 197 |
+
max_repetitions: int = 10,
|
| 198 |
+
ngram_sizes: List[int] = None,
|
| 199 |
+
window_size: int | None = None
|
| 200 |
+
):
|
| 201 |
+
self.tokenizer = tokenizer
|
| 202 |
+
self.max_repetitions = max_repetitions
|
| 203 |
+
# Check various n-gram sizes for consecutive repetition
|
| 204 |
+
# n=1 catches single token repetition like <x_0.0><x_0.0><x_0.0>...
|
| 205 |
+
# n=2,3,4,5 catch phrase repetition like "ABC ABC ABC..."
|
| 206 |
+
self.ngram_sizes = ngram_sizes if ngram_sizes is not None else [1, 2, 3, 4, 5]
|
| 207 |
+
max_ngram = max((n for n in self.ngram_sizes if n > 0), default=1)
|
| 208 |
+
default_window = 2 * max_ngram * (self.max_repetitions + 1)
|
| 209 |
+
self.window_size = int(window_size) if window_size is not None else int(default_window)
|
| 210 |
+
|
| 211 |
+
# Build set of token IDs
|
| 212 |
+
self._build_token_sets()
|
| 213 |
+
|
| 214 |
+
# State tracking - once triggered, stay in cooldown until we see <class_...>
|
| 215 |
+
self._in_cooldown = {} # batch_idx -> bool
|
| 216 |
+
# Track the start of the "current object segment" (tokens after the last <class_...>)
|
| 217 |
+
# We only look for repetition inside this segment to avoid old repetition re-triggering immediately.
|
| 218 |
+
self._segment_start = {} # batch_idx -> int (seq index)
|
| 219 |
+
|
| 220 |
+
def _build_token_sets(self):
|
| 221 |
+
"""Build sets of token IDs for detection."""
|
| 222 |
+
vocab = self.tokenizer.get_vocab()
|
| 223 |
+
|
| 224 |
+
# Collect all x_ coordinate tokens
|
| 225 |
+
self.x_token_ids: Set[int] = set()
|
| 226 |
+
self.x_token_list: List[int] = []
|
| 227 |
+
for token, token_id in vocab.items():
|
| 228 |
+
if token.startswith("<x_") and token.endswith(">"):
|
| 229 |
+
self.x_token_ids.add(token_id)
|
| 230 |
+
self.x_token_list.append(token_id)
|
| 231 |
+
|
| 232 |
+
# Sort and pick a middle x token as default
|
| 233 |
+
self.x_token_list.sort()
|
| 234 |
+
self.default_x_token = self.x_token_list[len(self.x_token_list) // 2] if self.x_token_list else None
|
| 235 |
+
|
| 236 |
+
# Collect all class tokens (used to reset cooldown)
|
| 237 |
+
self.class_token_ids: Set[int] = set()
|
| 238 |
+
for token, token_id in vocab.items():
|
| 239 |
+
if token.startswith("<class_") and token.endswith(">"):
|
| 240 |
+
self.class_token_ids.add(token_id)
|
| 241 |
+
|
| 242 |
+
def _count_consecutive_repetitions(self, seq: List[int], n: int) -> int:
|
| 243 |
+
"""
|
| 244 |
+
Count the maximum number of times an n-gram repeats CONSECUTIVELY.
|
| 245 |
+
|
| 246 |
+
For example, if seq contains "ABC ABC ABC ABC DEF ABC":
|
| 247 |
+
- The n-gram "ABC" appears 4 times consecutively at the start
|
| 248 |
+
- Returns 4 (the max consecutive count)
|
| 249 |
+
|
| 250 |
+
Returns:
|
| 251 |
+
Maximum consecutive repetition count for any n-gram
|
| 252 |
+
"""
|
| 253 |
+
if len(seq) < n:
|
| 254 |
+
return 0
|
| 255 |
+
|
| 256 |
+
max_consecutive = 1
|
| 257 |
+
current_consecutive = 1
|
| 258 |
+
|
| 259 |
+
# Slide through the sequence comparing adjacent n-grams
|
| 260 |
+
prev_ngram = tuple(seq[0:n])
|
| 261 |
+
|
| 262 |
+
i = n
|
| 263 |
+
while i <= len(seq) - n:
|
| 264 |
+
current_ngram = tuple(seq[i:i + n])
|
| 265 |
+
|
| 266 |
+
if current_ngram == prev_ngram:
|
| 267 |
+
current_consecutive += 1
|
| 268 |
+
max_consecutive = max(max_consecutive, current_consecutive)
|
| 269 |
+
i += n # Skip by n to check the next occurrence
|
| 270 |
+
else:
|
| 271 |
+
current_consecutive = 1
|
| 272 |
+
prev_ngram = current_ngram
|
| 273 |
+
i += 1 # Move by 1 to find new patterns
|
| 274 |
+
|
| 275 |
+
return max_consecutive
|
| 276 |
+
|
| 277 |
+
def _has_excessive_repetition(self, seq: List[int]) -> bool:
|
| 278 |
+
"""
|
| 279 |
+
Check if the sequence has excessive CONSECUTIVE repetition of any n-gram.
|
| 280 |
+
|
| 281 |
+
This detects hallucination patterns like:
|
| 282 |
+
- <x_0.0><x_0.0><x_0.0><x_0.0>... (same token repeating)
|
| 283 |
+
- "ABC ABC ABC ABC..." (same phrase repeating)
|
| 284 |
+
|
| 285 |
+
But NOT patterns like:
|
| 286 |
+
- "& 1 & 2 & 3 & 4..." (delimiters with different content)
|
| 287 |
+
|
| 288 |
+
Returns:
|
| 289 |
+
True if any n-gram repeats consecutively more than max_repetitions times
|
| 290 |
+
"""
|
| 291 |
+
# Only check the recent window
|
| 292 |
+
check_seq = seq[-self.window_size:] if len(seq) > self.window_size else seq
|
| 293 |
+
|
| 294 |
+
for n in self.ngram_sizes:
|
| 295 |
+
consecutive_count = self._count_consecutive_repetitions(check_seq, n)
|
| 296 |
+
if consecutive_count > self.max_repetitions:
|
| 297 |
+
return True
|
| 298 |
+
|
| 299 |
+
return False
|
| 300 |
+
|
| 301 |
+
def reset(self):
|
| 302 |
+
"""Reset the processor state for a new generation."""
|
| 303 |
+
self._in_cooldown = {}
|
| 304 |
+
self._segment_start = {}
|
| 305 |
+
|
| 306 |
+
def __call__(
|
| 307 |
+
self,
|
| 308 |
+
input_ids: torch.LongTensor,
|
| 309 |
+
scores: torch.FloatTensor
|
| 310 |
+
) -> torch.FloatTensor:
|
| 311 |
+
"""
|
| 312 |
+
Process logits to force <x_...> token when repetition is detected.
|
| 313 |
+
|
| 314 |
+
Args:
|
| 315 |
+
input_ids: (batch_size, seq_len) - tokens generated so far
|
| 316 |
+
scores: (batch_size, vocab_size) - logits for next token
|
| 317 |
+
|
| 318 |
+
Returns:
|
| 319 |
+
Modified scores with forced <x_...> tokens where repetition detected
|
| 320 |
+
"""
|
| 321 |
+
batch_size = input_ids.shape[0]
|
| 322 |
+
|
| 323 |
+
for batch_idx in range(batch_size):
|
| 324 |
+
seq = input_ids[batch_idx].tolist()
|
| 325 |
+
seq_len = len(seq)
|
| 326 |
+
|
| 327 |
+
# Check if we just saw a class token - reset cooldown and start a new segment
|
| 328 |
+
if seq_len > 0 and seq[-1] in self.class_token_ids:
|
| 329 |
+
self._in_cooldown[batch_idx] = False
|
| 330 |
+
self._segment_start[batch_idx] = seq_len
|
| 331 |
+
|
| 332 |
+
# If in cooldown, don't check for repetition (let model generate naturally)
|
| 333 |
+
if self._in_cooldown.get(batch_idx, False):
|
| 334 |
+
continue
|
| 335 |
+
|
| 336 |
+
# Check if repetition threshold is exceeded
|
| 337 |
+
segment_start = self._segment_start.get(batch_idx, 0)
|
| 338 |
+
segment_seq = seq[segment_start:]
|
| 339 |
+
if self._has_excessive_repetition(segment_seq):
|
| 340 |
+
# Enter cooldown to avoid repeated triggering
|
| 341 |
+
self._in_cooldown[batch_idx] = True
|
| 342 |
+
|
| 343 |
+
# Force one of the <x_...> tokens to start closing the object
|
| 344 |
+
# Keep the model's original preferences for which x_ coordinate to use
|
| 345 |
+
original_scores = scores[batch_idx].clone()
|
| 346 |
+
scores[batch_idx] = torch.full_like(scores[batch_idx], float('-inf'))
|
| 347 |
+
|
| 348 |
+
# Restore original logits for x_ tokens only
|
| 349 |
+
for x_token_id in self.x_token_ids:
|
| 350 |
+
scores[batch_idx, x_token_id] = original_scores[x_token_id]
|
| 351 |
+
|
| 352 |
+
return scores
|
hf_nemotron_parse_config.py
ADDED
|
@@ -0,0 +1,133 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
| 2 |
+
# SPDX-License-Identifier: OpenMDW-1.1
|
| 3 |
+
|
| 4 |
+
from transformers import PretrainedConfig
|
| 5 |
+
from typing import List, Optional
|
| 6 |
+
|
| 7 |
+
|
| 8 |
+
class NemotronParseTextConfig(PretrainedConfig):
|
| 9 |
+
"""
|
| 10 |
+
Configuration class for NemotronParse text decoder (mBART-based).
|
| 11 |
+
"""
|
| 12 |
+
model_type = "nemotron_parse_text"
|
| 13 |
+
|
| 14 |
+
def __init__(
|
| 15 |
+
self,
|
| 16 |
+
vocab_size: int = 250027,
|
| 17 |
+
d_model: int = 1024,
|
| 18 |
+
encoder_layers: int = 12,
|
| 19 |
+
decoder_layers: int = 12,
|
| 20 |
+
encoder_attention_heads: int = 16,
|
| 21 |
+
decoder_attention_heads: int = 16,
|
| 22 |
+
decoder_ffn_dim: int = 4096,
|
| 23 |
+
encoder_ffn_dim: int = 4096,
|
| 24 |
+
activation_function: str = "gelu",
|
| 25 |
+
dropout: float = 0.1,
|
| 26 |
+
attention_dropout: float = 0.0,
|
| 27 |
+
activation_dropout: float = 0.0,
|
| 28 |
+
classifier_dropout: float = 0.0,
|
| 29 |
+
init_std: float = 0.02,
|
| 30 |
+
encoder_layerdrop: float = 0.0,
|
| 31 |
+
decoder_layerdrop: float = 0.0,
|
| 32 |
+
scale_embedding: bool = False,
|
| 33 |
+
use_cache: bool = True,
|
| 34 |
+
num_labels: int = 3,
|
| 35 |
+
forced_eos_token_id: int = 2,
|
| 36 |
+
add_cross_attention: bool = True, # Enable cross-attention for vision-encoder-decoder
|
| 37 |
+
is_decoder: bool = True, # This is a decoder
|
| 38 |
+
max_sequence_length: int = 9000,
|
| 39 |
+
**kwargs
|
| 40 |
+
):
|
| 41 |
+
super().__init__(**kwargs)
|
| 42 |
+
self.vocab_size = vocab_size
|
| 43 |
+
self.d_model = d_model
|
| 44 |
+
self.encoder_layers = encoder_layers
|
| 45 |
+
self.decoder_layers = decoder_layers
|
| 46 |
+
self.encoder_attention_heads = encoder_attention_heads
|
| 47 |
+
self.decoder_attention_heads = decoder_attention_heads
|
| 48 |
+
self.decoder_ffn_dim = decoder_ffn_dim
|
| 49 |
+
self.encoder_ffn_dim = encoder_ffn_dim
|
| 50 |
+
self.activation_function = activation_function
|
| 51 |
+
self.dropout = dropout
|
| 52 |
+
self.attention_dropout = attention_dropout
|
| 53 |
+
self.activation_dropout = activation_dropout
|
| 54 |
+
self.classifier_dropout = classifier_dropout
|
| 55 |
+
self.init_std = init_std
|
| 56 |
+
self.encoder_layerdrop = encoder_layerdrop
|
| 57 |
+
self.decoder_layerdrop = decoder_layerdrop
|
| 58 |
+
self.scale_embedding = scale_embedding
|
| 59 |
+
self.use_cache = use_cache
|
| 60 |
+
self.num_labels = num_labels
|
| 61 |
+
self.add_cross_attention = add_cross_attention
|
| 62 |
+
self.is_decoder = is_decoder
|
| 63 |
+
|
| 64 |
+
# Add hidden_size as alias for d_model (for compatibility)
|
| 65 |
+
self.hidden_size = self.d_model
|
| 66 |
+
self.forced_eos_token_id = forced_eos_token_id
|
| 67 |
+
self.num_attention_heads = self.encoder_attention_heads
|
| 68 |
+
|
| 69 |
+
self.max_sequence_length = max_sequence_length
|
| 70 |
+
|
| 71 |
+
|
| 72 |
+
class NemotronParseConfig(PretrainedConfig):
|
| 73 |
+
"""
|
| 74 |
+
Configuration class for NemotronParse model.
|
| 75 |
+
|
| 76 |
+
This configuration class is used to store the configuration of a [`NemotronParseForConditionalGeneration`] model.
|
| 77 |
+
It is used to instantiate an NemotronParse model according to the specified arguments, defining the vision and text model configs.
|
| 78 |
+
"""
|
| 79 |
+
model_type = "nemotron_parse"
|
| 80 |
+
is_composition = True
|
| 81 |
+
max_sequence_length = 9000
|
| 82 |
+
|
| 83 |
+
def __init__(
|
| 84 |
+
self,
|
| 85 |
+
encoder: Optional[dict] = None,
|
| 86 |
+
decoder: Optional[dict] = None,
|
| 87 |
+
tie_word_embeddings: bool = False,
|
| 88 |
+
decoder_start_token_id: int = 2,
|
| 89 |
+
pad_token_id: int = 1,
|
| 90 |
+
eos_token_id: int = 2,
|
| 91 |
+
bos_token_id: int = 0,
|
| 92 |
+
image_size: List[int] = [2048, 1664],
|
| 93 |
+
is_encoder_decoder: bool = True,
|
| 94 |
+
max_sequence_length: int = 9000,
|
| 95 |
+
**kwargs
|
| 96 |
+
):
|
| 97 |
+
super().__init__(
|
| 98 |
+
tie_word_embeddings=tie_word_embeddings,
|
| 99 |
+
decoder_start_token_id=decoder_start_token_id,
|
| 100 |
+
pad_token_id=pad_token_id,
|
| 101 |
+
eos_token_id=eos_token_id,
|
| 102 |
+
bos_token_id=bos_token_id,
|
| 103 |
+
max_sequence_length=max_sequence_length,
|
| 104 |
+
**kwargs
|
| 105 |
+
)
|
| 106 |
+
|
| 107 |
+
|
| 108 |
+
if decoder is None:
|
| 109 |
+
decoder = {}
|
| 110 |
+
|
| 111 |
+
# Encoder tensors and architecture are part of this release. Avoid
|
| 112 |
+
# importing a second model repository while constructing the config.
|
| 113 |
+
self.encoder = PretrainedConfig(**(encoder or {}))
|
| 114 |
+
|
| 115 |
+
decoder["max_sequence_length"] = max_sequence_length
|
| 116 |
+
self.decoder = NemotronParseTextConfig(**decoder)
|
| 117 |
+
self.image_size = image_size
|
| 118 |
+
|
| 119 |
+
# Initialize vocab size from text config
|
| 120 |
+
self.vocab_size = self.decoder.vocab_size
|
| 121 |
+
self.is_encoder_decoder = is_encoder_decoder
|
| 122 |
+
self.max_sequence_length = max_sequence_length
|
| 123 |
+
|
| 124 |
+
def to_dict(self):
|
| 125 |
+
"""
|
| 126 |
+
Serializes this instance to a Python dictionary. Override the default [`~PretrainedConfig.to_dict`].
|
| 127 |
+
"""
|
| 128 |
+
output = super().to_dict()
|
| 129 |
+
output["encoder"] = self.encoder.to_dict()
|
| 130 |
+
output["decoder"] = self.decoder.to_dict()
|
| 131 |
+
output["model_type"] = self.model_type
|
| 132 |
+
output["is_encoder_decoder"] = self.is_encoder_decoder
|
| 133 |
+
return output
|
hf_nemotron_parse_modeling.py
ADDED
|
@@ -0,0 +1,275 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
| 2 |
+
# SPDX-License-Identifier: OpenMDW-1.1
|
| 3 |
+
|
| 4 |
+
"""Native Transformers inference for Nemotron 3.5 Parse.
|
| 5 |
+
|
| 6 |
+
The complete C-RADIOv4-SO400M encoder, content-normalized NaViT geometry,
|
| 7 |
+
causal mBART visual prefix, and trained one-step MTP tensors are loaded from
|
| 8 |
+
this checkpoint. No second encoder model or remote encoder code is fetched.
|
| 9 |
+
"""
|
| 10 |
+
|
| 11 |
+
import math
|
| 12 |
+
from functools import partial
|
| 13 |
+
|
| 14 |
+
import torch
|
| 15 |
+
from torch import nn
|
| 16 |
+
from torch.nn import functional as F
|
| 17 |
+
from timm.models.vision_transformer import Block
|
| 18 |
+
from transformers import GenerationMixin, PreTrainedModel
|
| 19 |
+
from transformers.cache_utils import DynamicCache
|
| 20 |
+
from transformers.modeling_outputs import CausalLMOutputWithPast
|
| 21 |
+
|
| 22 |
+
from .hf_nemotron_parse_config import NemotronParseConfig
|
| 23 |
+
from .hf_prefix_geometry import _normalized_2d_position_pages
|
| 24 |
+
|
| 25 |
+
|
| 26 |
+
class _ClassTokens(nn.Module):
|
| 27 |
+
def __init__(self):
|
| 28 |
+
super().__init__()
|
| 29 |
+
self.token = nn.Parameter(torch.empty(10, 1152))
|
| 30 |
+
|
| 31 |
+
|
| 32 |
+
class _PatchGenerator(nn.Module):
|
| 33 |
+
def __init__(self):
|
| 34 |
+
super().__init__()
|
| 35 |
+
self.cls_token = _ClassTokens()
|
| 36 |
+
self.embedder = nn.Linear(768, 1152, bias=False)
|
| 37 |
+
self.pos_embed = nn.Parameter(torch.empty(1, 16384, 1152))
|
| 38 |
+
|
| 39 |
+
|
| 40 |
+
class _VisionTransformer(nn.Module):
|
| 41 |
+
def __init__(self):
|
| 42 |
+
super().__init__()
|
| 43 |
+
self.patch_generator = _PatchGenerator()
|
| 44 |
+
self.blocks = nn.ModuleList([
|
| 45 |
+
Block(1152, 16, mlp_ratio=4304 / 1152, qkv_bias=True,
|
| 46 |
+
norm_layer=partial(nn.LayerNorm, eps=1e-6))
|
| 47 |
+
for _ in range(27)
|
| 48 |
+
])
|
| 49 |
+
|
| 50 |
+
def forward(self, patches, image_size, content_size):
|
| 51 |
+
pg = self.patch_generator
|
| 52 |
+
patches = pg.embedder(patches).unsqueeze(0)
|
| 53 |
+
patches = _normalized_2d_position_pages(
|
| 54 |
+
torch, patches=patches, pos_embed=pg.pos_embed,
|
| 55 |
+
source_rows=128, source_cols=128, image_sizes=[image_size],
|
| 56 |
+
patch_size=16, content_sizes=[content_size], content_normalized=True,
|
| 57 |
+
)
|
| 58 |
+
hidden = torch.cat((pg.cls_token.token.unsqueeze(0), patches), dim=1)
|
| 59 |
+
for block in self.blocks:
|
| 60 |
+
hidden = block(hidden)
|
| 61 |
+
return hidden
|
| 62 |
+
|
| 63 |
+
|
| 64 |
+
class _InputConditioner(nn.Module):
|
| 65 |
+
def __init__(self):
|
| 66 |
+
super().__init__()
|
| 67 |
+
self.register_buffer("norm_mean", torch.zeros(3, 1, 1))
|
| 68 |
+
self.register_buffer("norm_std", torch.ones(3, 1, 1))
|
| 69 |
+
|
| 70 |
+
|
| 71 |
+
class _RadioModel(nn.Module):
|
| 72 |
+
def __init__(self):
|
| 73 |
+
super().__init__()
|
| 74 |
+
self.register_buffer("summary_idxs", torch.tensor([0, 1]))
|
| 75 |
+
self.input_conditioner = _InputConditioner()
|
| 76 |
+
self.model = _VisionTransformer()
|
| 77 |
+
|
| 78 |
+
|
| 79 |
+
class _RadioWrapper(nn.Module):
|
| 80 |
+
def __init__(self):
|
| 81 |
+
super().__init__()
|
| 82 |
+
self.radio_model = _RadioModel()
|
| 83 |
+
|
| 84 |
+
|
| 85 |
+
class RadioWithNeck(nn.Module):
|
| 86 |
+
def __init__(self, config):
|
| 87 |
+
super().__init__()
|
| 88 |
+
self.model_encoder = _RadioWrapper()
|
| 89 |
+
self.conv1 = nn.Conv1d(1152, 1024, 1)
|
| 90 |
+
self.layer_norm1 = nn.LayerNorm(1024, eps=1e-6)
|
| 91 |
+
self.conv2 = nn.Conv2d(1024, 1024, (1, 4), stride=(1, 4), bias=False)
|
| 92 |
+
self.layer_norm2 = nn.LayerNorm(1024, eps=1e-6)
|
| 93 |
+
self.sum_proj = nn.Linear(2304, 1024)
|
| 94 |
+
self.layer_norm3 = nn.LayerNorm(1024, eps=1e-6)
|
| 95 |
+
|
| 96 |
+
def forward(self, pixel_values, image_sizes, content_sizes):
|
| 97 |
+
# The shipped processor has already applied CLIP normalization and
|
| 98 |
+
# flattened each 16x16 patch in channel/row/column order.
|
| 99 |
+
if image_sizes is None or content_sizes is None:
|
| 100 |
+
raise ValueError("Packed inference requires image_sizes and content_sizes from AutoProcessor")
|
| 101 |
+
if pixel_values.ndim == 3 and pixel_values.shape[0] == 1:
|
| 102 |
+
pixel_values = pixel_values[0]
|
| 103 |
+
if pixel_values.ndim != 2 or pixel_values.shape[1] != 768:
|
| 104 |
+
raise ValueError(f"Expected packed raw patches [patches,768], got {tuple(pixel_values.shape)}")
|
| 105 |
+
pixel_values = pixel_values.to(dtype=self.conv1.weight.dtype)
|
| 106 |
+
pages, cursor = [], 0
|
| 107 |
+
radio = self.model_encoder.radio_model
|
| 108 |
+
for size, content in zip(image_sizes.tolist(), content_sizes.tolist()):
|
| 109 |
+
height, width = map(int, size)
|
| 110 |
+
count = (height // 16) * (width // 16)
|
| 111 |
+
hidden = radio.model(pixel_values[cursor:cursor + count], size, content)
|
| 112 |
+
cursor += count
|
| 113 |
+
summary = hidden[:, radio.summary_idxs].flatten(1)
|
| 114 |
+
features = hidden[:, 10:]
|
| 115 |
+
features = self.layer_norm1(self.conv1(features.transpose(1, 2)).transpose(1, 2))
|
| 116 |
+
features = features.transpose(1, 2).reshape(1, 1024, height // 16, width // 16)
|
| 117 |
+
features = self.conv2(features).flatten(2).transpose(1, 2)
|
| 118 |
+
features = self.layer_norm2(features)
|
| 119 |
+
summary = self.layer_norm3(self.sum_proj(summary)).unsqueeze(1)
|
| 120 |
+
pages.append(torch.cat((features, summary), dim=1)[0])
|
| 121 |
+
if cursor != len(pixel_values):
|
| 122 |
+
raise ValueError("Image metadata did not consume all packed patches")
|
| 123 |
+
max_length = max(len(page) for page in pages)
|
| 124 |
+
output = pages[0].new_zeros(len(pages), max_length, 1024)
|
| 125 |
+
mask = torch.zeros(len(pages), max_length, device=output.device, dtype=torch.bool)
|
| 126 |
+
for i, page in enumerate(pages):
|
| 127 |
+
output[i, :len(page)] = page
|
| 128 |
+
mask[i, :len(page)] = True
|
| 129 |
+
return output, mask
|
| 130 |
+
|
| 131 |
+
|
| 132 |
+
class _ScaledEmbedding(nn.Embedding):
|
| 133 |
+
def __init__(self, config):
|
| 134 |
+
super().__init__(config.vocab_size, config.d_model, padding_idx=config.pad_token_id)
|
| 135 |
+
self.embed_scale = math.sqrt(config.d_model) if config.scale_embedding else 1.0
|
| 136 |
+
|
| 137 |
+
def forward(self, ids):
|
| 138 |
+
return super().forward(ids) * self.embed_scale
|
| 139 |
+
|
| 140 |
+
|
| 141 |
+
class _GQAAttention(nn.Module):
|
| 142 |
+
def __init__(self, config, index):
|
| 143 |
+
super().__init__()
|
| 144 |
+
self.index = index
|
| 145 |
+
self.num_heads = config.decoder_attention_heads
|
| 146 |
+
self.num_kv_heads = config.self_attention_kv_heads
|
| 147 |
+
self.head_dim = config.d_model // self.num_heads
|
| 148 |
+
self.q_proj = nn.Linear(config.d_model, config.d_model)
|
| 149 |
+
self.k_proj = nn.Linear(config.d_model, self.num_kv_heads * self.head_dim)
|
| 150 |
+
self.v_proj = nn.Linear(config.d_model, self.num_kv_heads * self.head_dim)
|
| 151 |
+
self.out_proj = nn.Linear(config.d_model, config.d_model)
|
| 152 |
+
|
| 153 |
+
def forward(self, hidden, cache, key_mask):
|
| 154 |
+
batch, length, _ = hidden.shape
|
| 155 |
+
q = self.q_proj(hidden).view(batch, length, self.num_heads, self.head_dim).transpose(1, 2)
|
| 156 |
+
k = self.k_proj(hidden).view(batch, length, self.num_kv_heads, self.head_dim).transpose(1, 2)
|
| 157 |
+
v = self.v_proj(hidden).view(batch, length, self.num_kv_heads, self.head_dim).transpose(1, 2)
|
| 158 |
+
if cache is not None:
|
| 159 |
+
k, v = cache.update(k, v, self.index)
|
| 160 |
+
past = k.shape[-2] - length
|
| 161 |
+
allowed = torch.arange(k.shape[-2], device=hidden.device)[None, :] <= (
|
| 162 |
+
torch.arange(length, device=hidden.device)[:, None] + past
|
| 163 |
+
)
|
| 164 |
+
allowed = allowed[None, None] & key_mask[:, None, None, :]
|
| 165 |
+
repeat = self.num_heads // self.num_kv_heads
|
| 166 |
+
output = F.scaled_dot_product_attention(
|
| 167 |
+
q, k.repeat_interleave(repeat, dim=1), v.repeat_interleave(repeat, dim=1),
|
| 168 |
+
attn_mask=allowed, dropout_p=0.0,
|
| 169 |
+
)
|
| 170 |
+
return self.out_proj(output.transpose(1, 2).reshape(batch, length, -1))
|
| 171 |
+
|
| 172 |
+
|
| 173 |
+
class _DecoderLayer(nn.Module):
|
| 174 |
+
def __init__(self, config, index):
|
| 175 |
+
super().__init__()
|
| 176 |
+
self.self_attn = _GQAAttention(config, index)
|
| 177 |
+
self.self_attn_layer_norm = nn.LayerNorm(config.d_model)
|
| 178 |
+
self.final_layer_norm = nn.LayerNorm(config.d_model)
|
| 179 |
+
self.fc1 = nn.Linear(config.d_model, config.decoder_ffn_dim)
|
| 180 |
+
self.fc2 = nn.Linear(config.decoder_ffn_dim, config.d_model)
|
| 181 |
+
|
| 182 |
+
def forward(self, hidden, cache, mask):
|
| 183 |
+
hidden = hidden + self.self_attn(self.self_attn_layer_norm(hidden), cache, mask)
|
| 184 |
+
return hidden + self.fc2(F.gelu(self.fc1(self.final_layer_norm(hidden))))
|
| 185 |
+
|
| 186 |
+
|
| 187 |
+
class NemotronParseDecoder(nn.Module):
|
| 188 |
+
def __init__(self, config, num_extra_heads):
|
| 189 |
+
super().__init__()
|
| 190 |
+
self.embed_tokens = _ScaledEmbedding(config)
|
| 191 |
+
self.layers = nn.ModuleList([_DecoderLayer(config, i) for i in range(config.decoder_layers)])
|
| 192 |
+
self.layernorm_embedding = nn.LayerNorm(config.d_model)
|
| 193 |
+
self.layer_norm = nn.LayerNorm(config.d_model)
|
| 194 |
+
self.extra_heads = nn.ModuleList([nn.Linear(config.d_model, config.d_model) for _ in range(num_extra_heads)])
|
| 195 |
+
self.extra_proj = nn.ModuleList([nn.Linear(config.d_model, config.d_model) for _ in range(num_extra_heads)])
|
| 196 |
+
|
| 197 |
+
def forward(self, embeddings, cache, mask):
|
| 198 |
+
hidden = self.layernorm_embedding(embeddings)
|
| 199 |
+
for layer in self.layers:
|
| 200 |
+
hidden = layer(hidden, cache, mask)
|
| 201 |
+
return self.layer_norm(hidden)
|
| 202 |
+
|
| 203 |
+
|
| 204 |
+
class NemotronParseForConditionalGeneration(PreTrainedModel, GenerationMixin):
|
| 205 |
+
config_class = NemotronParseConfig
|
| 206 |
+
main_input_name = "input_ids"
|
| 207 |
+
_tied_weights_keys = {"lm_head.weight": "decoder.embed_tokens.weight"}
|
| 208 |
+
_no_split_modules = ["RadioWithNeck", "Block", "_DecoderLayer"]
|
| 209 |
+
_supports_sdpa = True
|
| 210 |
+
|
| 211 |
+
def __init__(self, config):
|
| 212 |
+
super().__init__(config)
|
| 213 |
+
if not getattr(config, "legacy_causal_vision_prefix", False):
|
| 214 |
+
raise ValueError("This release implements the trained causal visual-prefix architecture")
|
| 215 |
+
self.encoder = RadioWithNeck(config.encoder)
|
| 216 |
+
self.decoder = NemotronParseDecoder(config.decoder, getattr(config, "num_extra_heads", 0))
|
| 217 |
+
self.lm_head = nn.Linear(config.decoder.d_model, config.decoder.vocab_size, bias=False)
|
| 218 |
+
self.post_init()
|
| 219 |
+
|
| 220 |
+
def get_input_embeddings(self):
|
| 221 |
+
return self.decoder.embed_tokens
|
| 222 |
+
|
| 223 |
+
def get_output_embeddings(self):
|
| 224 |
+
return self.lm_head
|
| 225 |
+
|
| 226 |
+
def _tie_weights(self):
|
| 227 |
+
self._tie_or_clone_weights(self.lm_head, self.decoder.embed_tokens)
|
| 228 |
+
|
| 229 |
+
def forward(self, input_ids=None, decoder_input_ids=None, pixel_values=None,
|
| 230 |
+
image_sizes=None, content_sizes=None, attention_mask=None,
|
| 231 |
+
past_key_values=None, use_cache=True, return_dict=True, **kwargs):
|
| 232 |
+
del kwargs
|
| 233 |
+
if input_ids is None:
|
| 234 |
+
input_ids = decoder_input_ids
|
| 235 |
+
if input_ids is None:
|
| 236 |
+
raise ValueError("input_ids are required")
|
| 237 |
+
if past_key_values is None and use_cache:
|
| 238 |
+
past_key_values = DynamicCache()
|
| 239 |
+
past_length = past_key_values.get_seq_length() if past_key_values is not None else 0
|
| 240 |
+
text_embeddings = self.decoder.embed_tokens(input_ids)
|
| 241 |
+
text_offset = 0
|
| 242 |
+
if past_length == 0:
|
| 243 |
+
if pixel_values is None:
|
| 244 |
+
raise ValueError("pixel_values are required for the first generation step")
|
| 245 |
+
vision, visual_mask = self.encoder(pixel_values, image_sizes, content_sizes)
|
| 246 |
+
embeddings = torch.cat((vision, text_embeddings), dim=1)
|
| 247 |
+
text_offset = vision.shape[1]
|
| 248 |
+
if past_key_values is not None:
|
| 249 |
+
past_key_values._nemotron_visual_mask = visual_mask
|
| 250 |
+
else:
|
| 251 |
+
embeddings = text_embeddings
|
| 252 |
+
visual_mask = past_key_values._nemotron_visual_mask
|
| 253 |
+
if attention_mask is None:
|
| 254 |
+
text_length = past_length + input_ids.shape[1] - visual_mask.shape[1]
|
| 255 |
+
if past_length == 0:
|
| 256 |
+
text_length = input_ids.shape[1]
|
| 257 |
+
attention_mask = torch.ones(input_ids.shape[0], text_length, device=input_ids.device, dtype=torch.bool)
|
| 258 |
+
mask = torch.cat((visual_mask, attention_mask.bool()), dim=1)
|
| 259 |
+
if mask.shape[1] > self.config.max_sequence_length:
|
| 260 |
+
raise ValueError("Visual prefix plus text exceeds the model context limit")
|
| 261 |
+
hidden = self.decoder(embeddings, past_key_values if use_cache else None, mask)
|
| 262 |
+
logits = self.lm_head(hidden[:, text_offset:])
|
| 263 |
+
output = CausalLMOutputWithPast(logits=logits, past_key_values=past_key_values if use_cache else None)
|
| 264 |
+
return output if return_dict else (output.logits, output.past_key_values)
|
| 265 |
+
|
| 266 |
+
def prepare_inputs_for_generation(self, input_ids, past_key_values=None,
|
| 267 |
+
attention_mask=None, pixel_values=None,
|
| 268 |
+
image_sizes=None, content_sizes=None, **kwargs):
|
| 269 |
+
del kwargs
|
| 270 |
+
if past_key_values is not None and past_key_values.get_seq_length() > 0:
|
| 271 |
+
input_ids = input_ids[:, -1:]
|
| 272 |
+
pixel_values = None
|
| 273 |
+
return {"input_ids": input_ids, "attention_mask": attention_mask,
|
| 274 |
+
"past_key_values": past_key_values, "pixel_values": pixel_values,
|
| 275 |
+
"image_sizes": image_sizes, "content_sizes": content_sizes, "use_cache": True}
|
hf_nemotron_parse_processor.py
ADDED
|
@@ -0,0 +1,501 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
| 2 |
+
# SPDX-License-Identifier: OpenMDW-1.1
|
| 3 |
+
|
| 4 |
+
import numpy as np
|
| 5 |
+
from PIL import Image
|
| 6 |
+
from typing import List, Optional, Union, Dict, Any
|
| 7 |
+
import torch
|
| 8 |
+
from torchvision import transforms as T
|
| 9 |
+
import json
|
| 10 |
+
|
| 11 |
+
from transformers import BatchFeature, ProcessorMixin, BaseImageProcessor, ImageProcessingMixin
|
| 12 |
+
from transformers.image_utils import ChannelDimension, ImageInput, PILImageResampling, infer_channel_dimension_format
|
| 13 |
+
from transformers.utils import TensorType
|
| 14 |
+
|
| 15 |
+
OPENAI_CLIP_MEAN = (0.48145466, 0.4578275, 0.40821073)
|
| 16 |
+
OPENAI_CLIP_STD = (0.26862954, 0.26130258, 0.27577711)
|
| 17 |
+
|
| 18 |
+
|
| 19 |
+
from transformers.processing_utils import ImagesKwargs
|
| 20 |
+
|
| 21 |
+
class NemotronParseImagesKwargs(ImagesKwargs, total=False):
|
| 22 |
+
content_hw: Optional[List[int]]
|
| 23 |
+
|
| 24 |
+
class NemotronParseImageProcessor(BaseImageProcessor, ImageProcessingMixin):
|
| 25 |
+
"""
|
| 26 |
+
Image processor for NemotronParse model.
|
| 27 |
+
|
| 28 |
+
This processor inherits from BaseImageProcessor to be compatible with transformers AutoImageProcessor.
|
| 29 |
+
"""
|
| 30 |
+
|
| 31 |
+
valid_kwargs = NemotronParseImagesKwargs
|
| 32 |
+
model_input_names = ["pixel_values", "image_sizes", "content_sizes"]
|
| 33 |
+
|
| 34 |
+
def __init__(
|
| 35 |
+
self,
|
| 36 |
+
final_size: tuple = (2048, 1664),
|
| 37 |
+
do_normalize: bool = False,
|
| 38 |
+
image_mean: Optional[Union[List[float], tuple]] = None,
|
| 39 |
+
image_std: Optional[Union[List[float], tuple]] = None,
|
| 40 |
+
**kwargs,
|
| 41 |
+
):
|
| 42 |
+
self.navit_packing = bool(kwargs.pop('navit_packing', False))
|
| 43 |
+
self.navit_padding_mode = kwargs.pop('navit_padding_mode', 'top_left_white')
|
| 44 |
+
self.navit_coordinate_frame = kwargs.pop('navit_coordinate_frame', 'content_normalized')
|
| 45 |
+
clean_kwargs = {}
|
| 46 |
+
for k, v in kwargs.items():
|
| 47 |
+
if not k.startswith('_') and k not in ['transform', 'torch_transform']:
|
| 48 |
+
clean_kwargs[k] = v
|
| 49 |
+
|
| 50 |
+
if 'size' in clean_kwargs:
|
| 51 |
+
size_config = clean_kwargs.pop('size')
|
| 52 |
+
if isinstance(size_config, dict):
|
| 53 |
+
if 'longest_edge' in size_config:
|
| 54 |
+
longest_edge = size_config['longest_edge']
|
| 55 |
+
if isinstance(longest_edge, (list, tuple)):
|
| 56 |
+
final_size = tuple(int(x) for x in longest_edge)
|
| 57 |
+
else:
|
| 58 |
+
final_size = (int(longest_edge), int(longest_edge))
|
| 59 |
+
elif 'height' in size_config and 'width' in size_config:
|
| 60 |
+
final_size = (int(size_config['height']), int(size_config['width']))
|
| 61 |
+
|
| 62 |
+
super().__init__(**clean_kwargs)
|
| 63 |
+
self.do_normalize = bool(do_normalize)
|
| 64 |
+
self.image_mean = list(image_mean or OPENAI_CLIP_MEAN)
|
| 65 |
+
self.image_std = list(image_std or OPENAI_CLIP_STD)
|
| 66 |
+
|
| 67 |
+
if isinstance(final_size, (list, tuple)) and len(final_size) >= 2:
|
| 68 |
+
self.final_size = (int(final_size[0]), int(final_size[1]))
|
| 69 |
+
elif isinstance(final_size, (int, float)):
|
| 70 |
+
self.final_size = (int(final_size), int(final_size))
|
| 71 |
+
else:
|
| 72 |
+
self.final_size = (2048, 1664) # Default fallback
|
| 73 |
+
|
| 74 |
+
self._create_transforms()
|
| 75 |
+
|
| 76 |
+
def _create_transforms(self):
|
| 77 |
+
"""Create transform objects (not serialized to JSON)."""
|
| 78 |
+
if isinstance(self.final_size, (list, tuple)):
|
| 79 |
+
self.target_height, self.target_width = int(self.final_size[0]), int(self.final_size[1])
|
| 80 |
+
else:
|
| 81 |
+
self.target_height = self.target_width = int(self.final_size)
|
| 82 |
+
|
| 83 |
+
self.transform = None
|
| 84 |
+
|
| 85 |
+
self.torch_transform = T.Compose([
|
| 86 |
+
T.ToTensor(),
|
| 87 |
+
])
|
| 88 |
+
|
| 89 |
+
def to_dict(self):
|
| 90 |
+
"""Override to exclude non-serializable transforms."""
|
| 91 |
+
output = super().to_dict()
|
| 92 |
+
output.pop('transform', None)
|
| 93 |
+
output.pop('torch_transform', None)
|
| 94 |
+
return output
|
| 95 |
+
|
| 96 |
+
@classmethod
|
| 97 |
+
def from_dict(cls, config_dict: dict, **kwargs):
|
| 98 |
+
"""Override to recreate transforms after loading."""
|
| 99 |
+
config_dict = config_dict.copy()
|
| 100 |
+
config_dict.pop('transform', None)
|
| 101 |
+
config_dict.pop('torch_transform', None)
|
| 102 |
+
|
| 103 |
+
# Clean any problematic entries
|
| 104 |
+
for key in list(config_dict.keys()):
|
| 105 |
+
if key.startswith('_') or config_dict[key] is None:
|
| 106 |
+
config_dict.pop(key, None)
|
| 107 |
+
|
| 108 |
+
# Ensure numeric types are correct
|
| 109 |
+
if 'final_size' in config_dict:
|
| 110 |
+
final_size = config_dict['final_size']
|
| 111 |
+
if isinstance(final_size, (list, tuple)):
|
| 112 |
+
config_dict['final_size'] = tuple(int(x) for x in final_size)
|
| 113 |
+
if 'image_mean' in config_dict:
|
| 114 |
+
config_dict['image_mean'] = [float(x) for x in config_dict['image_mean']]
|
| 115 |
+
if 'image_std' in config_dict:
|
| 116 |
+
config_dict['image_std'] = [float(x) for x in config_dict['image_std']]
|
| 117 |
+
|
| 118 |
+
try:
|
| 119 |
+
return cls(**config_dict, **kwargs)
|
| 120 |
+
except Exception as e:
|
| 121 |
+
print(f"Warning: Error in from_dict: {e}")
|
| 122 |
+
print("Using default parameters...")
|
| 123 |
+
return cls(**kwargs)
|
| 124 |
+
|
| 125 |
+
def save_pretrained(self, save_directory, **kwargs):
|
| 126 |
+
"""Save image processor configuration."""
|
| 127 |
+
import os
|
| 128 |
+
import json
|
| 129 |
+
|
| 130 |
+
os.makedirs(save_directory, exist_ok=True)
|
| 131 |
+
|
| 132 |
+
# Save preprocessor config in standard HuggingFace format
|
| 133 |
+
config = {
|
| 134 |
+
"feature_extractor_type": "NemotronParseImageProcessor",
|
| 135 |
+
"image_processor_type": "NemotronParseImageProcessor",
|
| 136 |
+
"processor_class": "NemotronParseProcessor",
|
| 137 |
+
"size": {
|
| 138 |
+
"height": self.final_size[0],
|
| 139 |
+
"width": self.final_size[1],
|
| 140 |
+
"longest_edge": self.final_size
|
| 141 |
+
},
|
| 142 |
+
"final_size": self.final_size,
|
| 143 |
+
"do_normalize": self.do_normalize,
|
| 144 |
+
"navit_packing": self.navit_packing,
|
| 145 |
+
"navit_padding_mode": self.navit_padding_mode,
|
| 146 |
+
"navit_coordinate_frame": self.navit_coordinate_frame,
|
| 147 |
+
"image_mean": self.image_mean,
|
| 148 |
+
"image_std": self.image_std,
|
| 149 |
+
}
|
| 150 |
+
|
| 151 |
+
config_path = os.path.join(save_directory, "preprocessor_config.json")
|
| 152 |
+
with open(config_path, 'w') as f:
|
| 153 |
+
json.dump(config, f, indent=2)
|
| 154 |
+
|
| 155 |
+
def _resize_with_aspect_ratio(self, image: np.ndarray) -> np.ndarray:
|
| 156 |
+
"""Resize image maintaining aspect ratio (exact replica of original LongestMaxSizeHW)."""
|
| 157 |
+
height, width = image.shape[:2]
|
| 158 |
+
max_size_height = self.target_height
|
| 159 |
+
max_size_width = self.target_width
|
| 160 |
+
|
| 161 |
+
# Original LongestMaxSizeHW algorithm from custom_augmentations.py
|
| 162 |
+
aspect_ratio = width / height
|
| 163 |
+
new_height = height
|
| 164 |
+
new_width = width
|
| 165 |
+
|
| 166 |
+
if height > max_size_height:
|
| 167 |
+
new_height = max_size_height
|
| 168 |
+
new_width = int(new_height * aspect_ratio)
|
| 169 |
+
|
| 170 |
+
if new_width > max_size_width:
|
| 171 |
+
new_width = max_size_width
|
| 172 |
+
new_height = int(new_width / aspect_ratio)
|
| 173 |
+
|
| 174 |
+
if new_height == height and new_width == width:
|
| 175 |
+
return image
|
| 176 |
+
resampling = getattr(Image, "Resampling", Image).BILINEAR
|
| 177 |
+
return np.asarray(Image.fromarray(image).resize((new_width, new_height), resampling))
|
| 178 |
+
|
| 179 |
+
def _pad_to_size(self, image: np.ndarray) -> np.ndarray:
|
| 180 |
+
"""Center-pad image to target size with white padding."""
|
| 181 |
+
h, w = image.shape[:2]
|
| 182 |
+
min_height, min_width = self.target_height, self.target_width
|
| 183 |
+
|
| 184 |
+
pad_h = max(0, min_height - h)
|
| 185 |
+
pad_w = max(0, min_width - w)
|
| 186 |
+
|
| 187 |
+
if pad_h == 0 and pad_w == 0:
|
| 188 |
+
return image
|
| 189 |
+
|
| 190 |
+
pad_top = pad_h // 2
|
| 191 |
+
pad_bottom = pad_h - pad_top
|
| 192 |
+
pad_left = pad_w // 2
|
| 193 |
+
pad_right = pad_w - pad_left
|
| 194 |
+
|
| 195 |
+
if len(image.shape) == 3:
|
| 196 |
+
return np.pad(
|
| 197 |
+
image,
|
| 198 |
+
((pad_top, pad_bottom), (pad_left, pad_right), (0, 0)),
|
| 199 |
+
mode='constant',
|
| 200 |
+
constant_values=255,
|
| 201 |
+
)
|
| 202 |
+
return np.pad(
|
| 203 |
+
image,
|
| 204 |
+
((pad_top, pad_bottom), (pad_left, pad_right)),
|
| 205 |
+
mode='constant',
|
| 206 |
+
constant_values=255,
|
| 207 |
+
)
|
| 208 |
+
|
| 209 |
+
def _pad_to_navit_grid(self, image: np.ndarray) -> np.ndarray:
|
| 210 |
+
height, width = image.shape[:2]
|
| 211 |
+
target_height = ((height + 15) // 16) * 16
|
| 212 |
+
target_width = ((width + 63) // 64) * 64
|
| 213 |
+
pad_h, pad_w = target_height - height, target_width - width
|
| 214 |
+
if not pad_h and not pad_w:
|
| 215 |
+
return image
|
| 216 |
+
padding = ((0, pad_h), (0, pad_w), (0, 0)) if image.ndim == 3 else ((0, pad_h), (0, pad_w))
|
| 217 |
+
return np.pad(image, padding, mode="constant", constant_values=255)
|
| 218 |
+
|
| 219 |
+
def preprocess(
|
| 220 |
+
self,
|
| 221 |
+
images: ImageInput,
|
| 222 |
+
return_tensors: Optional[Union[str, TensorType]] = None,
|
| 223 |
+
**kwargs,
|
| 224 |
+
) -> Dict[str, torch.Tensor]:
|
| 225 |
+
"""
|
| 226 |
+
Preprocess an image or batch of images for the NemotronParse model.
|
| 227 |
+
|
| 228 |
+
Args:
|
| 229 |
+
images: Input image(s)
|
| 230 |
+
return_tensors: Type of tensors to return
|
| 231 |
+
"""
|
| 232 |
+
|
| 233 |
+
captured_hw = kwargs.pop("content_hw", None)
|
| 234 |
+
image_options = kwargs.pop("images_kwargs", {}) or {}
|
| 235 |
+
nested_hw = image_options.get("content_hw")
|
| 236 |
+
if captured_hw is not None and nested_hw is not None and list(captured_hw) != list(nested_hw):
|
| 237 |
+
raise ValueError("conflicting captured content dimensions")
|
| 238 |
+
if captured_hw is None:
|
| 239 |
+
captured_hw = nested_hw
|
| 240 |
+
if captured_hw is not None:
|
| 241 |
+
if not self.navit_packing or self.navit_coordinate_frame != "content_normalized":
|
| 242 |
+
raise ValueError("loader capture requires content-normalized NaViT")
|
| 243 |
+
if self.navit_padding_mode != "top_left_white":
|
| 244 |
+
raise ValueError("unsupported captured loader pixel contract")
|
| 245 |
+
pages = images if isinstance(images, list) else [images]
|
| 246 |
+
if len(pages) != 1 or len(captured_hw) != 2:
|
| 247 |
+
raise ValueError("one captured page and [content_height, content_width] required")
|
| 248 |
+
pixels = np.array(pages[0].convert("RGB"), copy=True)
|
| 249 |
+
height, width = pixels.shape[:2]
|
| 250 |
+
ch, cw = (int(value) for value in captured_hw)
|
| 251 |
+
if list(captured_hw) != [ch, cw] or not (0 < ch <= height and 0 < cw <= width):
|
| 252 |
+
raise ValueError("invalid captured content dimensions")
|
| 253 |
+
if height % 16 or width % 64:
|
| 254 |
+
raise ValueError("captured canvas is not aligned to the NaViT grid")
|
| 255 |
+
page = torch.from_numpy(pixels).permute(2, 0, 1).contiguous().float().div_(255.0)
|
| 256 |
+
# Native RADIO normalizes inside its encoder; the vLLM boundary
|
| 257 |
+
# expects the equivalent normalization here, as in ordinary serving.
|
| 258 |
+
if self.do_normalize:
|
| 259 |
+
mean = page.new_tensor(self.image_mean).view(-1, 1, 1)
|
| 260 |
+
std = page.new_tensor(self.image_std).view(-1, 1, 1)
|
| 261 |
+
page = (page - mean) / std
|
| 262 |
+
patches = page.unfold(1, 16, 16).unfold(2, 16, 16)
|
| 263 |
+
patches = patches.permute(1, 2, 0, 3, 4).reshape(-1, 3 * 16 * 16)
|
| 264 |
+
return BatchFeature(data={
|
| 265 |
+
"pixel_values": patches,
|
| 266 |
+
"image_sizes": torch.tensor([[height, width]], dtype=torch.long),
|
| 267 |
+
"content_sizes": torch.tensor([[ch, cw]], dtype=torch.long),
|
| 268 |
+
}, tensor_type=return_tensors)
|
| 269 |
+
|
| 270 |
+
# Ensure images is a list
|
| 271 |
+
if not isinstance(images, list):
|
| 272 |
+
images = [images]
|
| 273 |
+
|
| 274 |
+
# Ensure images are RGB
|
| 275 |
+
for i in range(len(images)):
|
| 276 |
+
images[i] = images[i].convert('RGB')
|
| 277 |
+
|
| 278 |
+
# Convert PIL images to numpy arrays if needed
|
| 279 |
+
processed_images = []
|
| 280 |
+
for image in images:
|
| 281 |
+
if isinstance(image, Image.Image):
|
| 282 |
+
image = np.asarray(image)
|
| 283 |
+
processed_images.append(image)
|
| 284 |
+
|
| 285 |
+
# Apply NemotronParse-specific transforms
|
| 286 |
+
pixel_values = []
|
| 287 |
+
content_sizes = []
|
| 288 |
+
for image in processed_images:
|
| 289 |
+
processed_image = self._resize_with_aspect_ratio(image)
|
| 290 |
+
content_sizes.append(
|
| 291 |
+
[int(processed_image.shape[0]), int(processed_image.shape[1])]
|
| 292 |
+
)
|
| 293 |
+
if self.navit_packing:
|
| 294 |
+
processed_image = self._pad_to_navit_grid(processed_image)
|
| 295 |
+
|
| 296 |
+
if self.transform is not None:
|
| 297 |
+
transformed = self.transform(image=processed_image)
|
| 298 |
+
processed_image = transformed["image"]
|
| 299 |
+
else:
|
| 300 |
+
# Fixed exports retain the historical full-canvas pad.
|
| 301 |
+
if not self.navit_packing:
|
| 302 |
+
processed_image = self._pad_to_size(processed_image)
|
| 303 |
+
|
| 304 |
+
pixel_values_tensor = self.torch_transform(processed_image)
|
| 305 |
+
|
| 306 |
+
if pixel_values_tensor.shape[0] == 1:
|
| 307 |
+
pixel_values_tensor = pixel_values_tensor.expand(3, -1, -1)
|
| 308 |
+
|
| 309 |
+
pixel_values.append(pixel_values_tensor)
|
| 310 |
+
|
| 311 |
+
if self.navit_packing:
|
| 312 |
+
content_sizes = torch.tensor(content_sizes, dtype=torch.long)
|
| 313 |
+
image_sizes = torch.tensor(
|
| 314 |
+
[[int(page.shape[-2]), int(page.shape[-1])] for page in pixel_values],
|
| 315 |
+
dtype=torch.long,
|
| 316 |
+
)
|
| 317 |
+
raw_patches = []
|
| 318 |
+
for page in pixel_values:
|
| 319 |
+
if self.do_normalize:
|
| 320 |
+
mean = page.new_tensor(self.image_mean).view(-1, 1, 1)
|
| 321 |
+
std = page.new_tensor(self.image_std).view(-1, 1, 1)
|
| 322 |
+
page = (page - mean) / std
|
| 323 |
+
patches = page.unfold(1, 16, 16).unfold(2, 16, 16)
|
| 324 |
+
patches = patches.permute(1, 2, 0, 3, 4).reshape(-1, 3 * 16 * 16)
|
| 325 |
+
raw_patches.append(patches)
|
| 326 |
+
pixel_values = torch.cat(raw_patches, dim=0)
|
| 327 |
+
data = {
|
| 328 |
+
"pixel_values": pixel_values,
|
| 329 |
+
"image_sizes": image_sizes,
|
| 330 |
+
"content_sizes": content_sizes,
|
| 331 |
+
}
|
| 332 |
+
else:
|
| 333 |
+
pixel_values = torch.stack(pixel_values)
|
| 334 |
+
if self.do_normalize:
|
| 335 |
+
mean = pixel_values.new_tensor(self.image_mean).view(1, -1, 1, 1)
|
| 336 |
+
std = pixel_values.new_tensor(self.image_std).view(1, -1, 1, 1)
|
| 337 |
+
pixel_values = (pixel_values - mean) / std
|
| 338 |
+
data = {"pixel_values": pixel_values}
|
| 339 |
+
|
| 340 |
+
if return_tensors is not None:
|
| 341 |
+
data = self._convert_output_format(data, return_tensors)
|
| 342 |
+
|
| 343 |
+
return data
|
| 344 |
+
|
| 345 |
+
def _convert_output_format(self, data: Dict[str, torch.Tensor], return_tensors: Union[str, TensorType]) -> Dict:
|
| 346 |
+
"""Convert output format based on return_tensors parameter."""
|
| 347 |
+
if return_tensors == "pt" or return_tensors == TensorType.PYTORCH:
|
| 348 |
+
return data
|
| 349 |
+
elif return_tensors == "np" or return_tensors == TensorType.NUMPY:
|
| 350 |
+
return {k: v.numpy() for k, v in data.items()}
|
| 351 |
+
else:
|
| 352 |
+
return data
|
| 353 |
+
|
| 354 |
+
def __call__(self, images: Union[Image.Image, List[Image.Image]], **kwargs) -> Dict[str, torch.Tensor]:
|
| 355 |
+
"""Process images for the model (backward compatibility)."""
|
| 356 |
+
return self.preprocess(images, **kwargs)
|
| 357 |
+
|
| 358 |
+
|
| 359 |
+
class NemotronParseProcessor(ProcessorMixin):
|
| 360 |
+
|
| 361 |
+
attributes = ["image_processor", "tokenizer"]
|
| 362 |
+
image_processor_class = "AutoImageProcessor"
|
| 363 |
+
tokenizer_class = ("PreTrainedTokenizer", "PreTrainedTokenizerFast")
|
| 364 |
+
|
| 365 |
+
def __init__(self, image_processor=None, tokenizer=None, **kwargs):
|
| 366 |
+
if image_processor is None:
|
| 367 |
+
image_processor = NemotronParseImageProcessor(**kwargs)
|
| 368 |
+
|
| 369 |
+
super().__init__(image_processor, tokenizer)
|
| 370 |
+
|
| 371 |
+
|
| 372 |
+
def __call__(
|
| 373 |
+
self,
|
| 374 |
+
images: Union[Image.Image, List[Image.Image]] = None,
|
| 375 |
+
text: Union[str, List[str]] = None,
|
| 376 |
+
add_special_tokens: bool = True,
|
| 377 |
+
padding: Union[bool, str] = False,
|
| 378 |
+
truncation: Union[bool, str] = False,
|
| 379 |
+
max_length: Optional[int] = None,
|
| 380 |
+
stride: int = 0,
|
| 381 |
+
pad_to_multiple_of: Optional[int] = None,
|
| 382 |
+
return_attention_mask: Optional[bool] = None,
|
| 383 |
+
return_overflowing_tokens: bool = False,
|
| 384 |
+
return_special_tokens_mask: bool = False,
|
| 385 |
+
return_offsets_mapping: bool = False,
|
| 386 |
+
return_token_type_ids: bool = False,
|
| 387 |
+
return_length: bool = False,
|
| 388 |
+
verbose: bool = True,
|
| 389 |
+
return_tensors: Optional[Union[str, "TensorType"]] = None,
|
| 390 |
+
content_hw: Optional[List[int]] = None,
|
| 391 |
+
**kwargs
|
| 392 |
+
) -> BatchFeature:
|
| 393 |
+
"""
|
| 394 |
+
Main method to prepare for the model one or several text(s) and image(s).
|
| 395 |
+
"""
|
| 396 |
+
|
| 397 |
+
# Process images
|
| 398 |
+
if images is not None:
|
| 399 |
+
image_inputs = self.image_processor(
|
| 400 |
+
images, content_hw=content_hw, **kwargs)
|
| 401 |
+
else:
|
| 402 |
+
image_inputs = {}
|
| 403 |
+
|
| 404 |
+
# Process text
|
| 405 |
+
if text is not None:
|
| 406 |
+
text_inputs = self.tokenizer(
|
| 407 |
+
text,
|
| 408 |
+
add_special_tokens=add_special_tokens,
|
| 409 |
+
padding=padding,
|
| 410 |
+
truncation=truncation,
|
| 411 |
+
max_length=max_length,
|
| 412 |
+
stride=stride,
|
| 413 |
+
pad_to_multiple_of=pad_to_multiple_of,
|
| 414 |
+
return_attention_mask=return_attention_mask,
|
| 415 |
+
return_overflowing_tokens=return_overflowing_tokens,
|
| 416 |
+
return_special_tokens_mask=return_special_tokens_mask,
|
| 417 |
+
return_offsets_mapping=return_offsets_mapping,
|
| 418 |
+
return_token_type_ids=return_token_type_ids,
|
| 419 |
+
return_length=return_length,
|
| 420 |
+
verbose=verbose,
|
| 421 |
+
return_tensors=return_tensors,
|
| 422 |
+
**kwargs,
|
| 423 |
+
)
|
| 424 |
+
else:
|
| 425 |
+
text_inputs = {}
|
| 426 |
+
|
| 427 |
+
# Combine inputs
|
| 428 |
+
return BatchFeature(data={**image_inputs, **text_inputs})
|
| 429 |
+
|
| 430 |
+
def decode(self, *args, **kwargs):
|
| 431 |
+
"""Decode token ids to strings."""
|
| 432 |
+
return self.tokenizer.decode(*args, **kwargs)
|
| 433 |
+
|
| 434 |
+
def batch_decode(self, *args, **kwargs):
|
| 435 |
+
"""Batch decode token ids to strings."""
|
| 436 |
+
return self.tokenizer.batch_decode(*args, **kwargs)
|
| 437 |
+
|
| 438 |
+
def post_process_generation(self, sequences, fix_markdown=False):
|
| 439 |
+
"""Post-process generated sequences."""
|
| 440 |
+
if hasattr(self.tokenizer, 'post_process_generation'):
|
| 441 |
+
return self.tokenizer.post_process_generation(sequences, fix_markdown=fix_markdown)
|
| 442 |
+
else:
|
| 443 |
+
# Fallback processing
|
| 444 |
+
if isinstance(sequences, str):
|
| 445 |
+
sequences = [sequences]
|
| 446 |
+
|
| 447 |
+
processed = []
|
| 448 |
+
for seq in sequences:
|
| 449 |
+
# Basic cleaning
|
| 450 |
+
seq = seq.replace('<s>', '').replace('</s>', '').strip()
|
| 451 |
+
processed.append(seq)
|
| 452 |
+
|
| 453 |
+
return processed[0] if len(processed) == 1 else processed
|
| 454 |
+
|
| 455 |
+
@classmethod
|
| 456 |
+
def from_pretrained(cls, pretrained_model_name_or_path, **kwargs):
|
| 457 |
+
"""
|
| 458 |
+
Load processor from pretrained model.
|
| 459 |
+
|
| 460 |
+
This method is compatible with AutoProcessor.from_pretrained().
|
| 461 |
+
"""
|
| 462 |
+
# Explicitly load subcomponents via Auto* to ensure remote auto_map is honored.
|
| 463 |
+
from transformers import AutoImageProcessor, AutoTokenizer
|
| 464 |
+
trust_remote_code = kwargs.get("trust_remote_code", None)
|
| 465 |
+
revision = kwargs.get("revision", None)
|
| 466 |
+
token = kwargs.get("token", None)
|
| 467 |
+
image_processor = AutoImageProcessor.from_pretrained(
|
| 468 |
+
pretrained_model_name_or_path,
|
| 469 |
+
trust_remote_code=trust_remote_code,
|
| 470 |
+
revision=revision,
|
| 471 |
+
token=token,
|
| 472 |
+
)
|
| 473 |
+
tokenizer = AutoTokenizer.from_pretrained(
|
| 474 |
+
pretrained_model_name_or_path,
|
| 475 |
+
trust_remote_code=trust_remote_code,
|
| 476 |
+
revision=revision,
|
| 477 |
+
token=token,
|
| 478 |
+
)
|
| 479 |
+
return cls(image_processor=image_processor, tokenizer=tokenizer)
|
| 480 |
+
|
| 481 |
+
def save_pretrained(self, save_directory, **kwargs):
|
| 482 |
+
"""
|
| 483 |
+
Save processor to directory.
|
| 484 |
+
|
| 485 |
+
This method is compatible with AutoProcessor/AutoImageProcessor loading.
|
| 486 |
+
"""
|
| 487 |
+
import os
|
| 488 |
+
os.makedirs(save_directory, exist_ok=True)
|
| 489 |
+
|
| 490 |
+
# Save tokenizer with proper configuration for AutoTokenizer
|
| 491 |
+
print("Saving tokenizer for AutoTokenizer compatibility...")
|
| 492 |
+
self.tokenizer.save_pretrained(save_directory, **kwargs)
|
| 493 |
+
|
| 494 |
+
# Save image processor
|
| 495 |
+
print("Saving image processor...")
|
| 496 |
+
self.image_processor.save_pretrained(save_directory, **kwargs)
|
| 497 |
+
|
| 498 |
+
# Use the parent class's save_pretrained method for processor config
|
| 499 |
+
super().save_pretrained(save_directory, **kwargs)
|
| 500 |
+
print(f"NemotronParseProcessor saved to {save_directory}")
|
| 501 |
+
print(f"AutoTokenizer.from_pretrained('{save_directory}') should now work!")
|
hf_prefix_geometry.py
ADDED
|
@@ -0,0 +1,149 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
| 2 |
+
# SPDX-License-Identifier: OpenMDW-1.1
|
| 3 |
+
|
| 4 |
+
"""Learned content-normalized positional interpolation shared with serving."""
|
| 5 |
+
|
| 6 |
+
def _normalized_2d_position_pages(
|
| 7 |
+
torch_module,
|
| 8 |
+
*,
|
| 9 |
+
patches,
|
| 10 |
+
pos_embed,
|
| 11 |
+
source_rows: int,
|
| 12 |
+
source_cols: int,
|
| 13 |
+
image_sizes,
|
| 14 |
+
patch_size: int,
|
| 15 |
+
content_sizes=None,
|
| 16 |
+
content_normalized: bool = False,
|
| 17 |
+
):
|
| 18 |
+
"""Mirror training's per-page normalized-2D position interpolation.
|
| 19 |
+
|
| 20 |
+
Keep this helper independent of vLLM imports so the export contract can be
|
| 21 |
+
exercised on the login node. vLLM supplies one concatenated raw-patch
|
| 22 |
+
sequence, whereas training embeds each page separately; splitting before
|
| 23 |
+
interpolation makes the two representations numerically identical.
|
| 24 |
+
"""
|
| 25 |
+
|
| 26 |
+
source_rows = int(source_rows)
|
| 27 |
+
source_cols = int(source_cols)
|
| 28 |
+
patch_size = int(patch_size)
|
| 29 |
+
if source_rows <= 0 or source_cols <= 0 or patch_size <= 0:
|
| 30 |
+
raise ValueError(
|
| 31 |
+
"normalized 2D positions require positive source grid and patch size"
|
| 32 |
+
)
|
| 33 |
+
if source_rows * source_cols != int(pos_embed.shape[1]):
|
| 34 |
+
raise ValueError(
|
| 35 |
+
"RADIO learned positional field shape is inconsistent: "
|
| 36 |
+
f"source={source_rows}x{source_cols} tokens={int(pos_embed.shape[1])}"
|
| 37 |
+
)
|
| 38 |
+
|
| 39 |
+
positional_grid = pos_embed.reshape(
|
| 40 |
+
1, source_rows, source_cols, -1
|
| 41 |
+
).permute(0, 3, 1, 2)
|
| 42 |
+
normalized_content_sizes = None
|
| 43 |
+
if content_sizes is not None:
|
| 44 |
+
normalized_content_sizes = [
|
| 45 |
+
tuple(int(value) for value in size) for size in content_sizes
|
| 46 |
+
]
|
| 47 |
+
if len(normalized_content_sizes) != len(image_sizes):
|
| 48 |
+
raise ValueError(
|
| 49 |
+
"Content-size count must match image-size count: "
|
| 50 |
+
f"content_sizes={len(normalized_content_sizes)} "
|
| 51 |
+
f"image_sizes={len(image_sizes)}"
|
| 52 |
+
)
|
| 53 |
+
if content_normalized and normalized_content_sizes is None:
|
| 54 |
+
raise ValueError("content-normalized 2D positions require content_sizes")
|
| 55 |
+
|
| 56 |
+
positioned_pages = []
|
| 57 |
+
cursor = 0
|
| 58 |
+
for image_index, (raw_height, raw_width) in enumerate(image_sizes):
|
| 59 |
+
height, width = int(raw_height), int(raw_width)
|
| 60 |
+
if height % patch_size or width % patch_size:
|
| 61 |
+
raise ValueError(
|
| 62 |
+
f"Page shape {(height, width)} is not aligned to patch size {patch_size}"
|
| 63 |
+
)
|
| 64 |
+
rows, cols = height // patch_size, width // patch_size
|
| 65 |
+
patch_count = rows * cols
|
| 66 |
+
page = patches[:, cursor : cursor + patch_count]
|
| 67 |
+
if int(page.shape[1]) != patch_count:
|
| 68 |
+
raise ValueError(
|
| 69 |
+
"Normalized position split ended inside a page: "
|
| 70 |
+
f"page={(height, width)} expected={patch_count} actual={int(page.shape[1])}"
|
| 71 |
+
)
|
| 72 |
+
cursor += patch_count
|
| 73 |
+
if content_normalized:
|
| 74 |
+
content_height, content_width = normalized_content_sizes[image_index]
|
| 75 |
+
if content_height <= 0 or content_width <= 0:
|
| 76 |
+
raise ValueError(
|
| 77 |
+
"content size must be positive, got "
|
| 78 |
+
f"{(content_height, content_width)}"
|
| 79 |
+
)
|
| 80 |
+
if content_height > height or content_width > width:
|
| 81 |
+
raise ValueError(
|
| 82 |
+
"content size cannot exceed padded image: "
|
| 83 |
+
f"content={(content_height, content_width)} image={(height, width)}"
|
| 84 |
+
)
|
| 85 |
+
content_rows = (content_height + patch_size - 1) // patch_size
|
| 86 |
+
content_cols = (content_width + patch_size - 1) // patch_size
|
| 87 |
+
if min(content_rows, content_cols) == 1:
|
| 88 |
+
# Match training's safe area-preserving path for augmentation
|
| 89 |
+
# crops whose visible content spans only one patch in an axis.
|
| 90 |
+
page_positions = torch_module.nn.functional.adaptive_avg_pool2d(
|
| 91 |
+
positional_grid.float(), (content_rows, content_cols)
|
| 92 |
+
)
|
| 93 |
+
else:
|
| 94 |
+
page_positions = torch_module.nn.functional.interpolate(
|
| 95 |
+
positional_grid.float(),
|
| 96 |
+
size=(content_rows, content_cols),
|
| 97 |
+
mode="bilinear",
|
| 98 |
+
align_corners=False,
|
| 99 |
+
antialias=True,
|
| 100 |
+
)
|
| 101 |
+
row_starts = (
|
| 102 |
+
torch_module.arange(rows, device=page.device, dtype=torch_module.float32)
|
| 103 |
+
* patch_size
|
| 104 |
+
)
|
| 105 |
+
col_starts = (
|
| 106 |
+
torch_module.arange(cols, device=page.device, dtype=torch_module.float32)
|
| 107 |
+
* patch_size
|
| 108 |
+
)
|
| 109 |
+
row_extents = (content_height - row_starts).clamp(
|
| 110 |
+
min=0.0, max=float(patch_size)
|
| 111 |
+
)
|
| 112 |
+
col_extents = (content_width - col_starts).clamp(
|
| 113 |
+
min=0.0, max=float(patch_size)
|
| 114 |
+
)
|
| 115 |
+
normalized_y = (
|
| 116 |
+
row_starts + 0.5 * row_extents
|
| 117 |
+
).clamp(max=float(content_height)) / float(content_height)
|
| 118 |
+
normalized_x = (
|
| 119 |
+
col_starts + 0.5 * col_extents
|
| 120 |
+
).clamp(max=float(content_width)) / float(content_width)
|
| 121 |
+
grid_y, grid_x = torch_module.meshgrid(
|
| 122 |
+
normalized_y, normalized_x, indexing="ij"
|
| 123 |
+
)
|
| 124 |
+
sampling_grid = torch_module.stack(
|
| 125 |
+
(2.0 * grid_x - 1.0, 2.0 * grid_y - 1.0), dim=-1
|
| 126 |
+
)
|
| 127 |
+
page_positions = torch_module.nn.functional.grid_sample(
|
| 128 |
+
page_positions,
|
| 129 |
+
sampling_grid.unsqueeze(0),
|
| 130 |
+
mode="bilinear",
|
| 131 |
+
padding_mode="border",
|
| 132 |
+
align_corners=False,
|
| 133 |
+
).to(dtype=page.dtype)
|
| 134 |
+
else:
|
| 135 |
+
page_positions = torch_module.nn.functional.interpolate(
|
| 136 |
+
positional_grid.float(),
|
| 137 |
+
size=(rows, cols),
|
| 138 |
+
mode="bilinear",
|
| 139 |
+
align_corners=False,
|
| 140 |
+
antialias=True,
|
| 141 |
+
).to(dtype=page.dtype)
|
| 142 |
+
page_positions = page_positions.flatten(2).permute(0, 2, 1)
|
| 143 |
+
positioned_pages.append(page + page_positions)
|
| 144 |
+
if cursor != int(patches.shape[1]):
|
| 145 |
+
raise ValueError(
|
| 146 |
+
f"Normalized position split consumed {cursor} of {int(patches.shape[1])} patches"
|
| 147 |
+
)
|
| 148 |
+
return torch_module.cat(positioned_pages, dim=1)
|
| 149 |
+
|
latex2html.py
ADDED
|
@@ -0,0 +1,11 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
| 2 |
+
# SPDX-License-Identifier: OpenMDW-1.1
|
| 3 |
+
"""Compatibility imports; implementation lives in nemotron_parse.tables."""
|
| 4 |
+
from nemotron_parse.tables import (
|
| 5 |
+
latex_table_to_html,
|
| 6 |
+
convert_single_table,
|
| 7 |
+
convert_html_tables_to_markdown,
|
| 8 |
+
convert_html_table_to_json,
|
| 9 |
+
convert_html_table_to_json_hierarchical,
|
| 10 |
+
convert_html_table_to_csv,
|
| 11 |
+
)
|
logitsprocs/nemotron_parse_vllm_logitprocs.py
ADDED
|
@@ -0,0 +1,350 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
| 2 |
+
# SPDX-License-Identifier: OpenMDW-1.1
|
| 3 |
+
|
| 4 |
+
"""
|
| 5 |
+
Custom vLLM v1 logits processors for NVIDIA Nemotron Parse.
|
| 6 |
+
|
| 7 |
+
Tuning via environment variables:
|
| 8 |
+
- NEMOTRON_PARSE_TABLE_PREFIX (default: \\begin{tabular})
|
| 9 |
+
- NEMOTRON_PARSE_REP_MAX (default: 10)
|
| 10 |
+
- NEMOTRON_PARSE_REP_WINDOW (default: 2 * max_ngram * (rep_max + 1))
|
| 11 |
+
- NEMOTRON_PARSE_REP_NGRAMS (default: 1,2,3,4,5)
|
| 12 |
+
"""
|
| 13 |
+
|
| 14 |
+
from __future__ import annotations
|
| 15 |
+
|
| 16 |
+
import os
|
| 17 |
+
import warnings
|
| 18 |
+
from dataclasses import dataclass
|
| 19 |
+
from typing import Dict, List, Set
|
| 20 |
+
|
| 21 |
+
import torch
|
| 22 |
+
from transformers import AutoTokenizer
|
| 23 |
+
|
| 24 |
+
from vllm.v1.sample.logits_processor.builtin import process_dict_updates
|
| 25 |
+
from vllm.v1.sample.logits_processor.interface import BatchUpdate, LogitsProcessor
|
| 26 |
+
|
| 27 |
+
|
| 28 |
+
def _strip_trailing_negative_token_ids(token_ids: List[int]) -> List[int]:
|
| 29 |
+
"""vLLM v1 keeps a trailing -1 placeholder in output_token_ids."""
|
| 30 |
+
i = len(token_ids)
|
| 31 |
+
while i > 0 and token_ids[i - 1] < 0:
|
| 32 |
+
i -= 1
|
| 33 |
+
return token_ids[:i]
|
| 34 |
+
|
| 35 |
+
|
| 36 |
+
def _env_int(name: str, default: int) -> int:
|
| 37 |
+
val = os.environ.get(name, "").strip()
|
| 38 |
+
if not val:
|
| 39 |
+
return default
|
| 40 |
+
try:
|
| 41 |
+
return int(val)
|
| 42 |
+
except ValueError:
|
| 43 |
+
return default
|
| 44 |
+
|
| 45 |
+
|
| 46 |
+
def _env_csv_ints(name: str, default: List[int]) -> List[int]:
|
| 47 |
+
val = os.environ.get(name, "").strip()
|
| 48 |
+
if not val:
|
| 49 |
+
return default
|
| 50 |
+
out: List[int] = []
|
| 51 |
+
for part in val.split(","):
|
| 52 |
+
part = part.strip()
|
| 53 |
+
if not part:
|
| 54 |
+
continue
|
| 55 |
+
try:
|
| 56 |
+
out.append(int(part))
|
| 57 |
+
except ValueError:
|
| 58 |
+
return default
|
| 59 |
+
return out or default
|
| 60 |
+
|
| 61 |
+
|
| 62 |
+
def _warn_if_unexpected_arch(vllm_config) -> None:
|
| 63 |
+
try:
|
| 64 |
+
arches = getattr(vllm_config.model_config, "architectures", None) or []
|
| 65 |
+
except Exception:
|
| 66 |
+
arches = []
|
| 67 |
+
if any(a == "NemotronParseForConditionalGeneration" for a in arches):
|
| 68 |
+
return
|
| 69 |
+
warnings.warn(
|
| 70 |
+
"Nemotron-Parse logits processors enabled for a model whose "
|
| 71 |
+
f"architectures={arches!r}. These processors assume Nemotron-Parse-style "
|
| 72 |
+
"special tokens like <x_...>, <y_...>, and <class_...>; behavior may be "
|
| 73 |
+
"unsupported for other models.",
|
| 74 |
+
RuntimeWarning,
|
| 75 |
+
stacklevel=2,
|
| 76 |
+
)
|
| 77 |
+
|
| 78 |
+
|
| 79 |
+
def _load_hf_tokenizer(vllm_config):
|
| 80 |
+
model_cfg = vllm_config.model_config
|
| 81 |
+
tok_name = getattr(model_cfg, "tokenizer", None) or getattr(model_cfg, "model", None)
|
| 82 |
+
trust_rc = bool(getattr(model_cfg, "trust_remote_code", False))
|
| 83 |
+
revision = getattr(model_cfg, "tokenizer_revision", None) or getattr(model_cfg, "revision", None)
|
| 84 |
+
token = getattr(model_cfg, "hf_token", None)
|
| 85 |
+
kwargs = {"trust_remote_code": trust_rc}
|
| 86 |
+
if revision:
|
| 87 |
+
kwargs["revision"] = revision
|
| 88 |
+
if token:
|
| 89 |
+
# HF uses `token=` in modern versions.
|
| 90 |
+
kwargs["token"] = token
|
| 91 |
+
return AutoTokenizer.from_pretrained(tok_name, **kwargs)
|
| 92 |
+
|
| 93 |
+
|
| 94 |
+
def _build_token_sets(tokenizer) -> tuple[Set[int], Set[int], Set[int]]:
|
| 95 |
+
"""Build token-id sets for <x_...>, <y_...>, and <class_...>."""
|
| 96 |
+
|
| 97 |
+
def _scan(items) -> tuple[Set[int], Set[int], Set[int]]:
|
| 98 |
+
x_ids: Set[int] = set()
|
| 99 |
+
y_ids: Set[int] = set()
|
| 100 |
+
class_ids: Set[int] = set()
|
| 101 |
+
for token, tid in items:
|
| 102 |
+
if token.startswith("<x_") and token.endswith(">") and token.count(">") == 1:
|
| 103 |
+
x_ids.add(tid)
|
| 104 |
+
elif token.startswith("<y_") and token.endswith(">") and token.count(">") == 1:
|
| 105 |
+
y_ids.add(tid)
|
| 106 |
+
elif token.startswith("<class_") and token.endswith(">") and token.count(">") == 1:
|
| 107 |
+
class_ids.add(tid)
|
| 108 |
+
return x_ids, y_ids, class_ids
|
| 109 |
+
|
| 110 |
+
added_vocab = {}
|
| 111 |
+
if hasattr(tokenizer, "get_added_vocab"):
|
| 112 |
+
try:
|
| 113 |
+
added_vocab = tokenizer.get_added_vocab() or {}
|
| 114 |
+
except Exception:
|
| 115 |
+
added_vocab = {}
|
| 116 |
+
|
| 117 |
+
x_ids, y_ids, class_ids = _scan(added_vocab.items())
|
| 118 |
+
|
| 119 |
+
if len(x_ids) < 100 or len(y_ids) < 100 or not class_ids:
|
| 120 |
+
vocab = tokenizer.get_vocab()
|
| 121 |
+
x2, y2, c2 = _scan(vocab.items())
|
| 122 |
+
x_ids, y_ids, class_ids = x2, y2, c2
|
| 123 |
+
|
| 124 |
+
return x_ids, y_ids, class_ids
|
| 125 |
+
|
| 126 |
+
|
| 127 |
+
@dataclass
|
| 128 |
+
class _TableReqState:
|
| 129 |
+
output_ids: List[int]
|
| 130 |
+
insertion_active: bool = False
|
| 131 |
+
insertion_pos: int = 0
|
| 132 |
+
expecting_end_coords: bool = False
|
| 133 |
+
_last_class_pos: int = -1
|
| 134 |
+
|
| 135 |
+
|
| 136 |
+
class NemotronParseTableInsertionLogitsProcessor(LogitsProcessor):
|
| 137 |
+
"""Force a table prefix right after each START <x_...><y_...> coordinate pair."""
|
| 138 |
+
|
| 139 |
+
def __init__(self, vllm_config, device: torch.device, is_pin_memory: bool) -> None:
|
| 140 |
+
self.enabled = True
|
| 141 |
+
self.req_states: Dict[int, _TableReqState] = {}
|
| 142 |
+
|
| 143 |
+
self.table_prefix = os.environ.get("NEMOTRON_PARSE_TABLE_PREFIX", r"\begin{tabular}")
|
| 144 |
+
self.table_prefix_ids: List[int] = []
|
| 145 |
+
self.x_token_ids: Set[int] = set()
|
| 146 |
+
self.y_token_ids: Set[int] = set()
|
| 147 |
+
self.class_token_ids: Set[int] = set()
|
| 148 |
+
|
| 149 |
+
if self.enabled:
|
| 150 |
+
_warn_if_unexpected_arch(vllm_config)
|
| 151 |
+
tok = _load_hf_tokenizer(vllm_config)
|
| 152 |
+
self.table_prefix_ids = tok.encode(self.table_prefix, add_special_tokens=False)
|
| 153 |
+
self.x_token_ids, self.y_token_ids, self.class_token_ids = _build_token_sets(tok)
|
| 154 |
+
|
| 155 |
+
if not self.x_token_ids or not self.y_token_ids or not self.table_prefix_ids:
|
| 156 |
+
self.enabled = False
|
| 157 |
+
|
| 158 |
+
@classmethod
|
| 159 |
+
def validate_params(cls, sampling_params):
|
| 160 |
+
return None
|
| 161 |
+
|
| 162 |
+
def is_argmax_invariant(self) -> bool:
|
| 163 |
+
return False
|
| 164 |
+
|
| 165 |
+
def update_state(self, batch_update: BatchUpdate | None) -> None:
|
| 166 |
+
if not self.enabled:
|
| 167 |
+
return
|
| 168 |
+
|
| 169 |
+
def _new_state(params, prompt_ids, output_ids):
|
| 170 |
+
return _TableReqState(output_ids=output_ids)
|
| 171 |
+
|
| 172 |
+
process_dict_updates(self.req_states, batch_update, _new_state)
|
| 173 |
+
|
| 174 |
+
for st in self.req_states.values():
|
| 175 |
+
out = st.output_ids
|
| 176 |
+
if not out:
|
| 177 |
+
continue
|
| 178 |
+
|
| 179 |
+
view = _strip_trailing_negative_token_ids(out)
|
| 180 |
+
if not view:
|
| 181 |
+
continue
|
| 182 |
+
|
| 183 |
+
# If we ended an object (a newly-generated <class_...>), next xy should be START.
|
| 184 |
+
if self.class_token_ids:
|
| 185 |
+
last_class_pos = -1
|
| 186 |
+
for i in range(len(view) - 1, -1, -1):
|
| 187 |
+
if view[i] in self.class_token_ids:
|
| 188 |
+
last_class_pos = i
|
| 189 |
+
break
|
| 190 |
+
if last_class_pos != -1 and last_class_pos != st._last_class_pos:
|
| 191 |
+
st._last_class_pos = last_class_pos
|
| 192 |
+
st.expecting_end_coords = False
|
| 193 |
+
|
| 194 |
+
if st.insertion_active:
|
| 195 |
+
continue
|
| 196 |
+
|
| 197 |
+
if len(view) >= 2 and view[-2] in self.x_token_ids and view[-1] in self.y_token_ids:
|
| 198 |
+
if not st.expecting_end_coords:
|
| 199 |
+
st.insertion_active = True
|
| 200 |
+
st.insertion_pos = 0
|
| 201 |
+
st.expecting_end_coords = True
|
| 202 |
+
else:
|
| 203 |
+
pass
|
| 204 |
+
|
| 205 |
+
def apply(self, logits: torch.Tensor) -> torch.Tensor:
|
| 206 |
+
if not self.enabled or not self.req_states:
|
| 207 |
+
return logits
|
| 208 |
+
|
| 209 |
+
for req_idx, st in self.req_states.items():
|
| 210 |
+
if not st.insertion_active:
|
| 211 |
+
continue
|
| 212 |
+
pos = st.insertion_pos
|
| 213 |
+
if pos >= len(self.table_prefix_ids):
|
| 214 |
+
st.insertion_active = False
|
| 215 |
+
continue
|
| 216 |
+
|
| 217 |
+
forced_tid = self.table_prefix_ids[pos]
|
| 218 |
+
logits[req_idx].fill_(-float("inf"))
|
| 219 |
+
logits[req_idx, forced_tid] = 0.0
|
| 220 |
+
|
| 221 |
+
st.insertion_pos = pos + 1
|
| 222 |
+
if st.insertion_pos >= len(self.table_prefix_ids):
|
| 223 |
+
st.insertion_active = False
|
| 224 |
+
|
| 225 |
+
return logits
|
| 226 |
+
|
| 227 |
+
|
| 228 |
+
@dataclass
|
| 229 |
+
class _RepReqState:
|
| 230 |
+
output_ids: List[int]
|
| 231 |
+
in_cooldown: bool = False
|
| 232 |
+
segment_start: int = 0
|
| 233 |
+
_last_class_pos: int = -1
|
| 234 |
+
|
| 235 |
+
|
| 236 |
+
class NemotronParseRepetitionStopLogitsProcessor(LogitsProcessor):
|
| 237 |
+
"""Force an <x_...> token when consecutive repetition exceeds a threshold."""
|
| 238 |
+
|
| 239 |
+
def __init__(self, vllm_config, device: torch.device, is_pin_memory: bool) -> None:
|
| 240 |
+
self.enabled = True
|
| 241 |
+
self.req_states: Dict[int, _RepReqState] = {}
|
| 242 |
+
|
| 243 |
+
self.max_repetitions = _env_int("NEMOTRON_PARSE_REP_MAX", 10)
|
| 244 |
+
self.ngram_sizes = _env_csv_ints("NEMOTRON_PARSE_REP_NGRAMS", [1, 2, 3, 4, 5])
|
| 245 |
+
max_ngram = max((n for n in self.ngram_sizes if n > 0), default=1)
|
| 246 |
+
default_window = 2 * max_ngram * (self.max_repetitions + 1)
|
| 247 |
+
self.window_size = _env_int("NEMOTRON_PARSE_REP_WINDOW", default_window)
|
| 248 |
+
|
| 249 |
+
self.x_token_ids: Set[int] = set()
|
| 250 |
+
self.class_token_ids: Set[int] = set()
|
| 251 |
+
|
| 252 |
+
if self.enabled:
|
| 253 |
+
_warn_if_unexpected_arch(vllm_config)
|
| 254 |
+
tok = _load_hf_tokenizer(vllm_config)
|
| 255 |
+
x_ids, _, class_ids = _build_token_sets(tok)
|
| 256 |
+
self.x_token_ids = x_ids
|
| 257 |
+
self.class_token_ids = class_ids
|
| 258 |
+
|
| 259 |
+
if not self.x_token_ids:
|
| 260 |
+
self.enabled = False
|
| 261 |
+
|
| 262 |
+
@classmethod
|
| 263 |
+
def validate_params(cls, sampling_params):
|
| 264 |
+
return None
|
| 265 |
+
|
| 266 |
+
def is_argmax_invariant(self) -> bool:
|
| 267 |
+
return False
|
| 268 |
+
|
| 269 |
+
def update_state(self, batch_update: BatchUpdate | None) -> None:
|
| 270 |
+
if not self.enabled:
|
| 271 |
+
return
|
| 272 |
+
|
| 273 |
+
def _new_state(params, prompt_ids, output_ids):
|
| 274 |
+
return _RepReqState(output_ids=output_ids)
|
| 275 |
+
|
| 276 |
+
process_dict_updates(self.req_states, batch_update, _new_state)
|
| 277 |
+
|
| 278 |
+
for st in self.req_states.values():
|
| 279 |
+
out = st.output_ids
|
| 280 |
+
if not out:
|
| 281 |
+
continue
|
| 282 |
+
view = _strip_trailing_negative_token_ids(out)
|
| 283 |
+
if not view:
|
| 284 |
+
continue
|
| 285 |
+
# Clear cooldown and advance segment on a newly-generated <class_...> token.
|
| 286 |
+
if self.class_token_ids:
|
| 287 |
+
last_class_pos = -1
|
| 288 |
+
for i in range(len(view) - 1, -1, -1):
|
| 289 |
+
if view[i] in self.class_token_ids:
|
| 290 |
+
last_class_pos = i
|
| 291 |
+
break
|
| 292 |
+
if last_class_pos != -1 and last_class_pos != st._last_class_pos:
|
| 293 |
+
st._last_class_pos = last_class_pos
|
| 294 |
+
st.in_cooldown = False
|
| 295 |
+
st.segment_start = last_class_pos + 1
|
| 296 |
+
|
| 297 |
+
@staticmethod
|
| 298 |
+
def _max_consecutive_repetitions(seq: List[int], n: int) -> int:
|
| 299 |
+
if len(seq) < n:
|
| 300 |
+
return 0
|
| 301 |
+
max_consec = 1
|
| 302 |
+
cur = 1
|
| 303 |
+
prev = tuple(seq[0:n])
|
| 304 |
+
i = n
|
| 305 |
+
while i <= len(seq) - n:
|
| 306 |
+
cur_ng = tuple(seq[i:i + n])
|
| 307 |
+
if cur_ng == prev:
|
| 308 |
+
cur += 1
|
| 309 |
+
if cur > max_consec:
|
| 310 |
+
max_consec = cur
|
| 311 |
+
i += n
|
| 312 |
+
else:
|
| 313 |
+
cur = 1
|
| 314 |
+
prev = cur_ng
|
| 315 |
+
i += 1
|
| 316 |
+
return max_consec
|
| 317 |
+
|
| 318 |
+
def _has_excessive_repetition(self, seq: List[int]) -> bool:
|
| 319 |
+
if not seq:
|
| 320 |
+
return False
|
| 321 |
+
check_seq = seq[-self.window_size:] if len(seq) > self.window_size else seq
|
| 322 |
+
for n in self.ngram_sizes:
|
| 323 |
+
if n <= 0:
|
| 324 |
+
continue
|
| 325 |
+
if self._max_consecutive_repetitions(check_seq, n) > self.max_repetitions:
|
| 326 |
+
return True
|
| 327 |
+
return False
|
| 328 |
+
|
| 329 |
+
def apply(self, logits: torch.Tensor) -> torch.Tensor:
|
| 330 |
+
if not self.enabled or not self.req_states:
|
| 331 |
+
return logits
|
| 332 |
+
|
| 333 |
+
for req_idx, st in self.req_states.items():
|
| 334 |
+
if st.in_cooldown:
|
| 335 |
+
continue
|
| 336 |
+
view = _strip_trailing_negative_token_ids(st.output_ids)
|
| 337 |
+
seg = view[st.segment_start:]
|
| 338 |
+
if not self._has_excessive_repetition(seg):
|
| 339 |
+
continue
|
| 340 |
+
|
| 341 |
+
st.in_cooldown = True
|
| 342 |
+
row = logits[req_idx]
|
| 343 |
+
original = row.clone()
|
| 344 |
+
row.fill_(-float("inf"))
|
| 345 |
+
for tid in self.x_token_ids:
|
| 346 |
+
row[tid] = original[tid]
|
| 347 |
+
|
| 348 |
+
return logits
|
| 349 |
+
|
| 350 |
+
|
model-cards/nemotron-3-5-parse-subcards/bias.md
ADDED
|
@@ -0,0 +1,7 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
## Bias
|
| 2 |
+
|
| 3 |
+
Field | Response
|
| 4 |
+
:---- | :----
|
| 5 |
+
Participation considerations from adversely impacted groups [protected classes](https://calcivilrights.ca.gov/disputeresolution/protected-characteristics/) in model design and testing: | None
|
| 6 |
+
Measures taken to mitigate against unwanted bias: | Dataset curation, filtering, and evaluation include review across document types, layouts, languages, and scripts where available. Users should evaluate performance on their target document domains before deployment.
|
| 7 |
+
Bias Metric (If Measured): | Not measured as a demographic fairness metric. Model quality is evaluated through document parsing, OCR, layout, table, chart, and reading-order metrics.
|
model-cards/nemotron-3-5-parse-subcards/explainability.md
ADDED
|
@@ -0,0 +1,15 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
## Explainability
|
| 2 |
+
|
| 3 |
+
Field | Response
|
| 4 |
+
:---- | :----
|
| 5 |
+
Intended Task/Domain: | Document understanding, OCR, layout parsing, table extraction, and text extraction with spatial grounding.
|
| 6 |
+
Model Type: | Transformer-based vision-encoder-decoder model.
|
| 7 |
+
Intended Users: | Developers and teams building document intelligence, retrieval-augmented generation (RAG), data-curation, extraction, and multimodal agent workflows.
|
| 8 |
+
Output: | Structured text containing document element classes, bounding boxes, reading-order information, chart/table annotations, and extracted text.
|
| 9 |
+
Describe how the model works: | Nemotron 3.5 Parse encodes a dynamically sized document image with C-RADIOv4-SO400M, compresses the visual features with a neck, and places them before the OCR prompt as a causal prefix to a 10-layer mBART decoder. The decoder uses grouped-query self-attention and produces structured document annotations with content-normalized boxes. The 72,505-token tokenizer includes multilingual, task-control, and chart-aware tokens such as `<class_Chart>`.
|
| 10 |
+
Name the adversely impacted groups this has been tested to deliver comparable outcomes regardless of: | Not Applicable
|
| 11 |
+
Technical Limitations & Mitigation: | Performance can vary for low-resolution scans, motion blur, handwriting, decorative fonts, dense tables, unusual layouts, unsupported languages or scripts, and documents with heavy compression artifacts. Users should evaluate outputs on representative documents and use human review for high-impact workflows.
|
| 12 |
+
Verified to have met prescribed NVIDIA quality standards: | Yes
|
| 13 |
+
Performance Metrics: | OCR accuracy, layout/class accuracy, table extraction quality, chart extraction quality, reading-order quality, grounding quality, latency, throughput, and qualitative visual inspection.
|
| 14 |
+
Potential Known Risks: | The model may miss text, hallucinate structure, assign incorrect classes, produce inaccurate bounding boxes, or incorrectly order elements. Downstream systems should treat output as model-generated extraction results and apply validation, confidence checks, and human review where appropriate.
|
| 15 |
+
Licensing: | Governing Terms: Your use of this model is governed by the [OpenMDW License Agreement, version 1.1](https://openmdw.ai/license/1-1/). Use of the tokenizer included in this model is governed by the [MIT license and tokenizer attribution](https://huggingface.co/nvidia/Nemotron-3.5-Parse/blob/main/TOKENIZER_ATTRIBUTION.md). See [third-party notices](https://huggingface.co/nvidia/Nemotron-3.5-Parse/blob/main/THIRD_PARTY_NOTICES.md) for separately installed dependencies.
|
model-cards/nemotron-3-5-parse-subcards/privacy.md
ADDED
|
@@ -0,0 +1,2 @@
|
|
|
|
|
|
|
|
|
|
| 1 |
+
## Privacy
|
| 2 |
+
The datasets used to train this model were reviewed before release. Please review NVIDIA’s Privacy Policy (https://www.nvidia.com/en-us/about-nvidia/privacy-policy/) for more information.
|
model-cards/nemotron-3-5-parse-subcards/safety_security.md
ADDED
|
@@ -0,0 +1,12 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
## Safety
|
| 2 |
+
|
| 3 |
+
Field | Response
|
| 4 |
+
:---- | :----
|
| 5 |
+
Model Application Field(s): | Document intelligence, information retrieval, enterprise automation, data extraction, multimodal data curation, and retrieval-augmented generation workflows.
|
| 6 |
+
Describe the life critical impact (if present). | Not Applicable. This model is not intended to make life-critical, legal, medical, financial, or safety-critical decisions without downstream validation and human oversight.
|
| 7 |
+
Use Case Restrictions: | Governing Terms: Your use of this model is governed by the [OpenMDW License Agreement, version 1.1](https://openmdw.ai/license/1-1/). Use of the tokenizer included in this model is governed by the [MIT license and tokenizer attribution](https://huggingface.co/nvidia/Nemotron-3.5-Parse/blob/main/TOKENIZER_ATTRIBUTION.md). See [third-party notices](https://huggingface.co/nvidia/Nemotron-3.5-Parse/blob/main/THIRD_PARTY_NOTICES.md) for separately installed dependencies.
|
| 8 |
+
Input permissions: | Users are responsible for ensuring they have rights and permissions for input documents and images.
|
| 9 |
+
Model and dataset restrictions: | The Principle of least privilege (PoLP) is applied limiting access for dataset generation and model development. Restrictions enforce dataset access during training, and dataset license constraints adhered to.
|
| 10 |
+
|
| 11 |
+
|
| 12 |
+
**You are responsible for ensuring that your use of NVIDIA-provided models complies with all applicable laws.**
|
model-cards/nemotron-3-5-parse.md
ADDED
|
@@ -0,0 +1,158 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# Nemotron 3.5 Parse Overview
|
| 2 |
+
|
| 3 |
+
## Description:
|
| 4 |
+
Nemotron 3.5 Parse is a document image parser that produces structured, machine-readable text with semantic layout classes, bounding boxes, and reading order. It combines a C-RADIOv4-SO400M vision encoder with a 10-layer mBART decoder that consumes the visual sequence as a causal prefix. The release uses dynamic rectangular NaViT preprocessing, content-normalized coordinates, grouped-query self-attention (16 query heads and 8 key/value heads), a 72,505-token multilingual tokenizer, and chart-aware output through the `<class_Chart>` token. It is intended for document understanding, information retrieval, data extraction, and multimodal data-curation workflows.
|
| 5 |
+
|
| 6 |
+
Nemotron 3.5 Parse was developed by NVIDIA as a part of Nemotron.
|
| 7 |
+
|
| 8 |
+
This model is ready for commercial or non-commercial use.
|
| 9 |
+
|
| 10 |
+
### License/Terms of Use:
|
| 11 |
+
Governing Terms: Your use of this model is governed by the [OpenMDW License Agreement, version 1.1](https://openmdw.ai/license/1-1/). Use of the tokenizer included in this model is governed by the [MIT license and tokenizer attribution](https://huggingface.co/nvidia/Nemotron-3.5-Parse/blob/main/TOKENIZER_ATTRIBUTION.md). See [third-party notices](https://huggingface.co/nvidia/Nemotron-3.5-Parse/blob/main/THIRD_PARTY_NOTICES.md) for separately installed dependencies.
|
| 12 |
+
|
| 13 |
+
### Deployment Geography:
|
| 14 |
+
Global
|
| 15 |
+
|
| 16 |
+
### Use Case:
|
| 17 |
+
Nemotron 3.5 Parse is designed for developers and teams building document intelligence, retrieval-augmented generation (RAG), curator, extractor, and agentic AI applications. It can convert scanned or rendered PDFs, presentation slides, forms, reports, tables, and mixed-content document pages into structured outputs for downstream indexing, retrieval, analytics, model training-data creation, and human-in-the-loop review. The model supports document understanding, optical character recognition (OCR), layout parsing, table and chart extraction, reading-order recovery, and text extraction with spatial grounding.
|
| 18 |
+
|
| 19 |
+
### Release Date:
|
| 20 |
+
Hugging Face 10/17/2026 via https://huggingface.co/nvidia/Nemotron-3.5-Parse
|
| 21 |
+
|
| 22 |
+
## Reference(s):
|
| 23 |
+
* [Hugging Face Transformers mBART documentation](https://huggingface.co/docs/transformers/en/model_doc/mbart)
|
| 24 |
+
* [NVIDIA C-RADIOv4-SO400M](https://huggingface.co/nvidia/C-RADIOv4-SO400M)
|
| 25 |
+
|
| 26 |
+
## Model Architecture:
|
| 27 |
+
**Architecture Type:** Transformer
|
| 28 |
+
|
| 29 |
+
**Network Architecture:**
|
| 30 |
+
* Vision Encoder: NVIDIA C-RADIOv4-SO400M (`vit_so400m_patch16_224`) with dynamic rectangular NaViT preprocessing
|
| 31 |
+
* Adapter Layer: 1D convolutions and normalization layers that compress the vision latent sequence before decoding
|
| 32 |
+
* Decoder: 10-layer mBART causal visual-prefix decoder with 16 query heads and 8 key/value heads
|
| 33 |
+
* Auxiliary Prediction Head: One trained one-step multi-token-prediction (MTP) head is included in `model.safetensors`. Transformers uses the normal decoder for generation; the vLLM plugin enables the bundled head for one-token speculative decoding.
|
| 34 |
+
* Tokenizer: 72,505 entries with multilingual and task/control tokens including `<predict_bbox>`, `<predict_classes>`, `<predict_text_in_pic>`, and `<class_Chart>`; the decoder embedding is padded to 72,512 rows. Use of the tokenizer included in this model is governed by the [MIT license and tokenizer attribution](https://huggingface.co/nvidia/Nemotron-3.5-Parse/blob/main/TOKENIZER_ATTRIBUTION.md).
|
| 35 |
+
**Number of model parameters:** 630,800,496 unique trainable parameters (approximately 0.631B), including the MTP head; shared decoder embedding/output weights are counted once
|
| 36 |
+
|
| 37 |
+
## Input:
|
| 38 |
+
**Input Type(s):** Image, Text
|
| 39 |
+
|
| 40 |
+
**Input Format(s):**
|
| 41 |
+
* Image: Red, Green, Blue (RGB)
|
| 42 |
+
* Text: Prompt string
|
| 43 |
+
|
| 44 |
+
**Input Parameters:**
|
| 45 |
+
* Image: Two-Dimensional (2D)
|
| 46 |
+
* Text: One-Dimensional (1D)
|
| 47 |
+
|
| 48 |
+
**Other Properties Related to Input:**
|
| 49 |
+
* Recommended maximum input resolution (Width, Height): 1664, 2048
|
| 50 |
+
* Channel Count: 3
|
| 51 |
+
* Prompt format: one task prompt composed from the supported control tokens. The default prompt extracts bounding boxes, semantic classes, and text in markdown format: `</s><s><predict_bbox><predict_classes><output_markdown><predict_text_in_pic>`.
|
| 52 |
+
|
| 53 |
+
## Output:
|
| 54 |
+
**Output Type(s):** Text
|
| 55 |
+
|
| 56 |
+
**Output Format(s):** String
|
| 57 |
+
|
| 58 |
+
**Output Parameters:** One-Dimensional (1D)
|
| 59 |
+
|
| 60 |
+
**Other Properties Related to Output:** Nemotron 3.5 Parse returns a string encoding document text, semantic element classes, and content-normalized bounding boxes. Scale x coordinates by the evaluated image width and y coordinates by its height. Postprocessing utilities can convert table or chart-associated text into LaTeX, HTML, markdown, JSON, or CSV where supported.
|
| 61 |
+
|
| 62 |
+
Our AI models are designed and/or optimized to run on NVIDIA GPU-accelerated systems. By leveraging NVIDIA hardware and software frameworks, the model achieves faster training and inference times compared to CPU-only solutions.
|
| 63 |
+
|
| 64 |
+
## Software Integration:
|
| 65 |
+
**Runtime Engine(s):**
|
| 66 |
+
* Transformers
|
| 67 |
+
* vLLM
|
| 68 |
+
|
| 69 |
+
**Reference Hardware:** NVIDIA Hopper H100, BF16.
|
| 70 |
+
|
| 71 |
+
**Supported Operating System(s):**
|
| 72 |
+
* Linux
|
| 73 |
+
|
| 74 |
+
The integration of foundation and fine-tuned models into AI systems requires additional testing using use-case-specific data to ensure safe and effective deployment. Following the V-model methodology, iterative testing and validation at both unit and system levels is essential to mitigate risks, meet technical and functional requirements, and ensure compliance with safety and ethical standards before deployment.
|
| 75 |
+
|
| 76 |
+
## Model Version(s):
|
| 77 |
+
Nemotron 3.5 Parse · October 17, 2026 release
|
| 78 |
+
|
| 79 |
+
This AI model can be embedded as an Application Programming Interface (API) call through the documented Transformers or vLLM serving interfaces.
|
| 80 |
+
|
| 81 |
+
## Training, Testing, and Evaluation Datasets:
|
| 82 |
+
### Training Dataset
|
| 83 |
+
|
| 84 |
+
**Data Modality:**
|
| 85 |
+
* Image
|
| 86 |
+
* Text
|
| 87 |
+
|
| 88 |
+
**Image Training Data Size:**
|
| 89 |
+
* 1 Million to 1 Billion Images
|
| 90 |
+
|
| 91 |
+
**Text Training Data Size:**
|
| 92 |
+
* 1 Billion to 10 Trillion Tokens
|
| 93 |
+
|
| 94 |
+
**Data Collection Method by dataset:**
|
| 95 |
+
* Automated, Manually-Collected, Synthetic
|
| 96 |
+
|
| 97 |
+
**Labeling Method by dataset:** Automated, Manually-Labeled, Synthetic
|
| 98 |
+
**Labeling Details:** Labels are produced through a combination of automated OCR and layout annotation, manual annotation, and synthetic generation. The combination varies by constituent dataset.
|
| 99 |
+
|
| 100 |
+
**Properties:** The training set contains millions of image-text items aggregated across large document, table, and layout datasets. The data consists of document-page and table images paired with OCR text, bounding boxes, and layout labels. Sources include rendered digital documents, scientific papers, PDFs, Wikipedia-style pages, and synthetic document, table, word, and character renderings. Annotations come from OCR and layout models, third-party OCR services, synthetic-generation pipelines, and human labeling.
|
| 101 |
+
|
| 102 |
+
### Testing Dataset
|
| 103 |
+
|
| 104 |
+
**Data Collection Method by dataset:** Automated, Manually-Collected, Synthetic
|
| 105 |
+
**Collection Details:** Sources combine automated collection, manually collected material, and synthetic generation; methods vary by constituent dataset.
|
| 106 |
+
|
| 107 |
+
**Labeling Method by dataset:** Automated, Manually-Labeled, Synthetic
|
| 108 |
+
**Labeling Details:** Labels are produced through a combination of automated OCR and layout annotation, manual annotation, and synthetic generation. The combination varies by constituent dataset.
|
| 109 |
+
|
| 110 |
+
**Properties:** Tested on internal and public document-understanding benchmarks using document-page images and reference text, layout classes, bounding boxes, and structure annotations. Coverage includes OCR quality, layout structure, table parsing, reading order, and visual grounding across document types and languages.
|
| 111 |
+
|
| 112 |
+
### Evaluation Dataset
|
| 113 |
+
|
| 114 |
+
**Data Collection Method by dataset:** Automated, Manually-Collected, Synthetic
|
| 115 |
+
**Collection Details:** Sources combine automated collection, manually collected material, and synthetic generation; methods vary by constituent dataset.
|
| 116 |
+
|
| 117 |
+
**Labeling Method by dataset:** Automated, Manually-Labeled, Synthetic
|
| 118 |
+
**Labeling Details:** Labels are produced through a combination of automated OCR and layout annotation, manual annotation, and synthetic generation. The combination varies by constituent dataset.
|
| 119 |
+
|
| 120 |
+
**Properties:** Evaluated on internal and public document-understanding benchmarks covering OCR quality, layout structure, table parsing, reading order, and visual grounding. The ParseBench reference evaluation covers 2,078 unique document pages. Additional benchmarks cover Indic-language OCR and multilingual synthetic OCR across the scripts listed below.
|
| 121 |
+
|
| 122 |
+
Benchmark coverage:
|
| 123 |
+
* ParseBench evaluates document parsing across text fidelity, semantic formatting, tables, charts, and visual grounding with layout boxes/classes.
|
| 124 |
+
* IndicVisionBench evaluates OCR quality for Indic-language pages, including Bengali, Gujarati, Hindi, Kannada, Malayalam, Marathi, Odia, Punjabi, Tamil, and Telugu.
|
| 125 |
+
* MOSCAR evaluates multilingual synthetic OCR across broad script coverage, including Latin, Arabic, Cyrillic, Han, Hangul, Japanese, Indic scripts, Hebrew, Thai, Greek, and others.
|
| 126 |
+
|
| 127 |
+
## ParseBench
|
| 128 |
+
|
| 129 |
+
| Pipeline | Inference stages | Overall score ↑ | Prepared pages/s ↑ |
|
| 130 |
+
| --- | --- | ---: | ---: |
|
| 131 |
+
| Single | Full-page parse | 72.74 | 15.42 |
|
| 132 |
+
| Double | Full-page parse and predicted Table/Chart crop refinement | 74.92 | 11.65 |
|
| 133 |
+
| Triple | Full-page parse, separate detection, and crop refinement using their region union | 75.20 | 10.62 |
|
| 134 |
+
|
| 135 |
+
Measured on 2,078 unique pages with one H100. Double captures most of the observed
|
| 136 |
+
accuracy gain; triple adds 0.28 points at about 9% lower throughput than double.
|
| 137 |
+
Timing includes inference waits and dependent crop preparation, excluding initial
|
| 138 |
+
preparation and final rendering/writes. Processed-page rates include token-limited
|
| 139 |
+
outputs. The refinement merge retains differing alternatives. See the
|
| 140 |
+
[measurement sources](https://huggingface.co/nvidia/Nemotron-3.5-Parse/blob/main/reference/parsebench/comparisons/README.md).
|
| 141 |
+
|
| 142 |
+
## Inference:
|
| 143 |
+
**Acceleration Engine:** Transformers, vLLM
|
| 144 |
+
**Test Hardware:**
|
| 145 |
+
* NVIDIA Hopper H100
|
| 146 |
+
|
| 147 |
+
## Limitations:
|
| 148 |
+
Performance can vary for low-resolution scans, motion blur, handwriting, decorative fonts, dense tables, unusual layouts, unsupported languages or scripts, and documents with heavy compression artifacts. The model may miss text, hallucinate structure, assign incorrect classes, produce inaccurate bounding boxes, or incorrectly order elements. Users should evaluate outputs on representative target documents and use validation and human review for high-impact workflows.
|
| 149 |
+
|
| 150 |
+
## Ethical Considerations:
|
| 151 |
+
NVIDIA believes Trustworthy AI is a shared responsibility and has established policies and practices to enable development for a wide array of AI applications. Developers should work with their internal model teams to ensure this model meets requirements for the relevant industry and use case and addresses unforeseen product misuse.
|
| 152 |
+
|
| 153 |
+
Please make sure you have proper rights and permissions for all input image and video content. If an input image contains people, personal health information, confidential business data, or intellectual property, the model may extract or reproduce visible text from that content.
|
| 154 |
+
|
| 155 |
+
For more detailed information on ethical considerations for this model, please see the Model Card++ [Bias](nemotron-3-5-parse-subcards/bias.md), [Explainability](nemotron-3-5-parse-subcards/explainability.md), [Safety & Security](nemotron-3-5-parse-subcards/safety_security.md), and [Privacy](nemotron-3-5-parse-subcards/privacy.md) subcards.
|
| 156 |
+
Please report model quality, risk, security vulnerabilities, or NVIDIA AI concerns [here](https://www.nvidia.com/en-us/support/submit-security-vulnerability/).
|
| 157 |
+
|
| 158 |
+
**Generated by NVIDIA Model Card Generator Toolkit.**
|
model.safetensors
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:d7139ecd62f830bac1ce83b77655f10710ee0da31f2810dd47a40eec368917a3
|
| 3 |
+
size 2523266632
|
nemotron_parse/SOURCE_MANIFEST.json
ADDED
|
@@ -0,0 +1,110 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"repository": "nvidia/Nemotron-3.5-Parse",
|
| 3 |
+
"source_revision": "0fad7099302dc11313b7e3157eb8fd6ae633176e",
|
| 4 |
+
"extracted_components": [
|
| 5 |
+
{
|
| 6 |
+
"source": "reference/parsebench/pipeline/engine/__init__.py",
|
| 7 |
+
"source_sha256": "717a1b9838b948a88d2fb5106069db57dce81dc54854cf2be5fb54150c7ac652",
|
| 8 |
+
"destination": "nemotron_parse/engine/__init__.py"
|
| 9 |
+
},
|
| 10 |
+
{
|
| 11 |
+
"source": "reference/parsebench/pipeline/engine/caption_relocation.py",
|
| 12 |
+
"source_sha256": "3a17677b1a43c793cd817f8f3f5745be807203ff75b8e3e0af5a4d51d547dd4b",
|
| 13 |
+
"destination": "nemotron_parse/engine/caption_relocation.py"
|
| 14 |
+
},
|
| 15 |
+
{
|
| 16 |
+
"source": "reference/parsebench/pipeline/engine/chart_caption.py",
|
| 17 |
+
"source_sha256": "fe82e0eca23be96d7cc2199657432bd2ddf464313d41fed3984bf5960fddb0e6",
|
| 18 |
+
"destination": "nemotron_parse/engine/chart_caption.py"
|
| 19 |
+
},
|
| 20 |
+
{
|
| 21 |
+
"source": "reference/parsebench/pipeline/engine/chart_context.py",
|
| 22 |
+
"source_sha256": "ffe3393bc13bd54b9b65235ca51a4c15a0b40fe6404a5897379eb0dd2748d868",
|
| 23 |
+
"destination": "nemotron_parse/engine/chart_context.py"
|
| 24 |
+
},
|
| 25 |
+
{
|
| 26 |
+
"source": "reference/parsebench/pipeline/engine/convert/__init__.py",
|
| 27 |
+
"source_sha256": "717a1b9838b948a88d2fb5106069db57dce81dc54854cf2be5fb54150c7ac652",
|
| 28 |
+
"destination": "nemotron_parse/engine/convert/__init__.py"
|
| 29 |
+
},
|
| 30 |
+
{
|
| 31 |
+
"source": "reference/parsebench/pipeline/engine/convert/hybrid_cleanup.py",
|
| 32 |
+
"source_sha256": "f2f319de9bdabcfc5e264cdef6ff9d13235c488a372e9f58f676017f28f7c5c0",
|
| 33 |
+
"destination": "nemotron_parse/engine/convert/hybrid_cleanup.py"
|
| 34 |
+
},
|
| 35 |
+
{
|
| 36 |
+
"source": "reference/parsebench/pipeline/engine/convert/hybrid_inline.py",
|
| 37 |
+
"source_sha256": "0296f7827e74ae3ab64e9c9efbd496f096f0a3bbde7c3ad64d017539f1db7cb6",
|
| 38 |
+
"destination": "nemotron_parse/engine/convert/hybrid_inline.py"
|
| 39 |
+
},
|
| 40 |
+
{
|
| 41 |
+
"source": "reference/parsebench/pipeline/engine/convert/latex2html.py",
|
| 42 |
+
"source_sha256": "d583b5fa7fdbe181b190d2c8c22034318a51ef25d47203f9c7504bf35b508a72",
|
| 43 |
+
"destination": "nemotron_parse/engine/convert/latex2html.py"
|
| 44 |
+
},
|
| 45 |
+
{
|
| 46 |
+
"source": "reference/parsebench/pipeline/engine/convert/rich_table_fallback.py",
|
| 47 |
+
"source_sha256": "c2abbc602b293bf14a499f680de9f7d1584fdcf1438a374e3c75cf3705c0b793",
|
| 48 |
+
"destination": "nemotron_parse/engine/convert/rich_table_fallback.py"
|
| 49 |
+
},
|
| 50 |
+
{
|
| 51 |
+
"source": "reference/parsebench/pipeline/engine/convert/rich_table_reader.py",
|
| 52 |
+
"source_sha256": "8a8b615681d3adcb46c5290ad631abffeccb2af0ee18f25e82f0cdb70f5befac",
|
| 53 |
+
"destination": "nemotron_parse/engine/convert/rich_table_reader.py"
|
| 54 |
+
},
|
| 55 |
+
{
|
| 56 |
+
"source": "reference/parsebench/pipeline/engine/convert/semantic_styles.py",
|
| 57 |
+
"source_sha256": "3d92146e6c713b9de67cb71e92fe254a1344d9054c3925a85cb778ec68c0d075",
|
| 58 |
+
"destination": "nemotron_parse/engine/convert/semantic_styles.py"
|
| 59 |
+
},
|
| 60 |
+
{
|
| 61 |
+
"source": "reference/parsebench/pipeline/engine/geometry.py",
|
| 62 |
+
"source_sha256": "76671df686b574a827205acb757eecf702e970470ffa15f81b3129cbcd95e242",
|
| 63 |
+
"destination": "nemotron_parse/engine/geometry.py"
|
| 64 |
+
},
|
| 65 |
+
{
|
| 66 |
+
"source": "reference/parsebench/pipeline/engine/header_inference_probe.py",
|
| 67 |
+
"source_sha256": "02950f1aec5e2dc88a96487bd9124c97e5ce2842fbcfc5264ac2d5f697865722",
|
| 68 |
+
"destination": "nemotron_parse/engine/header_inference_probe.py"
|
| 69 |
+
},
|
| 70 |
+
{
|
| 71 |
+
"source": "reference/parsebench/pipeline/engine/markdown.py",
|
| 72 |
+
"source_sha256": "fc57566255ac632b24825973d9533a13fd717dad4e38d096098bac33ecc79957",
|
| 73 |
+
"destination": "nemotron_parse/engine/markdown.py"
|
| 74 |
+
},
|
| 75 |
+
{
|
| 76 |
+
"source": "reference/parsebench/pipeline/engine/merge.py",
|
| 77 |
+
"source_sha256": "b1c2e749e8f772b716c5a8a2ebeba7348474eb9368b00d37f1c8987e26d24e4e",
|
| 78 |
+
"destination": "nemotron_parse/engine/merge.py"
|
| 79 |
+
},
|
| 80 |
+
{
|
| 81 |
+
"source": "reference/parsebench/pipeline/engine/policies.py",
|
| 82 |
+
"source_sha256": "4c1f78be4d8061e4224fb59c3dc5342dfc21d9d37826814ea3573f339777aebd",
|
| 83 |
+
"destination": "nemotron_parse/engine/policies.py"
|
| 84 |
+
},
|
| 85 |
+
{
|
| 86 |
+
"source": "reference/parsebench/pipeline/engine/regions.py",
|
| 87 |
+
"source_sha256": "33058eec02bb378050febc577a2d85de97d34083234a08a2a8fa9c4b8dc62239",
|
| 88 |
+
"destination": "nemotron_parse/engine/regions.py"
|
| 89 |
+
},
|
| 90 |
+
{
|
| 91 |
+
"source": "reference/parsebench/pipeline/engine/text_cleanup.py",
|
| 92 |
+
"source_sha256": "f02f472d7b3cb08eca5b11d3ecf80ee6ec0636857ade17fb4cd11c13f371c2e4",
|
| 93 |
+
"destination": "nemotron_parse/engine/text_cleanup.py"
|
| 94 |
+
}
|
| 95 |
+
],
|
| 96 |
+
"adaptations": [
|
| 97 |
+
"No dataset, scorer, or reference-answer dependencies",
|
| 98 |
+
"Headerless flat exports preserve all rows; collision-free field names",
|
| 99 |
+
"Exact duplicate cleanup preserves whitespace differences",
|
| 100 |
+
"Malformed table fallback preserves source",
|
| 101 |
+
"Presentation heuristics configurable",
|
| 102 |
+
"Image-based inference wrappers and guarded crop merging",
|
| 103 |
+
"Code and Formula blocks retain literal table examples",
|
| 104 |
+
"Enable guarded chart presentation rules by default, retaining opt-outs",
|
| 105 |
+
"Prepared image batches, bounded pipelined concurrency, and staged comparisons",
|
| 106 |
+
"Explicit partial-output accounting and per-stage timing",
|
| 107 |
+
"Retain differing overlapping structural alternatives by default; guarded replacement remains an explicit option. Scheduler unchanged.",
|
| 108 |
+
"Add first-request-to-last-response inference wall timing, retaining dependent crop preparation gaps and separate total pipeline metrics."
|
| 109 |
+
]
|
| 110 |
+
}
|
nemotron_parse/__init__.py
ADDED
|
@@ -0,0 +1,6 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
| 2 |
+
# SPDX-License-Identifier: OpenMDW-1.1
|
| 3 |
+
"""General document postprocessing and optional multipass inference."""
|
| 4 |
+
from .postprocessing import PostprocessOptions, postprocess_document, render_markdown
|
| 5 |
+
|
| 6 |
+
__all__ = ["PostprocessOptions", "postprocess_document", "render_markdown"]
|
nemotron_parse/batch.py
ADDED
|
@@ -0,0 +1,571 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
| 2 |
+
# SPDX-License-Identifier: OpenMDW-1.1
|
| 3 |
+
"""Prepared image batches with bounded, overlapping HTTP concurrency and CPU rendering."""
|
| 4 |
+
from concurrent.futures import (
|
| 5 |
+
ThreadPoolExecutor,
|
| 6 |
+
ProcessPoolExecutor,
|
| 7 |
+
wait,
|
| 8 |
+
FIRST_COMPLETED,
|
| 9 |
+
)
|
| 10 |
+
from dataclasses import dataclass, asdict
|
| 11 |
+
from io import BytesIO
|
| 12 |
+
from pathlib import Path
|
| 13 |
+
from types import SimpleNamespace
|
| 14 |
+
import base64
|
| 15 |
+
import collections
|
| 16 |
+
import multiprocessing
|
| 17 |
+
import time
|
| 18 |
+
|
| 19 |
+
from PIL import Image
|
| 20 |
+
from .engine import geometry
|
| 21 |
+
from .engine.regions import select_union_regions
|
| 22 |
+
from .postprocessing import clean_raw_output, postprocess_document
|
| 23 |
+
|
| 24 |
+
|
| 25 |
+
@dataclass
|
| 26 |
+
class PreparedPage:
|
| 27 |
+
image: Image.Image
|
| 28 |
+
data_url: str
|
| 29 |
+
source: str | None
|
| 30 |
+
|
| 31 |
+
|
| 32 |
+
@dataclass
|
| 33 |
+
class PreparedBatch:
|
| 34 |
+
pages: list[PreparedPage]
|
| 35 |
+
preprocessing_seconds: float
|
| 36 |
+
|
| 37 |
+
|
| 38 |
+
@dataclass
|
| 39 |
+
class BatchResult:
|
| 40 |
+
documents: list[dict]
|
| 41 |
+
metrics: dict
|
| 42 |
+
|
| 43 |
+
|
| 44 |
+
def _url(data, mime):
|
| 45 |
+
return f"data:{mime};base64," + base64.b64encode(data).decode("ascii")
|
| 46 |
+
|
| 47 |
+
|
| 48 |
+
def prepare_images(images, *, workers=8):
|
| 49 |
+
if workers < 1:
|
| 50 |
+
raise ValueError("workers must be positive")
|
| 51 |
+
start = time.perf_counter()
|
| 52 |
+
|
| 53 |
+
def one(value):
|
| 54 |
+
if isinstance(value, Image.Image):
|
| 55 |
+
image = value.convert("RGB")
|
| 56 |
+
stream = BytesIO()
|
| 57 |
+
image.save(stream, format="PNG")
|
| 58 |
+
return PreparedPage(image, _url(stream.getvalue(), "image/png"), None)
|
| 59 |
+
data = Path(value).read_bytes()
|
| 60 |
+
with Image.open(BytesIO(data)) as opened:
|
| 61 |
+
mime = {"JPEG": "image/jpeg", "PNG": "image/png", "WEBP": "image/webp"}.get(
|
| 62 |
+
opened.format
|
| 63 |
+
)
|
| 64 |
+
keep_encoded = mime is not None and opened.mode == "RGB"
|
| 65 |
+
image = opened.convert("RGB")
|
| 66 |
+
if not keep_encoded:
|
| 67 |
+
stream = BytesIO()
|
| 68 |
+
image.save(stream, format="PNG")
|
| 69 |
+
data, mime = stream.getvalue(), "image/png"
|
| 70 |
+
return PreparedPage(image, _url(data, mime), str(value))
|
| 71 |
+
|
| 72 |
+
with ThreadPoolExecutor(max_workers=workers) as pool:
|
| 73 |
+
pages = list(pool.map(one, images))
|
| 74 |
+
if not pages:
|
| 75 |
+
raise ValueError("No page images provided")
|
| 76 |
+
return PreparedBatch(pages, time.perf_counter() - start)
|
| 77 |
+
|
| 78 |
+
|
| 79 |
+
def _finish(task):
|
| 80 |
+
raw, crops, size, options, merge_policy = task
|
| 81 |
+
from .inference import merge_crops
|
| 82 |
+
|
| 83 |
+
raw = clean_raw_output(raw, deduplicate=options.deduplicate)
|
| 84 |
+
if crops:
|
| 85 |
+
raw = merge_crops(raw, crops, policy=merge_policy)
|
| 86 |
+
return postprocess_document(raw, image_size=size, options=options)
|
| 87 |
+
|
| 88 |
+
|
| 89 |
+
def _crop_preparer(parser, pages, triple_regions):
|
| 90 |
+
args = SimpleNamespace(
|
| 91 |
+
include_detector_labels="table,chart",
|
| 92 |
+
detector_order="model",
|
| 93 |
+
detector_score_min=0.0,
|
| 94 |
+
max_elements_per_page=parser.max_regions,
|
| 95 |
+
model_coordinate_frame="content_normalized",
|
| 96 |
+
canvas_width=1664.0,
|
| 97 |
+
canvas_height=2048.0,
|
| 98 |
+
pad_position="top-left",
|
| 99 |
+
)
|
| 100 |
+
|
| 101 |
+
def page_crops(index, direct_response, detection_response=None, start_region=0):
|
| 102 |
+
page = pages[index]
|
| 103 |
+
page_spec = geometry.PageImage(
|
| 104 |
+
"page", "image.png", Path("image.png"), *page.image.size
|
| 105 |
+
)
|
| 106 |
+
|
| 107 |
+
def boxes(response):
|
| 108 |
+
return geometry.parse_model_detection_predictions(
|
| 109 |
+
args,
|
| 110 |
+
[page_spec],
|
| 111 |
+
[{"key": "page", "raw_out": clean_raw_output(response["raw_out"])}],
|
| 112 |
+
)
|
| 113 |
+
|
| 114 |
+
selected = (
|
| 115 |
+
boxes(direct_response)
|
| 116 |
+
if direct_response is not None
|
| 117 |
+
else [{"key": "page", "detections": []}]
|
| 118 |
+
)
|
| 119 |
+
if detection_response is not None:
|
| 120 |
+
detected = boxes(detection_response)
|
| 121 |
+
selected = (
|
| 122 |
+
select_union_regions(detected, selected, max_regions=parser.max_regions)
|
| 123 |
+
if triple_regions == "union"
|
| 124 |
+
else detected
|
| 125 |
+
)
|
| 126 |
+
boxes_for_page = selected[0]["detections"]
|
| 127 |
+
items, metadata = [], []
|
| 128 |
+
for region_index, region in enumerate(
|
| 129 |
+
boxes_for_page[start_region:], start_region
|
| 130 |
+
):
|
| 131 |
+
pixels = geometry.crop_bbox_pixels(
|
| 132 |
+
tuple(region["bbox_norm"]),
|
| 133 |
+
*page.image.size,
|
| 134 |
+
margin_px=parser.crop_margin_px,
|
| 135 |
+
margin_ratio=parser.crop_margin_ratio,
|
| 136 |
+
)
|
| 137 |
+
if pixels[2] - pixels[0] < 8 or pixels[3] - pixels[1] < 8:
|
| 138 |
+
continue
|
| 139 |
+
stream = BytesIO()
|
| 140 |
+
page.image.crop(pixels).save(stream, format="JPEG", quality=95)
|
| 141 |
+
data = stream.getvalue()
|
| 142 |
+
# HTTP sends the prepared crop directly; custom backends receive decoded pixels.
|
| 143 |
+
image = None
|
| 144 |
+
if parser.backend is not None:
|
| 145 |
+
with Image.open(BytesIO(data)) as decoded:
|
| 146 |
+
image = decoded.convert("RGB")
|
| 147 |
+
items.append((image, _url(data, "image/jpeg")))
|
| 148 |
+
metadata.append(
|
| 149 |
+
{
|
| 150 |
+
"page_index": index,
|
| 151 |
+
"crop_id": f"page__{region_index:03d}",
|
| 152 |
+
"page_key": "page",
|
| 153 |
+
"detector_label": region["label"],
|
| 154 |
+
"detector_score": region["score"],
|
| 155 |
+
"detector_bbox_norm": region["bbox_norm"],
|
| 156 |
+
"crop_bbox_pixels": pixels,
|
| 157 |
+
"page_size": page.image.size,
|
| 158 |
+
}
|
| 159 |
+
)
|
| 160 |
+
return index, boxes_for_page, items, metadata
|
| 161 |
+
|
| 162 |
+
return page_crops
|
| 163 |
+
|
| 164 |
+
|
| 165 |
+
def _inference_timing(spans):
|
| 166 |
+
"""Wall time across model calls, including gaps needed to prepare later calls."""
|
| 167 |
+
elapsed = max(end for begin, end in spans.values()) - min(
|
| 168 |
+
begin for begin, end in spans.values()
|
| 169 |
+
)
|
| 170 |
+
return {
|
| 171 |
+
"inference_seconds": elapsed,
|
| 172 |
+
"inference_timing_includes": [
|
| 173 |
+
"all HTTP/model calls, server queueing and response transfer",
|
| 174 |
+
"gaps between inference calls, including dependent crop preparation",
|
| 175 |
+
],
|
| 176 |
+
"inference_timing_excludes": [
|
| 177 |
+
"initial image preparation",
|
| 178 |
+
"server startup",
|
| 179 |
+
"warmup requests",
|
| 180 |
+
"merge/render work after the final model response",
|
| 181 |
+
"output serialization",
|
| 182 |
+
],
|
| 183 |
+
"inference_timing_definition": "first model request start to last model response",
|
| 184 |
+
}
|
| 185 |
+
|
| 186 |
+
|
| 187 |
+
def _pipeline(parser, pages, mode, concurrency, workers, region_policy, allow_partial):
|
| 188 |
+
"""Keep one bounded HTTP window full; dependent crops take the next free slot.
|
| 189 |
+
|
| 190 |
+
Detection regions form the stable prefix of the union. They can be submitted
|
| 191 |
+
before full-page parsing finishes; parser-only additions follow in union order.
|
| 192 |
+
CPU crop encoding and document rendering never occupy an HTTP worker.
|
| 193 |
+
"""
|
| 194 |
+
from .inference import GenerationError
|
| 195 |
+
|
| 196 |
+
count = len(pages)
|
| 197 |
+
direct = [None] * count
|
| 198 |
+
detection = [None] * count if mode == "triple" else None
|
| 199 |
+
regions = [[] for _ in pages]
|
| 200 |
+
crops = [{} for _ in pages]
|
| 201 |
+
documents = [None] * count
|
| 202 |
+
prep_pending = [0] * count
|
| 203 |
+
crop_pending = [0] * count
|
| 204 |
+
sealed = [mode == "single"] * count
|
| 205 |
+
render_started = set()
|
| 206 |
+
ready = collections.deque()
|
| 207 |
+
initial = collections.deque(
|
| 208 |
+
(i, stage, (page.image, page.data_url), None)
|
| 209 |
+
for i, page in enumerate(pages)
|
| 210 |
+
for stage in (["fullpage", "detection"] if mode == "triple" else ["fullpage"])
|
| 211 |
+
)
|
| 212 |
+
futures = {}
|
| 213 |
+
http_active = 0
|
| 214 |
+
peak = 0
|
| 215 |
+
spans = {}
|
| 216 |
+
service = collections.Counter()
|
| 217 |
+
started = time.perf_counter()
|
| 218 |
+
page_crops = _crop_preparer(parser, pages, region_policy)
|
| 219 |
+
|
| 220 |
+
def generate(stage, item):
|
| 221 |
+
begin = time.perf_counter()
|
| 222 |
+
try:
|
| 223 |
+
result = parser._generate(item[0], stage, data_url=item[1])
|
| 224 |
+
except GenerationError as exc:
|
| 225 |
+
if not (
|
| 226 |
+
allow_partial
|
| 227 |
+
and exc.response.get("finish_reason") == "length"
|
| 228 |
+
and exc.response.get("raw_out")
|
| 229 |
+
):
|
| 230 |
+
raise
|
| 231 |
+
result = dict(exc.response)
|
| 232 |
+
end = time.perf_counter()
|
| 233 |
+
return dict(result, latency_s=end - begin), begin, end
|
| 234 |
+
|
| 235 |
+
# Small batches avoid process startup; threaded CPU work still cannot block HTTP.
|
| 236 |
+
render_executor = (
|
| 237 |
+
ProcessPoolExecutor(
|
| 238 |
+
max_workers=workers, mp_context=multiprocessing.get_context("spawn")
|
| 239 |
+
)
|
| 240 |
+
if workers > 1 and count >= 16
|
| 241 |
+
else ThreadPoolExecutor(max_workers=1)
|
| 242 |
+
)
|
| 243 |
+
with ThreadPoolExecutor(max_workers=concurrency) as http, ThreadPoolExecutor(
|
| 244 |
+
max_workers=min(8, count)
|
| 245 |
+
) as prep, render_executor as render:
|
| 246 |
+
|
| 247 |
+
def submit_prep(i, first, second, offset=0):
|
| 248 |
+
prep_pending[i] += 1
|
| 249 |
+
futures[prep.submit(page_crops, i, first, second, offset)] = (
|
| 250 |
+
"prep",
|
| 251 |
+
i,
|
| 252 |
+
None,
|
| 253 |
+
)
|
| 254 |
+
|
| 255 |
+
def maybe_render(i):
|
| 256 |
+
if (
|
| 257 |
+
i in render_started
|
| 258 |
+
or direct[i] is None
|
| 259 |
+
or not sealed[i]
|
| 260 |
+
or prep_pending[i]
|
| 261 |
+
or crop_pending[i]
|
| 262 |
+
):
|
| 263 |
+
return
|
| 264 |
+
render_started.add(i)
|
| 265 |
+
rows = [crops[i][j] for j in sorted(crops[i])]
|
| 266 |
+
task = (
|
| 267 |
+
direct[i]["raw_out"],
|
| 268 |
+
[r for r in rows if r.get("finish_reason") != "length"],
|
| 269 |
+
pages[i].image.size,
|
| 270 |
+
parser.options,
|
| 271 |
+
parser.merge_policy,
|
| 272 |
+
)
|
| 273 |
+
futures[render.submit(_finish, task)] = ("render", i, None)
|
| 274 |
+
|
| 275 |
+
while initial or ready or futures:
|
| 276 |
+
while http_active < concurrency and (ready or initial):
|
| 277 |
+
i, stage, item, metadata = (
|
| 278 |
+
ready.popleft() if ready else initial.popleft()
|
| 279 |
+
)
|
| 280 |
+
futures[http.submit(generate, stage, item)] = (stage, i, metadata)
|
| 281 |
+
http_active += 1
|
| 282 |
+
peak = max(peak, http_active)
|
| 283 |
+
if not futures:
|
| 284 |
+
break
|
| 285 |
+
done, _ = wait(futures, return_when=FIRST_COMPLETED)
|
| 286 |
+
for future in done:
|
| 287 |
+
kind, i, metadata = futures.pop(future)
|
| 288 |
+
value = future.result()
|
| 289 |
+
if kind == "prep":
|
| 290 |
+
prep_pending[i] -= 1
|
| 291 |
+
_, selected, items, rows = value
|
| 292 |
+
# Two prep jobs can finish out of order; retain the full union.
|
| 293 |
+
if len(selected) > len(regions[i]):
|
| 294 |
+
regions[i] = selected
|
| 295 |
+
for item, row in zip(items, rows):
|
| 296 |
+
ready.append((i, "crop", item, row))
|
| 297 |
+
crop_pending[i] += 1
|
| 298 |
+
elif kind == "render":
|
| 299 |
+
documents[i] = value
|
| 300 |
+
else:
|
| 301 |
+
http_active -= 1
|
| 302 |
+
response, begin, end = value
|
| 303 |
+
span = spans.setdefault(kind, [begin, end])
|
| 304 |
+
span[0], span[1] = min(span[0], begin), max(span[1], end)
|
| 305 |
+
service[kind] += end - begin
|
| 306 |
+
if kind == "crop":
|
| 307 |
+
crop_pending[i] -= 1
|
| 308 |
+
crops[i][metadata["crop_id"]] = {**metadata, **response}
|
| 309 |
+
elif kind == "fullpage":
|
| 310 |
+
direct[i] = response
|
| 311 |
+
if mode == "double":
|
| 312 |
+
sealed[i] = True
|
| 313 |
+
submit_prep(i, response, None)
|
| 314 |
+
elif detection is not None and detection[i] is not None:
|
| 315 |
+
sealed[i] = True
|
| 316 |
+
if region_policy == "union":
|
| 317 |
+
# The detector prefix is already queued. Only prepare additions.
|
| 318 |
+
detector_count = len(
|
| 319 |
+
page_crops(
|
| 320 |
+
i, None, detection[i], parser.max_regions
|
| 321 |
+
)[1]
|
| 322 |
+
)
|
| 323 |
+
submit_prep(i, response, detection[i], detector_count)
|
| 324 |
+
else:
|
| 325 |
+
detection[i] = response
|
| 326 |
+
submit_prep(i, None, response)
|
| 327 |
+
if region_policy == "detector":
|
| 328 |
+
sealed[i] = True
|
| 329 |
+
elif direct[i] is not None:
|
| 330 |
+
sealed[i] = True
|
| 331 |
+
detector_count = len(
|
| 332 |
+
page_crops(i, None, response, parser.max_regions)[1]
|
| 333 |
+
)
|
| 334 |
+
submit_prep(i, direct[i], response, detector_count)
|
| 335 |
+
maybe_render(i)
|
| 336 |
+
assert all(document is not None for document in documents)
|
| 337 |
+
return (
|
| 338 |
+
direct,
|
| 339 |
+
detection,
|
| 340 |
+
regions,
|
| 341 |
+
[[rows[j] for j in sorted(rows)] for rows in crops],
|
| 342 |
+
documents,
|
| 343 |
+
{
|
| 344 |
+
**_inference_timing(spans),
|
| 345 |
+
"peak_inflight_requests": peak,
|
| 346 |
+
"stage_times": {
|
| 347 |
+
stage + "_span_seconds": end - begin
|
| 348 |
+
for stage, (begin, end) in spans.items()
|
| 349 |
+
},
|
| 350 |
+
"request_service_seconds": dict(service),
|
| 351 |
+
"stage_first_request_seconds": {
|
| 352 |
+
stage: span[0] - started for stage, span in spans.items()
|
| 353 |
+
},
|
| 354 |
+
"overlapping_stage_times": True,
|
| 355 |
+
},
|
| 356 |
+
)
|
| 357 |
+
|
| 358 |
+
|
| 359 |
+
def parse_prepared(
|
| 360 |
+
parser,
|
| 361 |
+
prepared,
|
| 362 |
+
*,
|
| 363 |
+
mode="single",
|
| 364 |
+
concurrency=128,
|
| 365 |
+
postprocess_workers=8,
|
| 366 |
+
triple_regions="union",
|
| 367 |
+
allow_partial=False,
|
| 368 |
+
schedule="pipelined",
|
| 369 |
+
):
|
| 370 |
+
from .inference import GenerationError
|
| 371 |
+
|
| 372 |
+
if mode not in ("single", "double", "triple") or triple_regions not in (
|
| 373 |
+
"detector",
|
| 374 |
+
"union",
|
| 375 |
+
):
|
| 376 |
+
raise ValueError("Invalid pipeline mode or region policy")
|
| 377 |
+
if concurrency < 1 or postprocess_workers < 1:
|
| 378 |
+
raise ValueError("Worker counts must be positive")
|
| 379 |
+
if not isinstance(prepared, PreparedBatch) or not prepared.pages:
|
| 380 |
+
raise ValueError("Use parser.prepare() to prepare a nonempty image batch")
|
| 381 |
+
if schedule not in ("staged", "pipelined"):
|
| 382 |
+
raise ValueError("schedule must be staged or pipelined")
|
| 383 |
+
pages = prepared.pages
|
| 384 |
+
start = time.perf_counter()
|
| 385 |
+
stage_times = {}
|
| 386 |
+
request_spans = {}
|
| 387 |
+
|
| 388 |
+
def requests(items, stage):
|
| 389 |
+
stage_start = time.perf_counter()
|
| 390 |
+
|
| 391 |
+
def one(item):
|
| 392 |
+
image, url = item
|
| 393 |
+
request_start = time.perf_counter()
|
| 394 |
+
try:
|
| 395 |
+
response = parser._generate(image, stage, data_url=url)
|
| 396 |
+
except GenerationError as exc:
|
| 397 |
+
# Opt-in partial output is explicit and never counted as complete.
|
| 398 |
+
if (
|
| 399 |
+
not allow_partial
|
| 400 |
+
or exc.response.get("finish_reason") != "length"
|
| 401 |
+
or not exc.response.get("raw_out")
|
| 402 |
+
):
|
| 403 |
+
raise
|
| 404 |
+
response = dict(exc.response)
|
| 405 |
+
request_end = time.perf_counter()
|
| 406 |
+
response = dict(response, latency_s=request_end - request_start)
|
| 407 |
+
return response, request_start, request_end
|
| 408 |
+
|
| 409 |
+
with ThreadPoolExecutor(max_workers=concurrency) as pool:
|
| 410 |
+
result = list(pool.map(one, items))
|
| 411 |
+
stage_times[stage + "_seconds"] = time.perf_counter() - stage_start
|
| 412 |
+
request_spans[stage] = (
|
| 413 |
+
min(begin for response, begin, end in result),
|
| 414 |
+
max(end for response, begin, end in result),
|
| 415 |
+
)
|
| 416 |
+
return [response for response, begin, end in result]
|
| 417 |
+
|
| 418 |
+
if schedule == "pipelined":
|
| 419 |
+
direct, detection, regions, crops_by_page, documents, stream_metrics = (
|
| 420 |
+
_pipeline(
|
| 421 |
+
parser,
|
| 422 |
+
pages,
|
| 423 |
+
mode,
|
| 424 |
+
concurrency,
|
| 425 |
+
postprocess_workers,
|
| 426 |
+
triple_regions,
|
| 427 |
+
allow_partial,
|
| 428 |
+
)
|
| 429 |
+
)
|
| 430 |
+
stage_times.update(stream_metrics.pop("stage_times"))
|
| 431 |
+
crop_responses = [r for crops in crops_by_page for r in crops]
|
| 432 |
+
else:
|
| 433 |
+
direct = requests([(p.image, p.data_url) for p in pages], "fullpage")
|
| 434 |
+
detection = (
|
| 435 |
+
requests([(p.image, p.data_url) for p in pages], "detection")
|
| 436 |
+
if mode == "triple"
|
| 437 |
+
else None
|
| 438 |
+
)
|
| 439 |
+
regions = [[] for _ in pages]
|
| 440 |
+
crop_metadata = []
|
| 441 |
+
crop_items = []
|
| 442 |
+
crop_responses = []
|
| 443 |
+
if mode != "single":
|
| 444 |
+
prep_start = time.perf_counter()
|
| 445 |
+
page_crops = _crop_preparer(parser, pages, triple_regions)
|
| 446 |
+
with ThreadPoolExecutor(max_workers=min(8, len(pages))) as pool:
|
| 447 |
+
for index, selected, items, metadata in pool.map(
|
| 448 |
+
lambda i: page_crops(
|
| 449 |
+
i, direct[i], detection[i] if detection is not None else None
|
| 450 |
+
),
|
| 451 |
+
range(len(pages)),
|
| 452 |
+
):
|
| 453 |
+
regions[index] = selected
|
| 454 |
+
crop_items.extend(items)
|
| 455 |
+
crop_metadata.extend(metadata)
|
| 456 |
+
stage_times["crop_preprocessing_seconds"] = time.perf_counter() - prep_start
|
| 457 |
+
crop_responses = requests(crop_items, "crop") if crop_items else []
|
| 458 |
+
crops_by_page = [[] for _ in pages]
|
| 459 |
+
for metadata, response in zip(crop_metadata, crop_responses):
|
| 460 |
+
crops_by_page[metadata["page_index"]].append({**metadata, **response})
|
| 461 |
+
post_start = time.perf_counter()
|
| 462 |
+
tasks = [
|
| 463 |
+
(
|
| 464 |
+
response["raw_out"],
|
| 465 |
+
[r for r in crops if r.get("finish_reason") != "length"],
|
| 466 |
+
page.image.size,
|
| 467 |
+
parser.options,
|
| 468 |
+
parser.merge_policy,
|
| 469 |
+
)
|
| 470 |
+
for page, response, crops in zip(pages, direct, crops_by_page)
|
| 471 |
+
]
|
| 472 |
+
if postprocess_workers == 1 or len(tasks) < 16:
|
| 473 |
+
documents = list(map(_finish, tasks))
|
| 474 |
+
else:
|
| 475 |
+
with ProcessPoolExecutor(
|
| 476 |
+
max_workers=postprocess_workers,
|
| 477 |
+
mp_context=multiprocessing.get_context("spawn"),
|
| 478 |
+
) as pool:
|
| 479 |
+
documents = list(pool.map(_finish, tasks, chunksize=8))
|
| 480 |
+
stage_times["postprocessing_seconds"] = time.perf_counter() - post_start
|
| 481 |
+
stream_metrics = _inference_timing(request_spans)
|
| 482 |
+
capped = collections.Counter()
|
| 483 |
+
completed = 0
|
| 484 |
+
for index, document in enumerate(documents):
|
| 485 |
+
stages = {"fullpage": direct[index], "crops": crops_by_page[index]}
|
| 486 |
+
if detection is not None:
|
| 487 |
+
stages["detection"] = detection[index]
|
| 488 |
+
errors = []
|
| 489 |
+
for stage in ["fullpage", "detection"]:
|
| 490 |
+
if stage in stages and stages[stage].get("finish_reason") == "length":
|
| 491 |
+
capped[stage] += 1
|
| 492 |
+
errors.append(stage + " reached its token limit")
|
| 493 |
+
for crop in stages["crops"]:
|
| 494 |
+
if crop.get("finish_reason") == "length":
|
| 495 |
+
capped["crop"] += 1
|
| 496 |
+
errors.append(
|
| 497 |
+
"crop "
|
| 498 |
+
+ crop["crop_id"]
|
| 499 |
+
+ " reached its token limit; refinement omitted"
|
| 500 |
+
)
|
| 501 |
+
document["complete"] = not errors
|
| 502 |
+
document["generation_warnings"] = errors
|
| 503 |
+
document["stages"] = stages
|
| 504 |
+
document["source"] = pages[index].source
|
| 505 |
+
document["inference"] = {
|
| 506 |
+
"mode": mode,
|
| 507 |
+
"merge_policy": parser.merge_policy,
|
| 508 |
+
"triple_regions": triple_regions if mode == "triple" else None,
|
| 509 |
+
"model": parser.model,
|
| 510 |
+
"coordinate_frame": "content_normalized",
|
| 511 |
+
"regions": regions[index],
|
| 512 |
+
"sampling": {
|
| 513 |
+
name: asdict(value) for name, value in parser.settings.items()
|
| 514 |
+
},
|
| 515 |
+
"postprocessing": asdict(parser.options),
|
| 516 |
+
"crop_jpeg_quality": 95,
|
| 517 |
+
"crop_margin_ratio": parser.crop_margin_ratio,
|
| 518 |
+
"crop_margin_px": parser.crop_margin_px,
|
| 519 |
+
}
|
| 520 |
+
completed += not errors
|
| 521 |
+
elapsed = time.perf_counter() - start
|
| 522 |
+
all_responses = direct + (detection or []) + crop_responses
|
| 523 |
+
metrics = {
|
| 524 |
+
"mode": mode,
|
| 525 |
+
"merge_policy": parser.merge_policy,
|
| 526 |
+
"schedule": schedule,
|
| 527 |
+
**stream_metrics,
|
| 528 |
+
"pages": len(pages),
|
| 529 |
+
"complete_pages": completed,
|
| 530 |
+
"incomplete_pages": len(pages) - completed,
|
| 531 |
+
"concurrency": concurrency,
|
| 532 |
+
"postprocess_workers": postprocess_workers,
|
| 533 |
+
"preprocessing_seconds": prepared.preprocessing_seconds,
|
| 534 |
+
"inference_pages_per_second": (
|
| 535 |
+
len(pages) / stream_metrics["inference_seconds"]
|
| 536 |
+
if stream_metrics["inference_seconds"] > 0
|
| 537 |
+
else None
|
| 538 |
+
),
|
| 539 |
+
"complete_inference_pages_per_second": (
|
| 540 |
+
completed / stream_metrics["inference_seconds"]
|
| 541 |
+
if stream_metrics["inference_seconds"] > 0
|
| 542 |
+
else None
|
| 543 |
+
),
|
| 544 |
+
"pipeline_seconds": elapsed,
|
| 545 |
+
"processed_pages_per_second": len(pages) / elapsed,
|
| 546 |
+
"complete_pages_per_second": completed / elapsed,
|
| 547 |
+
"stage_times": stage_times,
|
| 548 |
+
"requests": {
|
| 549 |
+
"fullpage": len(direct),
|
| 550 |
+
"detection": len(detection or []),
|
| 551 |
+
"crop": len(crop_responses),
|
| 552 |
+
},
|
| 553 |
+
"token_limited_responses": dict(capped),
|
| 554 |
+
"completion_tokens": sum(
|
| 555 |
+
int(r.get("usage", {}).get("completion_tokens") or 0) for r in all_responses
|
| 556 |
+
),
|
| 557 |
+
"postprocessing": asdict(parser.options),
|
| 558 |
+
"allow_partial": allow_partial,
|
| 559 |
+
"timing_excludes": [
|
| 560 |
+
"initial image preparation",
|
| 561 |
+
"output serialization",
|
| 562 |
+
"server startup",
|
| 563 |
+
"warmup requests",
|
| 564 |
+
],
|
| 565 |
+
"timing_includes": [
|
| 566 |
+
"all mode inference stages",
|
| 567 |
+
"predicted-region crop preparation",
|
| 568 |
+
"merge and document rendering",
|
| 569 |
+
],
|
| 570 |
+
}
|
| 571 |
+
return BatchResult(documents, metrics)
|
nemotron_parse/engine/__init__.py
ADDED
|
@@ -0,0 +1,2 @@
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
| 2 |
+
# SPDX-License-Identifier: OpenMDW-1.1
|
nemotron_parse/engine/caption_relocation.py
ADDED
|
@@ -0,0 +1,89 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
| 2 |
+
# SPDX-License-Identifier: OpenMDW-1.1
|
| 3 |
+
"""Relocate only plain above-chart captions with geometric/preservation guards."""
|
| 4 |
+
import re
|
| 5 |
+
from bs4 import BeautifulSoup
|
| 6 |
+
|
| 7 |
+
BLOCK = re.compile(
|
| 8 |
+
r"(<x_([\d.]+)><y_([\d.]+)>)(.*?)(<x_([\d.]+)><y_([\d.]+)><class_([^>]+)>)", re.S
|
| 9 |
+
)
|
| 10 |
+
|
| 11 |
+
|
| 12 |
+
def plain(text):
|
| 13 |
+
text = text.replace("<tbc>", "").strip()
|
| 14 |
+
if any(marker in text for marker in ["**", "__", "~~", "`"]):
|
| 15 |
+
return None
|
| 16 |
+
if re.search(r"(?<!\w)_[^_\n]+_(?!\w)|<[/a-zA-Z!]|\\[A-Za-z&%_]|\]\(", text):
|
| 17 |
+
return None
|
| 18 |
+
return text if text and "\n" not in text else None
|
| 19 |
+
|
| 20 |
+
|
| 21 |
+
def relocate_above_captions(raw, markdown, render_base):
|
| 22 |
+
matches = list(BLOCK.finditer(raw))
|
| 23 |
+
eligible = {}
|
| 24 |
+
targets = {}
|
| 25 |
+
moved = set()
|
| 26 |
+
allowed = {"Chart"}
|
| 27 |
+
boxes = [tuple(map(float, [m[2], m[3], m[6], m[7]])) for m in matches]
|
| 28 |
+
for i, m in enumerate(matches):
|
| 29 |
+
if m[8] != "Caption" or plain(m[4]) is None:
|
| 30 |
+
continue
|
| 31 |
+
box = boxes[i]
|
| 32 |
+
candidates = []
|
| 33 |
+
for j in [i - 1, i + 1]:
|
| 34 |
+
if not 0 <= j < len(matches) or matches[j][8] not in allowed:
|
| 35 |
+
continue
|
| 36 |
+
other = boxes[j]
|
| 37 |
+
if box[3] > other[1]:
|
| 38 |
+
continue # Attach above-object captions only; retain notes below as standalone.
|
| 39 |
+
width = min(box[2] - box[0], other[2] - other[0])
|
| 40 |
+
overlap = max(0, min(box[2], other[2]) - max(box[0], other[0]))
|
| 41 |
+
gap = max(other[1] - box[3], box[1] - other[3])
|
| 42 |
+
if width > 0 and overlap / width >= 0.8 and 0 <= gap <= 0.03:
|
| 43 |
+
candidates.append((gap, j))
|
| 44 |
+
candidates.sort()
|
| 45 |
+
if not candidates or (
|
| 46 |
+
len(candidates) > 1 and abs(candidates[0][0] - candidates[1][0]) < 1e-9
|
| 47 |
+
):
|
| 48 |
+
continue
|
| 49 |
+
gap, j = candidates[0]
|
| 50 |
+
eligible.setdefault(j, []).append((gap, i))
|
| 51 |
+
for j, caps in eligible.items():
|
| 52 |
+
caps.sort()
|
| 53 |
+
if len(caps) > 1 and abs(caps[0][0] - caps[1][0]) < 1e-9:
|
| 54 |
+
continue
|
| 55 |
+
caption_index = caps[0][1]
|
| 56 |
+
label = plain(matches[caption_index][4])
|
| 57 |
+
m = matches[j]
|
| 58 |
+
body = m[4]
|
| 59 |
+
if len(re.findall(r"<chart\b", body)) != 1:
|
| 60 |
+
continue
|
| 61 |
+
title = re.search(r"(?m)^title:[ \t]*([^\n]*)", body)
|
| 62 |
+
if title:
|
| 63 |
+
if label not in title[1]:
|
| 64 |
+
body = (
|
| 65 |
+
body[: title.start()]
|
| 66 |
+
+ "title: "
|
| 67 |
+
+ title[1].strip()
|
| 68 |
+
+ " — "
|
| 69 |
+
+ label
|
| 70 |
+
+ body[title.end() :]
|
| 71 |
+
)
|
| 72 |
+
else:
|
| 73 |
+
body = re.sub(
|
| 74 |
+
r"(<chart\b[^>]*>)", lambda q: q[0] + "\ntitle: " + label, body, count=1
|
| 75 |
+
)
|
| 76 |
+
derived = m[1] + body + m[5]
|
| 77 |
+
rendered = render_base(derived)
|
| 78 |
+
# Do not move text unless it survives inside an actual table caption.
|
| 79 |
+
captions = BeautifulSoup(rendered, "html.parser").find_all("caption")
|
| 80 |
+
if not any(label in cap.get_text(" ", strip=True) for cap in captions):
|
| 81 |
+
continue
|
| 82 |
+
targets[m.start()] = derived
|
| 83 |
+
moved.add(matches[caption_index].start())
|
| 84 |
+
if not targets:
|
| 85 |
+
return markdown
|
| 86 |
+
derived = BLOCK.sub(
|
| 87 |
+
lambda m: "" if m.start() in moved else targets.get(m.start(), m[0]), raw
|
| 88 |
+
)
|
| 89 |
+
return render_base(derived)
|
nemotron_parse/engine/chart_caption.py
ADDED
|
@@ -0,0 +1,71 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
| 2 |
+
# SPDX-License-Identifier: OpenMDW-1.1
|
| 3 |
+
|
| 4 |
+
"""Render titles explicitly declared by native chart metadata as table captions."""
|
| 5 |
+
import re
|
| 6 |
+
from html import escape
|
| 7 |
+
|
| 8 |
+
|
| 9 |
+
def transform(text):
|
| 10 |
+
changed = 0
|
| 11 |
+
|
| 12 |
+
def chart(m):
|
| 13 |
+
nonlocal changed
|
| 14 |
+
body = m[2]
|
| 15 |
+
if "<chart" in body:
|
| 16 |
+
return m[0]
|
| 17 |
+
title = re.search(r"(?m)^title:[ \t]*([^\n]*)$", body)
|
| 18 |
+
if not title:
|
| 19 |
+
return m[0]
|
| 20 |
+
title_text = title[1].strip()
|
| 21 |
+
if not title_text:
|
| 22 |
+
return m[0]
|
| 23 |
+
count = 0
|
| 24 |
+
pattern = re.compile(
|
| 25 |
+
r"(?m)^\|[^\n]*\|[ \t]*\n\|[ \t:|\-]+\|[ \t]*\n(?:\|[^\n]*\|[ \t]*(?:\n|$))+"
|
| 26 |
+
)
|
| 27 |
+
|
| 28 |
+
def table(t):
|
| 29 |
+
nonlocal count
|
| 30 |
+
raw = t[0]
|
| 31 |
+
if r"\|" in raw or "`" in raw:
|
| 32 |
+
return raw
|
| 33 |
+
lines = [a.strip() for a in raw.strip().splitlines()]
|
| 34 |
+
cells = [[c.strip() for c in line.strip("|").split("|")] for line in lines]
|
| 35 |
+
if (
|
| 36 |
+
len(cells) < 3
|
| 37 |
+
or any(len(row) != len(cells[0]) for row in cells)
|
| 38 |
+
or any(
|
| 39 |
+
not re.fullmatch(r":?-+:?", c.replace(" ", "")) for c in cells[1]
|
| 40 |
+
)
|
| 41 |
+
):
|
| 42 |
+
return raw
|
| 43 |
+
count += 1
|
| 44 |
+
result = (
|
| 45 |
+
"<table><caption>"
|
| 46 |
+
+ escape(title_text)
|
| 47 |
+
+ "</caption><thead><tr>"
|
| 48 |
+
+ "".join('<th scope="col">' + escape(c) + "</th>" for c in cells[0])
|
| 49 |
+
+ "</tr></thead><tbody>"
|
| 50 |
+
)
|
| 51 |
+
result += (
|
| 52 |
+
"".join(
|
| 53 |
+
"<tr>"
|
| 54 |
+
+ "".join("<td>" + escape(c) + "</td>" for c in row)
|
| 55 |
+
+ "</tr>"
|
| 56 |
+
for row in cells[2:]
|
| 57 |
+
)
|
| 58 |
+
+ "</tbody></table>\n"
|
| 59 |
+
)
|
| 60 |
+
return result
|
| 61 |
+
|
| 62 |
+
new = pattern.sub(table, body)
|
| 63 |
+
if not count:
|
| 64 |
+
return m[0]
|
| 65 |
+
# The title is rendered in its associated captions, retaining text and roles.
|
| 66 |
+
new = re.sub(r"(?m)^title:[ \t]*[^\n]*\n?", "", new, count=1)
|
| 67 |
+
changed += count
|
| 68 |
+
return m[1] + new + m[3]
|
| 69 |
+
|
| 70 |
+
out = re.sub(r"(<chart\b[^>]*>)(.*?)(</chart>)", chart, text, flags=re.S)
|
| 71 |
+
return out, changed
|
nemotron_parse/engine/chart_context.py
ADDED
|
@@ -0,0 +1,40 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
| 2 |
+
# SPDX-License-Identifier: OpenMDW-1.1
|
| 3 |
+
"""Preserve explicit chart subtitles in their existing table-caption context."""
|
| 4 |
+
import re
|
| 5 |
+
from bs4 import BeautifulSoup, NavigableString
|
| 6 |
+
|
| 7 |
+
TABLE = re.compile(r"<table\b[^>]*>.*?</table>", re.S | re.I)
|
| 8 |
+
CHART = re.compile(r"(<chart\b[^>]*>)(.*?)(</chart>)", re.S | re.I)
|
| 9 |
+
|
| 10 |
+
|
| 11 |
+
def add_native_subtitles(text):
|
| 12 |
+
def chart(match):
|
| 13 |
+
body = match[2]
|
| 14 |
+
labels = [
|
| 15 |
+
m[1].strip()
|
| 16 |
+
for m in re.finditer(r"(?m)^subtitle:[ \t]*([^\n]*)", body)
|
| 17 |
+
if m[1].strip()
|
| 18 |
+
]
|
| 19 |
+
if not labels:
|
| 20 |
+
return match[0]
|
| 21 |
+
|
| 22 |
+
def table(m):
|
| 23 |
+
soup = BeautifulSoup(m[0], "html.parser")
|
| 24 |
+
node = soup.find("table")
|
| 25 |
+
cap = node.find("caption", recursive=False)
|
| 26 |
+
if not cap:
|
| 27 |
+
return m[0]
|
| 28 |
+
existing = cap.get_text(" ", strip=True)
|
| 29 |
+
added = []
|
| 30 |
+
for value in labels:
|
| 31 |
+
if value not in existing and value not in added:
|
| 32 |
+
added.append(value)
|
| 33 |
+
if not added:
|
| 34 |
+
return m[0]
|
| 35 |
+
cap.append(NavigableString(" — " + " — ".join(added)))
|
| 36 |
+
return str(node)
|
| 37 |
+
|
| 38 |
+
return match[1] + TABLE.sub(table, body) + match[3]
|
| 39 |
+
|
| 40 |
+
return CHART.sub(chart, text)
|
nemotron_parse/engine/convert/__init__.py
ADDED
|
@@ -0,0 +1,2 @@
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
| 2 |
+
# SPDX-License-Identifier: OpenMDW-1.1
|
nemotron_parse/engine/convert/hybrid_cleanup.py
ADDED
|
@@ -0,0 +1,67 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
| 2 |
+
# SPDX-License-Identifier: OpenMDW-1.1
|
| 3 |
+
|
| 4 |
+
"""Escape LaTeX cell literals while preserving balanced nested inline markup."""
|
| 5 |
+
import re
|
| 6 |
+
|
| 7 |
+
ESCAPES = {
|
| 8 |
+
"#": r"\#",
|
| 9 |
+
"%": r"\%",
|
| 10 |
+
"$": r"\$",
|
| 11 |
+
"_": r"\_",
|
| 12 |
+
"~": r"\textasciitilde{}",
|
| 13 |
+
"^": r"\textasciicircum{}",
|
| 14 |
+
}
|
| 15 |
+
MATH_OR_TAG = re.compile(r"\$\$.*?\$\$|\$[^$]*\$|\\\(.*?\\\)|\\\[.*?\\\]|<[^>]+>", re.S)
|
| 16 |
+
|
| 17 |
+
|
| 18 |
+
def _process_cell(text, escape_literals):
|
| 19 |
+
def scan(position, closing=None, depth=0):
|
| 20 |
+
if depth > 64:
|
| 21 |
+
raise ValueError("inline_markup_nesting_too_deep")
|
| 22 |
+
output = []
|
| 23 |
+
while position < len(text):
|
| 24 |
+
if closing and text.startswith(closing, position):
|
| 25 |
+
return "".join(output), position + len(closing), True
|
| 26 |
+
protected = MATH_OR_TAG.match(text, position)
|
| 27 |
+
if protected:
|
| 28 |
+
output.append(protected[0])
|
| 29 |
+
position = protected.end()
|
| 30 |
+
continue
|
| 31 |
+
if text[position] == "\\":
|
| 32 |
+
# Preserve escaped literals and command prefixes. Formatting
|
| 33 |
+
# inside command arguments is still scanned when reached.
|
| 34 |
+
output.append(text[position : position + 2])
|
| 35 |
+
position += 2
|
| 36 |
+
continue
|
| 37 |
+
marker = next(
|
| 38 |
+
(m for m in ("**", "__", "~~", "_") if text.startswith(m, position)),
|
| 39 |
+
None,
|
| 40 |
+
)
|
| 41 |
+
if marker:
|
| 42 |
+
body, end, balanced = scan(position + len(marker), marker, depth + 1)
|
| 43 |
+
if balanced:
|
| 44 |
+
if escape_literals:
|
| 45 |
+
output.extend((marker, body, marker))
|
| 46 |
+
else:
|
| 47 |
+
tag = {"**": "b", "__": "b", "_": "i", "~~": "s"}[marker]
|
| 48 |
+
output.extend(("<" + tag + ">", body, "</" + tag + ">"))
|
| 49 |
+
position = end
|
| 50 |
+
continue
|
| 51 |
+
output.append(
|
| 52 |
+
ESCAPES.get(text[position], text[position])
|
| 53 |
+
if escape_literals
|
| 54 |
+
else text[position]
|
| 55 |
+
)
|
| 56 |
+
position += 1
|
| 57 |
+
return "".join(output), position, closing is None
|
| 58 |
+
|
| 59 |
+
return scan(0)[0]
|
| 60 |
+
|
| 61 |
+
|
| 62 |
+
def escape_cell(text):
|
| 63 |
+
return _process_cell(text, True)
|
| 64 |
+
|
| 65 |
+
|
| 66 |
+
def decode_styles(text):
|
| 67 |
+
return _process_cell(text, False)
|
nemotron_parse/engine/convert/hybrid_inline.py
ADDED
|
@@ -0,0 +1,164 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
| 2 |
+
# SPDX-License-Identifier: OpenMDW-1.1
|
| 3 |
+
|
| 4 |
+
"""Bounded Markdown/markup inline dialect inside LaTeX table cells.
|
| 5 |
+
|
| 6 |
+
Escaped punctuation is literal. Formatting is parsed before literal decoding;
|
| 7 |
+
literal text is never reparsed as markup. Unknown syntax fails explicitly.
|
| 8 |
+
"""
|
| 9 |
+
import html
|
| 10 |
+
import re
|
| 11 |
+
from bs4 import NavigableString, BeautifulSoup
|
| 12 |
+
|
| 13 |
+
try:
|
| 14 |
+
from .semantic_styles import mark_open, ALIASES
|
| 15 |
+
except ImportError:
|
| 16 |
+
from semantic_styles import mark_open, ALIASES
|
| 17 |
+
try:
|
| 18 |
+
from .rich_table_reader import TAGS, LITERALS, argument
|
| 19 |
+
except ImportError:
|
| 20 |
+
from rich_table_reader import TAGS, LITERALS, argument
|
| 21 |
+
|
| 22 |
+
ESC = {
|
| 23 |
+
"&": r"\&",
|
| 24 |
+
"%": r"\%",
|
| 25 |
+
"$": r"\$",
|
| 26 |
+
"#": r"\#",
|
| 27 |
+
"_": r"\_",
|
| 28 |
+
"*": r"\*",
|
| 29 |
+
"{": r"\{",
|
| 30 |
+
"}": r"\}",
|
| 31 |
+
"\\": r"\textbackslash{}",
|
| 32 |
+
"~": r"\textasciitilde{}",
|
| 33 |
+
"^": r"\textasciicircum{}",
|
| 34 |
+
"<": r"\textless{}",
|
| 35 |
+
">": r"\textgreater{}",
|
| 36 |
+
}
|
| 37 |
+
|
| 38 |
+
|
| 39 |
+
def encode_node(node):
|
| 40 |
+
if isinstance(node, NavigableString):
|
| 41 |
+
return "".join(ESC.get(char, char) for char in str(node))
|
| 42 |
+
if node.name == "mark":
|
| 43 |
+
return (
|
| 44 |
+
mark_open(node)
|
| 45 |
+
+ "".join(encode_node(child) for child in node.contents)
|
| 46 |
+
+ "</mark>"
|
| 47 |
+
)
|
| 48 |
+
if node.name in ALIASES:
|
| 49 |
+
node = BeautifulSoup(str(node), "html.parser").find()
|
| 50 |
+
node.name = ALIASES[node.name]
|
| 51 |
+
if node.name == "br":
|
| 52 |
+
return r"\linebreak{}"
|
| 53 |
+
content = "".join(encode_node(child) for child in node.contents)
|
| 54 |
+
if node.name in {"b", "strong"}:
|
| 55 |
+
return "**" + content + "**"
|
| 56 |
+
if node.name in {"i", "em"}:
|
| 57 |
+
# The historical cleanup deliberately strips single punctuation emphasis.
|
| 58 |
+
# Markup preserves that uncommon case without changing global cleanup.
|
| 59 |
+
before = str(node.previous_sibling or "")[-1:]
|
| 60 |
+
after = str(node.next_sibling or "")[:1]
|
| 61 |
+
if (
|
| 62 |
+
re.fullmatch(r"\\?[\W_]", node.get_text())
|
| 63 |
+
or not node.get_text().strip()
|
| 64 |
+
or re.match(r"\w", before)
|
| 65 |
+
or re.match(r"\w", after)
|
| 66 |
+
):
|
| 67 |
+
return "<i>" + content + "</i>"
|
| 68 |
+
return "_" + content + "_"
|
| 69 |
+
if node.name == "s":
|
| 70 |
+
return "~~" + content + "~~"
|
| 71 |
+
if node.name == "u":
|
| 72 |
+
return "<u>" + content + "</u>"
|
| 73 |
+
if node.name in {"sup", "sub"}:
|
| 74 |
+
return "<" + node.name + ">" + content + "</" + node.name + ">"
|
| 75 |
+
return content
|
| 76 |
+
|
| 77 |
+
|
| 78 |
+
def inline_html(text):
|
| 79 |
+
def parse(position, closing=None):
|
| 80 |
+
output = []
|
| 81 |
+
while position < len(text):
|
| 82 |
+
if closing and text.startswith(closing, position):
|
| 83 |
+
return "".join(output), position + len(closing)
|
| 84 |
+
if text.startswith("<br>", position):
|
| 85 |
+
output.append("<br>")
|
| 86 |
+
position += 4
|
| 87 |
+
continue
|
| 88 |
+
if text[position] == "\n":
|
| 89 |
+
output.append("<br>")
|
| 90 |
+
position += 1
|
| 91 |
+
continue
|
| 92 |
+
mark = re.match(r"<mark(?:\s[^<>]*)?>", text[position:], re.I)
|
| 93 |
+
if mark:
|
| 94 |
+
opener = mark_open(
|
| 95 |
+
BeautifulSoup(mark[0] + "</mark>", "html.parser").mark
|
| 96 |
+
)
|
| 97 |
+
inner, position = parse(position + len(mark[0]), "</mark>")
|
| 98 |
+
output.append(opener + inner + "</mark>")
|
| 99 |
+
continue
|
| 100 |
+
marker = next(
|
| 101 |
+
(
|
| 102 |
+
m
|
| 103 |
+
for m in ("**", "_", "~~", "<u>", "<s>", "<i>", "<sup>", "<sub>")
|
| 104 |
+
if text.startswith(m, position)
|
| 105 |
+
),
|
| 106 |
+
None,
|
| 107 |
+
)
|
| 108 |
+
if marker:
|
| 109 |
+
tag = {
|
| 110 |
+
"**": "b",
|
| 111 |
+
"_": "i",
|
| 112 |
+
"~~": "s",
|
| 113 |
+
"<u>": "u",
|
| 114 |
+
"<s>": "s",
|
| 115 |
+
"<i>": "i",
|
| 116 |
+
"<sup>": "sup",
|
| 117 |
+
"<sub>": "sub",
|
| 118 |
+
}[marker]
|
| 119 |
+
end = marker if marker in {"**", "_", "~~"} else "</" + tag + ">"
|
| 120 |
+
inner, position = parse(position + len(marker), end)
|
| 121 |
+
output.append("<" + tag + ">" + inner + "</" + tag + ">")
|
| 122 |
+
continue
|
| 123 |
+
if text[position] != "\\":
|
| 124 |
+
if text[position] in "%$#_*{}^~<>":
|
| 125 |
+
raise ValueError("unsupported_unescaped_hybrid_special")
|
| 126 |
+
output.append(html.escape(text[position]))
|
| 127 |
+
position += 1
|
| 128 |
+
continue
|
| 129 |
+
if position + 1 < len(text) and text[position + 1] in "&%$#_*{}":
|
| 130 |
+
output.append(html.escape(text[position + 1]))
|
| 131 |
+
position += 2
|
| 132 |
+
continue
|
| 133 |
+
match = re.match(r"\\([A-Za-z]+)", text[position:])
|
| 134 |
+
if not match:
|
| 135 |
+
raise ValueError("unsupported_hybrid_escape")
|
| 136 |
+
command = match[1]
|
| 137 |
+
position += len(match[0])
|
| 138 |
+
body, position = argument(text, position)
|
| 139 |
+
if command == "colorbox":
|
| 140 |
+
color = body
|
| 141 |
+
body, position = argument(text, position)
|
| 142 |
+
node = BeautifulSoup("<mark></mark>", "html.parser").mark
|
| 143 |
+
node["background"] = color
|
| 144 |
+
output.append(mark_open(node) + inline_html(body) + "</mark>")
|
| 145 |
+
continue
|
| 146 |
+
if command == "hl":
|
| 147 |
+
output.append("<mark>" + inline_html(body) + "</mark>")
|
| 148 |
+
continue
|
| 149 |
+
if command == "linebreak" and not body:
|
| 150 |
+
output.append("<br>")
|
| 151 |
+
elif command == "mbox" and not body:
|
| 152 |
+
pass
|
| 153 |
+
elif command in TAGS:
|
| 154 |
+
tag = TAGS[command]
|
| 155 |
+
output.append("<" + tag + ">" + inline_html(body) + "</" + tag + ">")
|
| 156 |
+
elif command in LITERALS and not body:
|
| 157 |
+
output.append(html.escape(LITERALS[command]))
|
| 158 |
+
else:
|
| 159 |
+
raise ValueError("unsupported_hybrid_command_" + command)
|
| 160 |
+
if closing:
|
| 161 |
+
raise ValueError("unclosed_hybrid_style_" + closing)
|
| 162 |
+
return "".join(output), position
|
| 163 |
+
|
| 164 |
+
return parse(0)[0]
|
nemotron_parse/engine/convert/latex2html.py
ADDED
|
@@ -0,0 +1,653 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
| 2 |
+
# SPDX-License-Identifier: OpenMDW-1.1
|
| 3 |
+
|
| 4 |
+
import csv
|
| 5 |
+
import io
|
| 6 |
+
import json
|
| 7 |
+
import re
|
| 8 |
+
|
| 9 |
+
from bs4 import BeautifulSoup
|
| 10 |
+
|
| 11 |
+
|
| 12 |
+
def skip_whitespace(text, i):
|
| 13 |
+
"""Advance index i past any whitespace."""
|
| 14 |
+
while i < len(text) and text[i].isspace():
|
| 15 |
+
i += 1
|
| 16 |
+
return i
|
| 17 |
+
|
| 18 |
+
|
| 19 |
+
def parse_braced_argument(text, i):
|
| 20 |
+
"""
|
| 21 |
+
Given text and an index i that should point at an opening '{',
|
| 22 |
+
return a tuple (argument_content, new_index) where argument_content is the full
|
| 23 |
+
string inside the balanced braces and new_index is the position just after the matching '}'.
|
| 24 |
+
"""
|
| 25 |
+
if i >= len(text) or text[i] != "{":
|
| 26 |
+
raise ValueError("Expected '{' at position {}".format(i))
|
| 27 |
+
i += 1 # skip the opening brace
|
| 28 |
+
start = i
|
| 29 |
+
level = 1
|
| 30 |
+
while i < len(text) and level > 0:
|
| 31 |
+
if text[i] == "{":
|
| 32 |
+
level += 1
|
| 33 |
+
elif text[i] == "}":
|
| 34 |
+
level -= 1
|
| 35 |
+
i += 1
|
| 36 |
+
if level != 0:
|
| 37 |
+
raise ValueError("Unbalanced braces starting at position {}".format(start - 1))
|
| 38 |
+
# The argument content is from start to i-1 (excluding the closing brace)
|
| 39 |
+
return text[start : i - 1], i
|
| 40 |
+
|
| 41 |
+
|
| 42 |
+
def parse_command(text, i):
|
| 43 |
+
r"""
|
| 44 |
+
Parse a \multirow or \multicolumn command starting at index i.
|
| 45 |
+
This function assumes the command has exactly three braced arguments.
|
| 46 |
+
|
| 47 |
+
It processes each argument recursively. For the third argument, after recursive processing,
|
| 48 |
+
it replaces any unescaped & with \&.
|
| 49 |
+
|
| 50 |
+
Returns a tuple (command_text, new_index) where command_text is the reconstructed command.
|
| 51 |
+
"""
|
| 52 |
+
# Determine which command we have.
|
| 53 |
+
if text.startswith(r"\multirow", i):
|
| 54 |
+
command_name = r"\multirow"
|
| 55 |
+
i += len(r"\multirow")
|
| 56 |
+
elif text.startswith(r"\multicolumn", i):
|
| 57 |
+
command_name = r"\multicolumn"
|
| 58 |
+
i += len(r"\multicolumn")
|
| 59 |
+
else:
|
| 60 |
+
raise ValueError(
|
| 61 |
+
"Expected \\multirow or \\multicolumn at position {}".format(i)
|
| 62 |
+
)
|
| 63 |
+
|
| 64 |
+
# Skip whitespace between the command name and the first argument.
|
| 65 |
+
i = skip_whitespace(text, i)
|
| 66 |
+
args = []
|
| 67 |
+
# Expect exactly three arguments
|
| 68 |
+
for arg_index in range(3):
|
| 69 |
+
if i >= len(text) or text[i] != "{":
|
| 70 |
+
raise ValueError(
|
| 71 |
+
"Expected '{' for argument {} at position {}".format(arg_index + 1, i)
|
| 72 |
+
)
|
| 73 |
+
arg_content, i = parse_braced_argument(text, i)
|
| 74 |
+
# Process the content recursively to catch nested commands
|
| 75 |
+
processed_arg = clean_multi_cells(arg_content)
|
| 76 |
+
if arg_index == 2:
|
| 77 |
+
# For the cell text (third argument), replace any unescaped &
|
| 78 |
+
processed_arg = re.sub(r"(?<!\\)&", r"\\&", processed_arg)
|
| 79 |
+
args.append(processed_arg)
|
| 80 |
+
# Only skip whitespace between arguments, not after the last one.
|
| 81 |
+
if arg_index < 2:
|
| 82 |
+
i = skip_whitespace(text, i)
|
| 83 |
+
# Reconstruct the full command with its three arguments
|
| 84 |
+
command_text = f"{command_name}{{{args[0]}}}{{{args[1]}}}{{{args[2]}}}"
|
| 85 |
+
return command_text, i
|
| 86 |
+
|
| 87 |
+
|
| 88 |
+
def clean_multi_cells(text):
|
| 89 |
+
r"""
|
| 90 |
+
Process an arbitrary LaTeX text string and look for occurrences of \multirow or \multicolumn commands.
|
| 91 |
+
When found, the command is parsed (handling nested braces and nested commands) and its third argument is fixed.
|
| 92 |
+
|
| 93 |
+
Returns the processed text.
|
| 94 |
+
"""
|
| 95 |
+
result = []
|
| 96 |
+
i = 0
|
| 97 |
+
while i < len(text):
|
| 98 |
+
# Find next occurrence of either command.
|
| 99 |
+
idx_multi = text.find(r"\multirow", i)
|
| 100 |
+
idx_multiC = text.find(r"\multicolumn", i)
|
| 101 |
+
|
| 102 |
+
# Determine the next index among the two (if any)
|
| 103 |
+
if idx_multi == -1 and idx_multiC == -1:
|
| 104 |
+
result.append(text[i:])
|
| 105 |
+
break
|
| 106 |
+
if idx_multi == -1:
|
| 107 |
+
next_idx = idx_multiC
|
| 108 |
+
elif idx_multiC == -1:
|
| 109 |
+
next_idx = idx_multi
|
| 110 |
+
else:
|
| 111 |
+
next_idx = min(idx_multi, idx_multiC)
|
| 112 |
+
|
| 113 |
+
# Append text before the command (preserving any whitespace)
|
| 114 |
+
result.append(text[i:next_idx])
|
| 115 |
+
# Process the command starting at next_idx
|
| 116 |
+
command_text, new_index = parse_command(text, next_idx)
|
| 117 |
+
result.append(command_text)
|
| 118 |
+
i = new_index
|
| 119 |
+
return "".join(result)
|
| 120 |
+
|
| 121 |
+
|
| 122 |
+
def parse_brace(s, pos):
|
| 123 |
+
"""
|
| 124 |
+
Given a string s and an index pos pointing to an opening '{',
|
| 125 |
+
returns a tuple (content, new_pos) where content is the string
|
| 126 |
+
between the matching braces (handling nested braces) and new_pos is
|
| 127 |
+
the index just after the closing '}'.
|
| 128 |
+
"""
|
| 129 |
+
if pos >= len(s) or s[pos] != "{":
|
| 130 |
+
raise ValueError("Expected '{' at position %d" % pos)
|
| 131 |
+
pos += 1 # skip the opening brace
|
| 132 |
+
content = ""
|
| 133 |
+
depth = 1
|
| 134 |
+
while pos < len(s) and depth:
|
| 135 |
+
char = s[pos]
|
| 136 |
+
if char == "{":
|
| 137 |
+
depth += 1
|
| 138 |
+
content += char
|
| 139 |
+
elif char == "}":
|
| 140 |
+
depth -= 1
|
| 141 |
+
if depth:
|
| 142 |
+
content += char
|
| 143 |
+
else:
|
| 144 |
+
content += char
|
| 145 |
+
pos += 1
|
| 146 |
+
if depth != 0:
|
| 147 |
+
raise ValueError("Unmatched '{' in string.")
|
| 148 |
+
return content, pos
|
| 149 |
+
|
| 150 |
+
|
| 151 |
+
def parse_command_merge(s, pos):
|
| 152 |
+
"""
|
| 153 |
+
Parse a multirow or multicolumn command starting at s[pos]. If the content
|
| 154 |
+
of the command contains a nested command, then recursively parse the inner
|
| 155 |
+
command and merge its parameters with the outer ones. The merging is done
|
| 156 |
+
so that the outer multirow’s parameters (e.g. rowspan and width) are kept
|
| 157 |
+
while the inner command’s parameters (e.g. colspan, alignment) and its innermost
|
| 158 |
+
content are returned.
|
| 159 |
+
|
| 160 |
+
Returns a tuple (merged_dict, new_pos) where merged_dict is a dictionary
|
| 161 |
+
containing the combined parameters and new_pos is the updated index after
|
| 162 |
+
parsing the command.
|
| 163 |
+
"""
|
| 164 |
+
if s.startswith(r"\multirow", pos):
|
| 165 |
+
newpos = pos + len(r"\multirow")
|
| 166 |
+
# Parse the three required arguments for multirow: rowspan, width, and content.
|
| 167 |
+
rowspan, newpos = parse_brace(s, newpos)
|
| 168 |
+
width, newpos = parse_brace(s, newpos)
|
| 169 |
+
content, newpos = parse_brace(s, newpos)
|
| 170 |
+
# Look for a nested command (either \multirow or \multicolumn) in the content.
|
| 171 |
+
index_mr = content.find(r"\multirow")
|
| 172 |
+
index_mc = content.find(r"\multicolumn")
|
| 173 |
+
if index_mr == -1 and index_mc == -1:
|
| 174 |
+
# No nested command found; return this command’s details.
|
| 175 |
+
return {
|
| 176 |
+
"rowspan": rowspan.strip(),
|
| 177 |
+
"width": width.strip(),
|
| 178 |
+
"content": content.strip(),
|
| 179 |
+
}, newpos
|
| 180 |
+
else:
|
| 181 |
+
# At least one nested command is present. Pick the first occurrence.
|
| 182 |
+
indices = [i for i in (index_mr, index_mc) if i != -1]
|
| 183 |
+
first_index = min(indices)
|
| 184 |
+
# Parse the inner (nested) command from within the content.
|
| 185 |
+
inner, _ = parse_command_merge(content, first_index)
|
| 186 |
+
# Merge: keep the outer multirow’s parameters and add the inner ones.
|
| 187 |
+
merged = {"rowspan": rowspan.strip(), "width": width.strip()}
|
| 188 |
+
merged.update(inner)
|
| 189 |
+
return merged, newpos
|
| 190 |
+
|
| 191 |
+
elif s.startswith(r"\multicolumn", pos):
|
| 192 |
+
newpos = pos + len(r"\multicolumn")
|
| 193 |
+
# Parse the three arguments for multicolumn: colspan, alignment, and content.
|
| 194 |
+
colspan, newpos = parse_brace(s, newpos)
|
| 195 |
+
alignment, newpos = parse_brace(s, newpos)
|
| 196 |
+
content, newpos = parse_brace(s, newpos)
|
| 197 |
+
# Look for a nested command in the content.
|
| 198 |
+
index_mr = content.find(r"\multirow")
|
| 199 |
+
index_mc = content.find(r"\multicolumn")
|
| 200 |
+
if index_mr == -1 and index_mc == -1:
|
| 201 |
+
return {
|
| 202 |
+
"colspan": colspan.strip(),
|
| 203 |
+
"alignment": alignment.strip(),
|
| 204 |
+
"content": content.strip(),
|
| 205 |
+
}, newpos
|
| 206 |
+
else:
|
| 207 |
+
indices = [i for i in (index_mr, index_mc) if i != -1]
|
| 208 |
+
first_index = min(indices)
|
| 209 |
+
inner, _ = parse_command_merge(content, first_index)
|
| 210 |
+
merged = {"colspan": colspan.strip(), "alignment": alignment.strip()}
|
| 211 |
+
merged.update(inner)
|
| 212 |
+
return merged, newpos
|
| 213 |
+
|
| 214 |
+
# Not a recognized command starting at pos.
|
| 215 |
+
return None, pos
|
| 216 |
+
|
| 217 |
+
|
| 218 |
+
def extract_merged_commands(s):
|
| 219 |
+
"""
|
| 220 |
+
Scan through the LaTeX string s and extract merged multirow/multicolumn commands.
|
| 221 |
+
For each command found, if there is nesting the parser merges the outer and inner
|
| 222 |
+
parameters so that the final result includes both the rowspan (or width) and the colspan
|
| 223 |
+
(or alignment) along with the innermost content.
|
| 224 |
+
|
| 225 |
+
Returns a list of dictionaries.
|
| 226 |
+
"""
|
| 227 |
+
pos = 0
|
| 228 |
+
results = []
|
| 229 |
+
while pos < len(s):
|
| 230 |
+
if s[pos] == "\\":
|
| 231 |
+
res, newpos = parse_command_merge(s, pos)
|
| 232 |
+
if res is not None:
|
| 233 |
+
results.append(res)
|
| 234 |
+
pos = newpos
|
| 235 |
+
continue
|
| 236 |
+
pos += 1
|
| 237 |
+
return results
|
| 238 |
+
|
| 239 |
+
|
| 240 |
+
def remove_tags(html, tags_to_remove):
|
| 241 |
+
soup = BeautifulSoup(html, "html.parser")
|
| 242 |
+
# Loop through the tags to remove
|
| 243 |
+
for tag_name in tags_to_remove:
|
| 244 |
+
for tag in soup.find_all(tag_name):
|
| 245 |
+
# Move the children of the tag to the parent tag
|
| 246 |
+
tag.unwrap() # This removes the tag but keeps its contents
|
| 247 |
+
# Return the modified HTML as a string
|
| 248 |
+
return str(soup)
|
| 249 |
+
|
| 250 |
+
|
| 251 |
+
def convert_th_to_td(html):
|
| 252 |
+
"""Replace all th tags with td tags"""
|
| 253 |
+
soup = BeautifulSoup(html, "html.parser")
|
| 254 |
+
for th_tag in soup.find_all("th"):
|
| 255 |
+
th_tag.name = "td"
|
| 256 |
+
return str(soup)
|
| 257 |
+
|
| 258 |
+
|
| 259 |
+
def replace_italic(text):
|
| 260 |
+
pattern = re.compile(r"(?<!\\)_(.*?)(?<!\\)_")
|
| 261 |
+
|
| 262 |
+
def italic_replacer(match):
|
| 263 |
+
# Get the text inside the underscores.
|
| 264 |
+
content = match.group(1)
|
| 265 |
+
# Remove the escape (backslash) from any escaped underscores inside.
|
| 266 |
+
content = content.replace(r"\_", "_")
|
| 267 |
+
return f"<i>{content}</i>"
|
| 268 |
+
|
| 269 |
+
# Replace all occurrences of the pattern using the replacer function.
|
| 270 |
+
return pattern.sub(italic_replacer, text)
|
| 271 |
+
|
| 272 |
+
|
| 273 |
+
def replace_bold(text):
|
| 274 |
+
pattern = re.compile(r"(?<!\\)\*\*(.*?)(?<!\\)\*\*")
|
| 275 |
+
|
| 276 |
+
def bold_replacer(match):
|
| 277 |
+
content = match.group(1)
|
| 278 |
+
# Unescape any escaped asterisks within the captured text.
|
| 279 |
+
content = content.replace(r"\*", "*")
|
| 280 |
+
return f"<b>{content}</b>"
|
| 281 |
+
|
| 282 |
+
return pattern.sub(bold_replacer, text)
|
| 283 |
+
|
| 284 |
+
|
| 285 |
+
def latex_table_to_html(latex_str, add_head_body=False):
|
| 286 |
+
# Strict paired reader for rich distillation targets; legacy syntax follows.
|
| 287 |
+
try:
|
| 288 |
+
from .rich_table_reader import to_html as rich_to_html
|
| 289 |
+
except ImportError:
|
| 290 |
+
from rich_table_reader import to_html as rich_to_html
|
| 291 |
+
try:
|
| 292 |
+
from .rich_table_fallback import decode_known_inline, restore_header_cells
|
| 293 |
+
except ImportError:
|
| 294 |
+
from rich_table_fallback import decode_known_inline, restore_header_cells
|
| 295 |
+
# Pattern to match the entire tabular environment
|
| 296 |
+
table_pattern = r"\\begin{tabular}{([^}]*)}\s*(.*?)\\end{tabular}"
|
| 297 |
+
|
| 298 |
+
def process_cell(cell):
|
| 299 |
+
# Clean up cell content
|
| 300 |
+
cell = cell.strip()
|
| 301 |
+
if re.fullmatch(r"\{\s*\}", cell):
|
| 302 |
+
cell = ""
|
| 303 |
+
|
| 304 |
+
out = extract_merged_commands(cell)
|
| 305 |
+
if len(out) > 0:
|
| 306 |
+
cell = process_cell(out[0]["content"])["content"]
|
| 307 |
+
rowspan = int(out[0].get("rowspan", "1"))
|
| 308 |
+
colspan = int(out[0].get("colspan", "1"))
|
| 309 |
+
return {"content": cell, "colspan": colspan, "rowspan": rowspan}
|
| 310 |
+
|
| 311 |
+
# Replace latex and markdown formatting with HTML tags
|
| 312 |
+
cell = re.sub(r"\$([^$]*)\$", r"\1", cell) # Remove math mode
|
| 313 |
+
cell = re.sub(r"\\textbf{([^}]*)}", r"<b>\1</b>", cell) # Convert latex bold
|
| 314 |
+
cell = re.sub(r"\\textit{([^}]*)}", r"<i>\1</i>", cell) # Convert latex italic
|
| 315 |
+
cell = re.sub(r"\\emph{([^}]*)}", r"<i>\1</i>", cell)
|
| 316 |
+
cell = re.sub(r"\\underline{([^}]*)}", r"<u>\1</u>", cell)
|
| 317 |
+
cell = re.sub(r"\\(?:sout|st){([^}]*)}", r"<s>\1</s>", cell)
|
| 318 |
+
cell = replace_italic(cell)
|
| 319 |
+
cell = replace_bold(cell)
|
| 320 |
+
cell = (
|
| 321 |
+
cell.replace("\\$", "$")
|
| 322 |
+
.replace("\\%", "%")
|
| 323 |
+
.replace("\\newline", "<br>")
|
| 324 |
+
.replace("\\textless", "<")
|
| 325 |
+
.replace("\\textgreater", ">")
|
| 326 |
+
.replace("\\*", "*")
|
| 327 |
+
.replace("\\_", "_")
|
| 328 |
+
.replace("\\backslash", "\\")
|
| 329 |
+
)
|
| 330 |
+
|
| 331 |
+
# Replace \& with & in the cell text
|
| 332 |
+
cell = cell.replace(r"\&", "&")
|
| 333 |
+
cell = cell.replace("<tbc>", "")
|
| 334 |
+
cell = re.sub(r"(?:\r?\n|<br\s*/?>)+", "<br>", cell, flags=re.IGNORECASE)
|
| 335 |
+
# Preserve in-cell line breaks as <br>; clean other tokens.
|
| 336 |
+
cell = (
|
| 337 |
+
cell.replace("\\unknown", "")
|
| 338 |
+
.replace("\\<|unk|\\>", "")
|
| 339 |
+
.replace("<underline>", "<u>")
|
| 340 |
+
.replace("</underline>", "</u>")
|
| 341 |
+
)
|
| 342 |
+
return {"content": cell, "colspan": 1, "rowspan": 1}
|
| 343 |
+
|
| 344 |
+
def split_row(input_string):
|
| 345 |
+
# Use a regular expression to split on '&' that is not preceded by a backslash
|
| 346 |
+
return re.split(r"(?<!\\)&", input_string)
|
| 347 |
+
|
| 348 |
+
def convert_table(match):
|
| 349 |
+
try:
|
| 350 |
+
rich = rich_to_html(match.group(0))
|
| 351 |
+
except ValueError:
|
| 352 |
+
rich = None
|
| 353 |
+
if rich is not None:
|
| 354 |
+
if add_head_body:
|
| 355 |
+
soup = BeautifulSoup(rich, "html.parser")
|
| 356 |
+
rows = list(soup.table.find_all("tr", recursive=False))
|
| 357 |
+
head = soup.new_tag("thead")
|
| 358 |
+
body = soup.new_tag("tbody")
|
| 359 |
+
in_header = True
|
| 360 |
+
for row in rows:
|
| 361 |
+
cells = row.find_all(["td", "th"], recursive=False)
|
| 362 |
+
in_header = (
|
| 363 |
+
in_header and bool(cells) and all(c.name == "th" for c in cells)
|
| 364 |
+
)
|
| 365 |
+
(head if in_header else body).append(row.extract())
|
| 366 |
+
soup.table.append(head)
|
| 367 |
+
soup.table.append(body)
|
| 368 |
+
rich = str(soup)
|
| 369 |
+
return rich
|
| 370 |
+
# Extract table content
|
| 371 |
+
format_spec, content = match.groups()
|
| 372 |
+
content = decode_known_inline(content)
|
| 373 |
+
|
| 374 |
+
# Start building HTML table
|
| 375 |
+
html = ["<table>"]
|
| 376 |
+
|
| 377 |
+
# Track cells for multirow
|
| 378 |
+
multirow_tracker = set()
|
| 379 |
+
|
| 380 |
+
# Process rows
|
| 381 |
+
rows = re.split(r"\\\\", content)
|
| 382 |
+
current_row = 0
|
| 383 |
+
|
| 384 |
+
for row in rows:
|
| 385 |
+
if not row.strip():
|
| 386 |
+
continue
|
| 387 |
+
|
| 388 |
+
row = row.strip()
|
| 389 |
+
|
| 390 |
+
# Skip \hline
|
| 391 |
+
if "\\hline" in row:
|
| 392 |
+
row = row.replace("\\hline", "")
|
| 393 |
+
if not row.strip():
|
| 394 |
+
continue
|
| 395 |
+
|
| 396 |
+
row = clean_multi_cells(row)
|
| 397 |
+
|
| 398 |
+
# Process cells
|
| 399 |
+
cells = split_row(row)
|
| 400 |
+
processed_cells = [process_cell(cell) for cell in cells]
|
| 401 |
+
|
| 402 |
+
if add_head_body:
|
| 403 |
+
if current_row == 0:
|
| 404 |
+
html.append(" <thead>")
|
| 405 |
+
if current_row == 1:
|
| 406 |
+
html.append(" <tbody>")
|
| 407 |
+
html.append(" <tr>")
|
| 408 |
+
current_col = 0
|
| 409 |
+
|
| 410 |
+
for cell in processed_cells:
|
| 411 |
+
content_segment = cell["content"]
|
| 412 |
+
|
| 413 |
+
attrs = []
|
| 414 |
+
if cell["colspan"] > 1:
|
| 415 |
+
attrs.append(f'colspan="{cell["colspan"]}"')
|
| 416 |
+
if cell["rowspan"] > 1:
|
| 417 |
+
attrs.append(f'rowspan="{cell["rowspan"]}"')
|
| 418 |
+
for r in range(current_row + 1, current_row + cell["rowspan"]):
|
| 419 |
+
for c in range(current_col, current_col + cell["colspan"]):
|
| 420 |
+
multirow_tracker.add((r, c))
|
| 421 |
+
|
| 422 |
+
# Empty multicolumn placeholders under a prior multirow are
|
| 423 |
+
# covered cells too. Skip only a fully occupied, empty span;
|
| 424 |
+
# do not remove content or infer missing/partial occupancy.
|
| 425 |
+
if (
|
| 426 |
+
content_segment == ""
|
| 427 |
+
and cell["rowspan"] == 1
|
| 428 |
+
and cell["colspan"] >= 1
|
| 429 |
+
and all(
|
| 430 |
+
(current_row, col) in multirow_tracker
|
| 431 |
+
for col in range(current_col, current_col + cell["colspan"])
|
| 432 |
+
)
|
| 433 |
+
):
|
| 434 |
+
current_col += cell["colspan"]
|
| 435 |
+
continue
|
| 436 |
+
|
| 437 |
+
attr_str = " " + " ".join(attrs) if attrs else ""
|
| 438 |
+
cell_tag = "td"
|
| 439 |
+
html.append(f" <{cell_tag}{attr_str}>{content_segment}</{cell_tag}>")
|
| 440 |
+
current_col += cell["colspan"]
|
| 441 |
+
|
| 442 |
+
if add_head_body and current_row == 0:
|
| 443 |
+
html.append(" </thead>")
|
| 444 |
+
html.append(" </tr>")
|
| 445 |
+
current_row += 1
|
| 446 |
+
if add_head_body:
|
| 447 |
+
html.append(" </tbody>")
|
| 448 |
+
html.append("</table>")
|
| 449 |
+
return restore_header_cells("\n".join(html))
|
| 450 |
+
|
| 451 |
+
# Convert all tabular environments in the input
|
| 452 |
+
return re.sub(table_pattern, convert_table, latex_str, flags=re.DOTALL)
|
| 453 |
+
|
| 454 |
+
|
| 455 |
+
def convert_single_table(table):
|
| 456 |
+
"""
|
| 457 |
+
Convert a single HTML table to Markdown format.
|
| 458 |
+
|
| 459 |
+
Args:
|
| 460 |
+
table: BeautifulSoup table element
|
| 461 |
+
|
| 462 |
+
Returns:
|
| 463 |
+
str: Markdown table string
|
| 464 |
+
"""
|
| 465 |
+
markdown_lines = []
|
| 466 |
+
rows = table.find_all("tr")
|
| 467 |
+
|
| 468 |
+
for i, row in enumerate(rows):
|
| 469 |
+
cells = row.find_all(["td", "th"])
|
| 470 |
+
if not cells:
|
| 471 |
+
continue
|
| 472 |
+
|
| 473 |
+
# Convert cells to text, handling nested elements
|
| 474 |
+
row_data = []
|
| 475 |
+
for cell in cells:
|
| 476 |
+
# Get text content, handling nested elements
|
| 477 |
+
cell_text = cell.get_text(separator=" ", strip=True)
|
| 478 |
+
# Escape pipe characters
|
| 479 |
+
cell_text = cell_text.replace("|", "\\|")
|
| 480 |
+
row_data.append(cell_text)
|
| 481 |
+
|
| 482 |
+
# Add row to markdown
|
| 483 |
+
markdown_lines.append("| " + " | ".join(row_data) + " |")
|
| 484 |
+
|
| 485 |
+
# Add separator after header row
|
| 486 |
+
if i == 0:
|
| 487 |
+
separator = "| " + " | ".join(["---"] * len(cells)) + " |"
|
| 488 |
+
markdown_lines.append(separator)
|
| 489 |
+
|
| 490 |
+
return "\n".join(markdown_lines)
|
| 491 |
+
|
| 492 |
+
|
| 493 |
+
def convert_html_tables_to_markdown(html_content):
|
| 494 |
+
"""
|
| 495 |
+
Find all HTML tables and convert them to Markdown while preserving all other content.
|
| 496 |
+
|
| 497 |
+
Args:
|
| 498 |
+
html_content (str): HTML content that may contain tables
|
| 499 |
+
|
| 500 |
+
Returns:
|
| 501 |
+
str: HTML content with tables converted to Markdown
|
| 502 |
+
"""
|
| 503 |
+
soup = BeautifulSoup(html_content, "html.parser")
|
| 504 |
+
|
| 505 |
+
# Find all tables
|
| 506 |
+
tables = soup.find_all("table")
|
| 507 |
+
|
| 508 |
+
if not tables:
|
| 509 |
+
return html_content # Return original content unchanged
|
| 510 |
+
|
| 511 |
+
# Convert each table to markdown and replace it
|
| 512 |
+
for table in tables:
|
| 513 |
+
markdown_table = convert_single_table(table)
|
| 514 |
+
|
| 515 |
+
# Create a new element to replace the table
|
| 516 |
+
replacement = soup.new_string("\n" + markdown_table + "\n")
|
| 517 |
+
table.replace_with(replacement)
|
| 518 |
+
|
| 519 |
+
return str(soup)
|
| 520 |
+
|
| 521 |
+
|
| 522 |
+
def _expand_spans(table):
|
| 523 |
+
"""
|
| 524 |
+
Expand colspan/rowspan in a BeautifulSoup <table> into a regular grid.
|
| 525 |
+
Returns a list-of-lists of cell text strings.
|
| 526 |
+
"""
|
| 527 |
+
rows = table.find_all("tr")
|
| 528 |
+
if not rows:
|
| 529 |
+
return []
|
| 530 |
+
|
| 531 |
+
num_cols = 0
|
| 532 |
+
for row in rows:
|
| 533 |
+
col_count = sum(int(c.get("colspan", 1)) for c in row.find_all(["td", "th"]))
|
| 534 |
+
num_cols = max(num_cols, col_count)
|
| 535 |
+
|
| 536 |
+
grid = [[None] * num_cols for _ in range(len(rows))]
|
| 537 |
+
|
| 538 |
+
for r, row in enumerate(rows):
|
| 539 |
+
cells = row.find_all(["td", "th"])
|
| 540 |
+
c = 0
|
| 541 |
+
for cell in cells:
|
| 542 |
+
while c < num_cols and grid[r][c] is not None:
|
| 543 |
+
c += 1
|
| 544 |
+
if c >= num_cols:
|
| 545 |
+
break
|
| 546 |
+
|
| 547 |
+
colspan = int(cell.get("colspan", 1))
|
| 548 |
+
rowspan = int(cell.get("rowspan", 1))
|
| 549 |
+
text = cell.get_text(separator=" ", strip=True)
|
| 550 |
+
|
| 551 |
+
for dr in range(rowspan):
|
| 552 |
+
for dc in range(colspan):
|
| 553 |
+
rr, cc = r + dr, c + dc
|
| 554 |
+
if rr < len(grid) and cc < num_cols:
|
| 555 |
+
grid[rr][cc] = text
|
| 556 |
+
c += colspan
|
| 557 |
+
|
| 558 |
+
for row in grid:
|
| 559 |
+
for i in range(len(row)):
|
| 560 |
+
if row[i] is None:
|
| 561 |
+
row[i] = ""
|
| 562 |
+
|
| 563 |
+
return grid
|
| 564 |
+
|
| 565 |
+
|
| 566 |
+
def _grid_to_records(grid):
|
| 567 |
+
"""
|
| 568 |
+
Convert a grid (list-of-lists) to a list of dicts using the first row as
|
| 569 |
+
column headers. Duplicate or empty headers get disambiguated with a suffix.
|
| 570 |
+
"""
|
| 571 |
+
if len(grid) < 2:
|
| 572 |
+
return grid
|
| 573 |
+
|
| 574 |
+
raw_headers = grid[0]
|
| 575 |
+
headers = []
|
| 576 |
+
seen = {}
|
| 577 |
+
for h in raw_headers:
|
| 578 |
+
key = h if h else "column"
|
| 579 |
+
if key in seen:
|
| 580 |
+
seen[key] += 1
|
| 581 |
+
key = f"{key}_{seen[key]}"
|
| 582 |
+
else:
|
| 583 |
+
seen[key] = 0
|
| 584 |
+
headers.append(key)
|
| 585 |
+
|
| 586 |
+
return [dict(zip(headers, row)) for row in grid[1:]]
|
| 587 |
+
|
| 588 |
+
|
| 589 |
+
def convert_html_table_to_json(html_content, records=True):
|
| 590 |
+
"""
|
| 591 |
+
Convert HTML tables produced by ``latex_table_to_html`` into JSON.
|
| 592 |
+
|
| 593 |
+
Parameters
|
| 594 |
+
----------
|
| 595 |
+
html_content : str
|
| 596 |
+
HTML string (may contain one or more ``<table>`` elements and
|
| 597 |
+
surrounding non-table text).
|
| 598 |
+
records : bool
|
| 599 |
+
If *True* (default) the first row is treated as a header and each
|
| 600 |
+
subsequent row becomes a ``{header: value, ...}`` dict.
|
| 601 |
+
If *False* the table is returned as a plain list-of-lists.
|
| 602 |
+
|
| 603 |
+
Returns
|
| 604 |
+
-------
|
| 605 |
+
str
|
| 606 |
+
A JSON string. If the input contains a single table the top-level
|
| 607 |
+
value is a list (of dicts or lists). If there are multiple tables
|
| 608 |
+
the top-level value is a list of such lists.
|
| 609 |
+
"""
|
| 610 |
+
soup = BeautifulSoup(html_content, "html.parser")
|
| 611 |
+
tables = soup.find_all("table")
|
| 612 |
+
|
| 613 |
+
if not tables:
|
| 614 |
+
return html_content
|
| 615 |
+
|
| 616 |
+
results = []
|
| 617 |
+
for table in tables:
|
| 618 |
+
grid = _expand_spans(table)
|
| 619 |
+
if records:
|
| 620 |
+
results.append(_grid_to_records(grid))
|
| 621 |
+
else:
|
| 622 |
+
results.append(grid)
|
| 623 |
+
|
| 624 |
+
payload = results[0] if len(results) == 1 else results
|
| 625 |
+
return json.dumps(payload, ensure_ascii=False, indent=2)
|
| 626 |
+
|
| 627 |
+
|
| 628 |
+
def convert_html_table_to_csv(html_content):
|
| 629 |
+
"""
|
| 630 |
+
Convert HTML tables produced by ``latex_table_to_html`` into CSV.
|
| 631 |
+
|
| 632 |
+
Each table becomes a CSV block separated by a blank line.
|
| 633 |
+
|
| 634 |
+
Returns
|
| 635 |
+
-------
|
| 636 |
+
str
|
| 637 |
+
CSV-formatted string.
|
| 638 |
+
"""
|
| 639 |
+
soup = BeautifulSoup(html_content, "html.parser")
|
| 640 |
+
tables = soup.find_all("table")
|
| 641 |
+
|
| 642 |
+
if not tables:
|
| 643 |
+
return html_content
|
| 644 |
+
|
| 645 |
+
csv_blocks = []
|
| 646 |
+
for table in tables:
|
| 647 |
+
grid = _expand_spans(table)
|
| 648 |
+
buf = io.StringIO()
|
| 649 |
+
writer = csv.writer(buf)
|
| 650 |
+
writer.writerows(grid)
|
| 651 |
+
csv_blocks.append(buf.getvalue().rstrip("\r\n"))
|
| 652 |
+
|
| 653 |
+
return "\n\n".join(csv_blocks)
|
nemotron_parse/engine/convert/rich_table_fallback.py
ADDED
|
@@ -0,0 +1,105 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
| 2 |
+
# SPDX-License-Identifier: OpenMDW-1.1
|
| 3 |
+
|
| 4 |
+
"""Decode known rich inline commands for the legacy tabular fallback.
|
| 5 |
+
|
| 6 |
+
Used only after strict rich-grid parsing fails. This preserves unknown commands
|
| 7 |
+
and malformed arguments; it never invents cell contents or repairs table spans.
|
| 8 |
+
"""
|
| 9 |
+
import html
|
| 10 |
+
import re
|
| 11 |
+
|
| 12 |
+
try:
|
| 13 |
+
from .hybrid_cleanup import decode_styles
|
| 14 |
+
except ImportError:
|
| 15 |
+
from hybrid_cleanup import decode_styles
|
| 16 |
+
|
| 17 |
+
try:
|
| 18 |
+
from .rich_table_reader import TAGS, LITERALS, argument
|
| 19 |
+
except ImportError:
|
| 20 |
+
from rich_table_reader import TAGS, LITERALS, argument
|
| 21 |
+
|
| 22 |
+
|
| 23 |
+
def decode_known_inline(text, _depth=0):
|
| 24 |
+
"""Decode balanced, recognized inline arguments without exposing delimiters."""
|
| 25 |
+
if _depth > 64:
|
| 26 |
+
return text
|
| 27 |
+
out = []
|
| 28 |
+
i = 0
|
| 29 |
+
known = set(TAGS) | set(LITERALS) | {"thead", "linebreak", "mbox", "hl", "colorbox"}
|
| 30 |
+
while i < len(text):
|
| 31 |
+
# Preserve escaped literal characters. In particular, an escaped
|
| 32 |
+
# backslash followed by 'thead' must not become a formatting command.
|
| 33 |
+
if text[i] == "\\" and i + 1 < len(text) and not text[i + 1].isalpha():
|
| 34 |
+
literal = text[i + 1]
|
| 35 |
+
out.append(
|
| 36 |
+
"\\&#" + str(ord(literal)) + ";"
|
| 37 |
+
if literal in "&%$#_*{}"
|
| 38 |
+
else text[i : i + 2]
|
| 39 |
+
)
|
| 40 |
+
i += 2
|
| 41 |
+
continue
|
| 42 |
+
match = re.match(r"\\([A-Za-z]+)\{", text[i:])
|
| 43 |
+
if not match or match[1] not in known:
|
| 44 |
+
out.append(text[i])
|
| 45 |
+
i += 1
|
| 46 |
+
continue
|
| 47 |
+
command = match[1]
|
| 48 |
+
try:
|
| 49 |
+
body, end = argument(text, i + len(command) + 1)
|
| 50 |
+
except ValueError:
|
| 51 |
+
out.append(text[i])
|
| 52 |
+
i += 1
|
| 53 |
+
continue
|
| 54 |
+
if command in {"hl", "colorbox"}:
|
| 55 |
+
try:
|
| 56 |
+
from .semantic_styles import mark_open
|
| 57 |
+
except ImportError:
|
| 58 |
+
from semantic_styles import mark_open
|
| 59 |
+
from bs4 import BeautifulSoup
|
| 60 |
+
|
| 61 |
+
node = BeautifulSoup("<mark></mark>", "html.parser").mark
|
| 62 |
+
if command == "colorbox":
|
| 63 |
+
node["background"] = body
|
| 64 |
+
try:
|
| 65 |
+
body, end = argument(text, end)
|
| 66 |
+
except ValueError:
|
| 67 |
+
out.append(text[i:end])
|
| 68 |
+
i = end
|
| 69 |
+
continue
|
| 70 |
+
out.append(
|
| 71 |
+
mark_open(node) + decode_known_inline(body, _depth + 1) + "</mark>"
|
| 72 |
+
)
|
| 73 |
+
elif command in TAGS or command == "thead":
|
| 74 |
+
inner = decode_known_inline(body, _depth + 1)
|
| 75 |
+
tag = "nvp-header-cell" if command == "thead" else TAGS[command]
|
| 76 |
+
out.append("<" + tag + ">" + inner + "</" + tag + ">")
|
| 77 |
+
elif command in LITERALS and not body:
|
| 78 |
+
# Entities cannot be mistaken for LaTeX row/column delimiters.
|
| 79 |
+
out.append("\\&#" + str(ord(LITERALS[command])) + ";")
|
| 80 |
+
elif command == "linebreak" and not body:
|
| 81 |
+
out.append("<nvp-line-break></nvp-line-break>")
|
| 82 |
+
elif command == "mbox" and not body:
|
| 83 |
+
pass
|
| 84 |
+
else:
|
| 85 |
+
out.append(text[i:end])
|
| 86 |
+
i = end
|
| 87 |
+
return decode_styles("".join(out))
|
| 88 |
+
|
| 89 |
+
|
| 90 |
+
def restore_header_cells(html_text):
|
| 91 |
+
from bs4 import BeautifulSoup
|
| 92 |
+
|
| 93 |
+
# Restore explicit breaks after legacy whitespace normalization, which
|
| 94 |
+
# otherwise collapses consecutive linebreak commands.
|
| 95 |
+
html_text = html_text.replace("<nvp-line-break></nvp-line-break>", "<br>")
|
| 96 |
+
if "<nvp-header-cell>" not in html_text:
|
| 97 |
+
return html_text
|
| 98 |
+
soup = BeautifulSoup(html_text, "html.parser")
|
| 99 |
+
for cell in soup.find_all(["td", "th"]):
|
| 100 |
+
markers = cell.find_all("nvp-header-cell")
|
| 101 |
+
if markers:
|
| 102 |
+
cell.name = "th"
|
| 103 |
+
for marker in markers:
|
| 104 |
+
marker.unwrap()
|
| 105 |
+
return str(soup)
|
nemotron_parse/engine/convert/rich_table_reader.py
ADDED
|
@@ -0,0 +1,164 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
| 2 |
+
# SPDX-License-Identifier: OpenMDW-1.1
|
| 3 |
+
|
| 4 |
+
"""Strict rich_tabular_v3 reverse reader. Other LaTeX uses the legacy parser."""
|
| 5 |
+
import html, re
|
| 6 |
+
|
| 7 |
+
TAGS = {
|
| 8 |
+
"textstrong": "b",
|
| 9 |
+
"textem": "i",
|
| 10 |
+
"textbf": "b",
|
| 11 |
+
"textit": "i",
|
| 12 |
+
"underline": "u",
|
| 13 |
+
"sout": "s",
|
| 14 |
+
"textsuperscript": "sup",
|
| 15 |
+
"textsubscript": "sub",
|
| 16 |
+
}
|
| 17 |
+
|
| 18 |
+
LITERALS = {
|
| 19 |
+
"textbackslash": "\\",
|
| 20 |
+
"textasciitilde": "~",
|
| 21 |
+
"textasciicircum": "^",
|
| 22 |
+
"textless": "<",
|
| 23 |
+
"textgreater": ">",
|
| 24 |
+
}
|
| 25 |
+
|
| 26 |
+
|
| 27 |
+
def split_top(text, delimiter):
|
| 28 |
+
parts = []
|
| 29 |
+
start = 0
|
| 30 |
+
depth = 0
|
| 31 |
+
i = 0
|
| 32 |
+
while i < len(text):
|
| 33 |
+
if depth == 0 and text.startswith(delimiter, i):
|
| 34 |
+
parts.append(text[start:i])
|
| 35 |
+
i += len(delimiter)
|
| 36 |
+
start = i
|
| 37 |
+
continue
|
| 38 |
+
if text[i] == "\\":
|
| 39 |
+
i += 2
|
| 40 |
+
continue
|
| 41 |
+
if text[i] == "{":
|
| 42 |
+
depth += 1
|
| 43 |
+
elif text[i] == "}":
|
| 44 |
+
depth -= 1
|
| 45 |
+
if depth < 0:
|
| 46 |
+
raise ValueError("unbalanced_generated_latex")
|
| 47 |
+
i += 1
|
| 48 |
+
if depth:
|
| 49 |
+
raise ValueError("unbalanced_generated_latex")
|
| 50 |
+
parts.append(text[start:])
|
| 51 |
+
return parts
|
| 52 |
+
|
| 53 |
+
|
| 54 |
+
def argument(text, pos):
|
| 55 |
+
if pos >= len(text) or text[pos] != "{":
|
| 56 |
+
raise ValueError("expected_braced_argument")
|
| 57 |
+
start = pos + 1
|
| 58 |
+
depth = 1
|
| 59 |
+
pos += 1
|
| 60 |
+
while pos < len(text):
|
| 61 |
+
if text[pos] == "\\":
|
| 62 |
+
pos += 2
|
| 63 |
+
continue
|
| 64 |
+
if text[pos] == "{":
|
| 65 |
+
depth += 1
|
| 66 |
+
elif text[pos] == "}":
|
| 67 |
+
depth -= 1
|
| 68 |
+
if depth == 0:
|
| 69 |
+
return text[start:pos], pos + 1
|
| 70 |
+
pos += 1
|
| 71 |
+
raise ValueError("unbalanced_generated_latex")
|
| 72 |
+
|
| 73 |
+
|
| 74 |
+
def inline_html(text):
|
| 75 |
+
try:
|
| 76 |
+
from .hybrid_inline import inline_html as decode
|
| 77 |
+
except ImportError:
|
| 78 |
+
from hybrid_inline import inline_html as decode
|
| 79 |
+
return decode(text)
|
| 80 |
+
|
| 81 |
+
|
| 82 |
+
def cell_parts(text):
|
| 83 |
+
text = text.strip()
|
| 84 |
+
rs = cs = 1
|
| 85 |
+
while text.startswith((r"\multicolumn", r"\multirow")):
|
| 86 |
+
command = "multicolumn" if text.startswith(r"\multicolumn") else "multirow"
|
| 87 |
+
pos = len(command) + 1
|
| 88 |
+
count, pos = argument(text, pos)
|
| 89 |
+
_, pos = argument(text, pos)
|
| 90 |
+
body, pos = argument(text, pos)
|
| 91 |
+
if text[pos:].strip():
|
| 92 |
+
raise ValueError("trailing_span_content")
|
| 93 |
+
if command == "multicolumn":
|
| 94 |
+
cs = int(count)
|
| 95 |
+
else:
|
| 96 |
+
rs = int(count)
|
| 97 |
+
text = body
|
| 98 |
+
header = False
|
| 99 |
+
if text.startswith(r"\thead{"):
|
| 100 |
+
body, end = argument(text, len(r"\thead"))
|
| 101 |
+
if text[end:].strip():
|
| 102 |
+
raise ValueError("trailing_header_content")
|
| 103 |
+
text = body
|
| 104 |
+
header = True
|
| 105 |
+
return rs, cs, inline_html(text), header
|
| 106 |
+
|
| 107 |
+
|
| 108 |
+
def to_html(latex):
|
| 109 |
+
m = re.fullmatch(
|
| 110 |
+
r"\\begin\{tabular\}\{[clrx]+\}\n(.*?)\n\\end\{tabular\}", latex, re.S
|
| 111 |
+
)
|
| 112 |
+
if not m:
|
| 113 |
+
raise ValueError("unsupported_generated_tabular")
|
| 114 |
+
rows = split_top(m.group(1), r"\\")
|
| 115 |
+
if rows[-1].strip():
|
| 116 |
+
raise ValueError("missing_final_row_separator")
|
| 117 |
+
columns = len(re.match(r"\\begin\{tabular\}\{([clrx]+)\}", latex).group(1))
|
| 118 |
+
rows.pop()
|
| 119 |
+
occupied = set()
|
| 120 |
+
out = ["<table>"]
|
| 121 |
+
if not rows or len(rows) > 4096 or columns > 512:
|
| 122 |
+
raise ValueError("generated_grid_size")
|
| 123 |
+
for r, row in enumerate(rows):
|
| 124 |
+
out.append("<tr>")
|
| 125 |
+
c = 0
|
| 126 |
+
for part in split_top(row, "&"):
|
| 127 |
+
rs, cs, content, header = cell_parts(part)
|
| 128 |
+
if rs < 1 or cs < 1 or r + rs > len(rows) or c + cs > columns:
|
| 129 |
+
raise ValueError("invalid_generated_span")
|
| 130 |
+
if (r, c) in occupied:
|
| 131 |
+
if content or header or rs != 1:
|
| 132 |
+
raise ValueError("nonempty_span_placeholder")
|
| 133 |
+
if any((r, cc) not in occupied for cc in range(c, c + cs)):
|
| 134 |
+
raise ValueError("invalid_span_placeholder")
|
| 135 |
+
c += cs
|
| 136 |
+
continue
|
| 137 |
+
if any(
|
| 138 |
+
(rr, cc) in occupied
|
| 139 |
+
for rr in range(r, r + rs)
|
| 140 |
+
for cc in range(c, c + cs)
|
| 141 |
+
):
|
| 142 |
+
raise ValueError("overlapping_generated_span")
|
| 143 |
+
tag = "th" if header else "td"
|
| 144 |
+
out.append(
|
| 145 |
+
"<"
|
| 146 |
+
+ tag
|
| 147 |
+
+ ' rowspan="'
|
| 148 |
+
+ str(rs)
|
| 149 |
+
+ '" colspan="'
|
| 150 |
+
+ str(cs)
|
| 151 |
+
+ '">'
|
| 152 |
+
+ content
|
| 153 |
+
+ "</"
|
| 154 |
+
+ tag
|
| 155 |
+
+ ">"
|
| 156 |
+
)
|
| 157 |
+
for rr in range(r, r + rs):
|
| 158 |
+
for cc in range(c, c + cs):
|
| 159 |
+
occupied.add((rr, cc))
|
| 160 |
+
c += cs
|
| 161 |
+
if c != columns:
|
| 162 |
+
raise ValueError("generated_column_count")
|
| 163 |
+
out.append("</tr>")
|
| 164 |
+
return "".join(out) + "</table>"
|
nemotron_parse/engine/convert/semantic_styles.py
ADDED
|
@@ -0,0 +1,68 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
| 2 |
+
# SPDX-License-Identifier: OpenMDW-1.1
|
| 3 |
+
|
| 4 |
+
"""Lossless, bounded semantic inline styles shared by table and page conversion."""
|
| 5 |
+
import html
|
| 6 |
+
import re
|
| 7 |
+
from bs4 import BeautifulSoup
|
| 8 |
+
|
| 9 |
+
ALIASES = {
|
| 10 |
+
"strong": "b",
|
| 11 |
+
"em": "i",
|
| 12 |
+
"ins": "u",
|
| 13 |
+
"underline": "u",
|
| 14 |
+
"del": "s",
|
| 15 |
+
"strike": "s",
|
| 16 |
+
}
|
| 17 |
+
|
| 18 |
+
|
| 19 |
+
# Highlight color is deliberately outside the training target contract.
|
| 20 |
+
def mark_color(node):
|
| 21 |
+
return None
|
| 22 |
+
|
| 23 |
+
|
| 24 |
+
def mark_open(node):
|
| 25 |
+
return "<mark>"
|
| 26 |
+
|
| 27 |
+
|
| 28 |
+
def normalize_semantic_html(value):
|
| 29 |
+
soup = BeautifulSoup(str(value), "html.parser")
|
| 30 |
+
for node in list(soup.find_all(True)):
|
| 31 |
+
node.name = ALIASES.get(node.name, node.name)
|
| 32 |
+
if node.name == "mark":
|
| 33 |
+
node.attrs = {}
|
| 34 |
+
if node.name != "span":
|
| 35 |
+
continue
|
| 36 |
+
props = {
|
| 37 |
+
k.strip().lower(): v.strip().lower()
|
| 38 |
+
for part in str(node.get("style", "")).split(";")
|
| 39 |
+
for k, sep, v in [part.partition(":")]
|
| 40 |
+
if sep
|
| 41 |
+
}
|
| 42 |
+
styles = []
|
| 43 |
+
if props.get("font-weight") in ["bold", "bolder", "600", "700", "800", "900"]:
|
| 44 |
+
styles.append("b")
|
| 45 |
+
if props.get("font-style") in ["italic", "oblique"]:
|
| 46 |
+
styles.append("i")
|
| 47 |
+
decoration = (
|
| 48 |
+
props.get("text-decoration", "")
|
| 49 |
+
+ " "
|
| 50 |
+
+ props.get("text-decoration-line", "")
|
| 51 |
+
)
|
| 52 |
+
if "underline" in decoration:
|
| 53 |
+
styles.append("u")
|
| 54 |
+
if "line-through" in decoration:
|
| 55 |
+
styles.append("s")
|
| 56 |
+
if props.get("vertical-align") in ["super", "sub"]:
|
| 57 |
+
styles.append("sup" if props["vertical-align"] == "super" else "sub")
|
| 58 |
+
if props.get("background-color") or props.get("background"):
|
| 59 |
+
styles.append("mark")
|
| 60 |
+
if not styles:
|
| 61 |
+
continue
|
| 62 |
+
for name in reversed(styles):
|
| 63 |
+
inner = soup.new_tag(name)
|
| 64 |
+
for child in list(node.contents):
|
| 65 |
+
inner.append(child.extract())
|
| 66 |
+
node.append(inner)
|
| 67 |
+
node.unwrap()
|
| 68 |
+
return soup
|
nemotron_parse/engine/geometry.py
ADDED
|
@@ -0,0 +1,438 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
| 2 |
+
# SPDX-License-Identifier: OpenMDW-1.1
|
| 3 |
+
|
| 4 |
+
from __future__ import annotations
|
| 5 |
+
|
| 6 |
+
import argparse
|
| 7 |
+
|
| 8 |
+
|
| 9 |
+
import base64
|
| 10 |
+
|
| 11 |
+
|
| 12 |
+
import json
|
| 13 |
+
|
| 14 |
+
|
| 15 |
+
import math
|
| 16 |
+
|
| 17 |
+
|
| 18 |
+
import mimetypes
|
| 19 |
+
|
| 20 |
+
|
| 21 |
+
import re
|
| 22 |
+
|
| 23 |
+
|
| 24 |
+
from dataclasses import dataclass
|
| 25 |
+
|
| 26 |
+
|
| 27 |
+
from pathlib import Path
|
| 28 |
+
|
| 29 |
+
|
| 30 |
+
from typing import Any
|
| 31 |
+
|
| 32 |
+
|
| 33 |
+
import numpy as np
|
| 34 |
+
|
| 35 |
+
|
| 36 |
+
from PIL import Image
|
| 37 |
+
|
| 38 |
+
|
| 39 |
+
BLOCK_RE = re.compile(
|
| 40 |
+
r"<x_(\d+(?:\.\d+)?)><y_(\d+(?:\.\d+)?)>(.*?)"
|
| 41 |
+
r"<x_(\d+(?:\.\d+)?)><y_(\d+(?:\.\d+)?)><class_([^>]+)>",
|
| 42 |
+
re.MULTILINE | re.DOTALL,
|
| 43 |
+
)
|
| 44 |
+
|
| 45 |
+
|
| 46 |
+
@dataclass(frozen=True)
|
| 47 |
+
class PageImage:
|
| 48 |
+
key: str
|
| 49 |
+
image_name: str
|
| 50 |
+
image_path: Path
|
| 51 |
+
width: int
|
| 52 |
+
height: int
|
| 53 |
+
category: str | None = None
|
| 54 |
+
|
| 55 |
+
|
| 56 |
+
@dataclass(frozen=True)
|
| 57 |
+
class CropSpec:
|
| 58 |
+
crop_id: str
|
| 59 |
+
page_key: str
|
| 60 |
+
image_name: str
|
| 61 |
+
page_image_path: str
|
| 62 |
+
crop_image_name: str
|
| 63 |
+
crop_image_path: str
|
| 64 |
+
element_index: int
|
| 65 |
+
detector_label: str
|
| 66 |
+
detector_score: float
|
| 67 |
+
detector_bbox_norm: tuple[float, float, float, float]
|
| 68 |
+
crop_bbox_pixels: tuple[int, int, int, int]
|
| 69 |
+
page_size: tuple[int, int]
|
| 70 |
+
|
| 71 |
+
|
| 72 |
+
def write_jsonl(path: Path, rows: list[dict[str, Any]]) -> None:
|
| 73 |
+
path.parent.mkdir(parents=True, exist_ok=True)
|
| 74 |
+
with path.open("w", encoding="utf-8") as f:
|
| 75 |
+
for row in rows:
|
| 76 |
+
f.write(json.dumps(row, ensure_ascii=False) + "\n")
|
| 77 |
+
|
| 78 |
+
|
| 79 |
+
def normalize_detector_label(label: Any, labels: Any = None) -> str:
|
| 80 |
+
if isinstance(label, (int, float, np.integer, np.floating)) and labels is not None:
|
| 81 |
+
try:
|
| 82 |
+
label_idx = int(label)
|
| 83 |
+
if isinstance(labels, dict):
|
| 84 |
+
label = labels.get(label_idx, labels.get(str(label_idx), label))
|
| 85 |
+
else:
|
| 86 |
+
label = labels[label_idx]
|
| 87 |
+
except Exception:
|
| 88 |
+
pass
|
| 89 |
+
text = str(label).strip().lower()
|
| 90 |
+
text = (
|
| 91 |
+
text.replace("&", "and").replace("/", "_").replace("-", "_").replace(" ", "_")
|
| 92 |
+
)
|
| 93 |
+
text = text.replace("__", "_")
|
| 94 |
+
text = re.sub(r"_+", "_", text).strip("_")
|
| 95 |
+
if text in {
|
| 96 |
+
"headers_and_footers",
|
| 97 |
+
"header_and_footer",
|
| 98 |
+
"header_footer",
|
| 99 |
+
"headers_footers",
|
| 100 |
+
}:
|
| 101 |
+
return "header_footer"
|
| 102 |
+
if text in {
|
| 103 |
+
"header",
|
| 104 |
+
"footer",
|
| 105 |
+
"page_header",
|
| 106 |
+
"page_footer",
|
| 107 |
+
"header_image",
|
| 108 |
+
"footer_image",
|
| 109 |
+
}:
|
| 110 |
+
return "header_footer"
|
| 111 |
+
if text in {"tables"}:
|
| 112 |
+
return "table"
|
| 113 |
+
if text in {"charts"}:
|
| 114 |
+
return "chart"
|
| 115 |
+
if text in {
|
| 116 |
+
"texts",
|
| 117 |
+
"paragraph",
|
| 118 |
+
"paragraphs",
|
| 119 |
+
"paragraph_title",
|
| 120 |
+
"content",
|
| 121 |
+
"abstract",
|
| 122 |
+
"references",
|
| 123 |
+
"reference",
|
| 124 |
+
"reference_content",
|
| 125 |
+
"footnote",
|
| 126 |
+
"page_number",
|
| 127 |
+
"number",
|
| 128 |
+
"formula",
|
| 129 |
+
"formula_number",
|
| 130 |
+
"inline_formula",
|
| 131 |
+
"display_formula",
|
| 132 |
+
"algorithm",
|
| 133 |
+
"aside_text",
|
| 134 |
+
"sidebar_text",
|
| 135 |
+
"list",
|
| 136 |
+
"lists",
|
| 137 |
+
"vertical_text",
|
| 138 |
+
}:
|
| 139 |
+
return "text"
|
| 140 |
+
if text in {
|
| 141 |
+
"title",
|
| 142 |
+
"titles",
|
| 143 |
+
"document_title",
|
| 144 |
+
"section_header",
|
| 145 |
+
"figure_title",
|
| 146 |
+
"table_title",
|
| 147 |
+
"chart_title",
|
| 148 |
+
"figure_table_title",
|
| 149 |
+
"figure_caption",
|
| 150 |
+
"table_caption",
|
| 151 |
+
"image_caption",
|
| 152 |
+
}:
|
| 153 |
+
return "title"
|
| 154 |
+
if text in {
|
| 155 |
+
"image",
|
| 156 |
+
"images",
|
| 157 |
+
"figure",
|
| 158 |
+
"figures",
|
| 159 |
+
"picture",
|
| 160 |
+
"pictures",
|
| 161 |
+
"seal",
|
| 162 |
+
"stamp",
|
| 163 |
+
}:
|
| 164 |
+
return "infographic"
|
| 165 |
+
if text in {"infographics"}:
|
| 166 |
+
return "infographic"
|
| 167 |
+
return text
|
| 168 |
+
|
| 169 |
+
|
| 170 |
+
def parse_label_filter(value: str) -> set[str]:
|
| 171 |
+
if value.strip().lower() in {"", "*", "all"}:
|
| 172 |
+
return set()
|
| 173 |
+
return {normalize_detector_label(item) for item in value.split(",") if item.strip()}
|
| 174 |
+
|
| 175 |
+
|
| 176 |
+
def clamp(value: float, low: float = 0.0, high: float = 1.0) -> float:
|
| 177 |
+
return min(high, max(low, value))
|
| 178 |
+
|
| 179 |
+
|
| 180 |
+
def crop_bbox_pixels(
|
| 181 |
+
bbox_norm: tuple[float, float, float, float],
|
| 182 |
+
width: int,
|
| 183 |
+
height: int,
|
| 184 |
+
*,
|
| 185 |
+
margin_px: int,
|
| 186 |
+
margin_ratio: float,
|
| 187 |
+
) -> tuple[int, int, int, int]:
|
| 188 |
+
x1, y1, x2, y2 = bbox_norm
|
| 189 |
+
px1, py1, px2, py2 = x1 * width, y1 * height, x2 * width, y2 * height
|
| 190 |
+
mx = margin_px + (px2 - px1) * margin_ratio
|
| 191 |
+
my = margin_px + (py2 - py1) * margin_ratio
|
| 192 |
+
left = max(0, math.floor(px1 - mx))
|
| 193 |
+
top = max(0, math.floor(py1 - my))
|
| 194 |
+
right = min(width, math.ceil(px2 + mx))
|
| 195 |
+
bottom = min(height, math.ceil(py2 + my))
|
| 196 |
+
return left, top, right, bottom
|
| 197 |
+
|
| 198 |
+
|
| 199 |
+
def filter_detection_rows(
|
| 200 |
+
args: argparse.Namespace, rows: list[dict[str, Any]]
|
| 201 |
+
) -> list[dict[str, Any]]:
|
| 202 |
+
include_labels = parse_label_filter(args.include_detector_labels)
|
| 203 |
+
filtered_rows = []
|
| 204 |
+
for row in rows:
|
| 205 |
+
detections = []
|
| 206 |
+
for det in row.get("detections") or []:
|
| 207 |
+
label = normalize_detector_label(det.get("label") or "")
|
| 208 |
+
if include_labels and label not in include_labels:
|
| 209 |
+
continue
|
| 210 |
+
if float(det.get("score") or 0.0) < args.detector_score_min:
|
| 211 |
+
continue
|
| 212 |
+
kept = dict(det)
|
| 213 |
+
kept["label"] = label
|
| 214 |
+
detections.append(kept)
|
| 215 |
+
if args.detector_order == "spatial":
|
| 216 |
+
detections.sort(
|
| 217 |
+
key=lambda item: (
|
| 218 |
+
item["bbox_norm"][1],
|
| 219 |
+
item["bbox_norm"][0],
|
| 220 |
+
-float(item.get("score") or 0.0),
|
| 221 |
+
)
|
| 222 |
+
)
|
| 223 |
+
if args.max_elements_per_page > 0:
|
| 224 |
+
detections = detections[: args.max_elements_per_page]
|
| 225 |
+
out = dict(row)
|
| 226 |
+
out["detections"] = detections
|
| 227 |
+
filtered_rows.append(out)
|
| 228 |
+
return filtered_rows
|
| 229 |
+
|
| 230 |
+
|
| 231 |
+
def make_fallback_detection(page: PageImage) -> dict[str, Any]:
|
| 232 |
+
return {
|
| 233 |
+
"label": "text",
|
| 234 |
+
"score": 0.0,
|
| 235 |
+
"bbox_norm": [0.0, 0.0, 1.0, 1.0],
|
| 236 |
+
"fallback": True,
|
| 237 |
+
}
|
| 238 |
+
|
| 239 |
+
|
| 240 |
+
def build_crops(
|
| 241 |
+
args: argparse.Namespace,
|
| 242 |
+
pages: list[PageImage],
|
| 243 |
+
detections: list[dict[str, Any]],
|
| 244 |
+
crops_dir: Path,
|
| 245 |
+
crops_manifest_path: Path,
|
| 246 |
+
) -> list[CropSpec]:
|
| 247 |
+
pages_by_name = {page.image_name: page for page in pages}
|
| 248 |
+
pages_by_key = {page.key: page for page in pages}
|
| 249 |
+
crops: list[CropSpec] = []
|
| 250 |
+
for row in detections:
|
| 251 |
+
page = pages_by_key.get(str(row.get("key"))) or pages_by_name.get(
|
| 252 |
+
str(row.get("image_name"))
|
| 253 |
+
)
|
| 254 |
+
if page is None:
|
| 255 |
+
continue
|
| 256 |
+
items = list(row.get("detections") or [])
|
| 257 |
+
if not items and args.fallback_full_page:
|
| 258 |
+
items = [make_fallback_detection(page)]
|
| 259 |
+
if not items:
|
| 260 |
+
continue
|
| 261 |
+
with Image.open(page.image_path) as image:
|
| 262 |
+
image = image.convert("RGB")
|
| 263 |
+
for element_index, det in enumerate(items):
|
| 264 |
+
bbox_norm = tuple(float(v) for v in det["bbox_norm"])
|
| 265 |
+
crop_px = crop_bbox_pixels(
|
| 266 |
+
bbox_norm,
|
| 267 |
+
page.width,
|
| 268 |
+
page.height,
|
| 269 |
+
margin_px=args.crop_margin_px,
|
| 270 |
+
margin_ratio=args.crop_margin_ratio,
|
| 271 |
+
)
|
| 272 |
+
left, top, right, bottom = crop_px
|
| 273 |
+
if (
|
| 274 |
+
right - left < args.min_crop_size
|
| 275 |
+
or bottom - top < args.min_crop_size
|
| 276 |
+
):
|
| 277 |
+
continue
|
| 278 |
+
crop_image_name = f"{page.key}__pe_{element_index:03d}.jpg"
|
| 279 |
+
crop_image_path = crops_dir / crop_image_name
|
| 280 |
+
crop_image_path.parent.mkdir(parents=True, exist_ok=True)
|
| 281 |
+
if not crop_image_path.is_file() or crop_image_path.stat().st_size == 0:
|
| 282 |
+
image.crop((left, top, right, bottom)).save(
|
| 283 |
+
crop_image_path, quality=95
|
| 284 |
+
)
|
| 285 |
+
crops.append(
|
| 286 |
+
CropSpec(
|
| 287 |
+
crop_id=f"{page.key}__pe_{element_index:03d}",
|
| 288 |
+
page_key=page.key,
|
| 289 |
+
image_name=page.image_name,
|
| 290 |
+
page_image_path=str(page.image_path),
|
| 291 |
+
crop_image_name=crop_image_name,
|
| 292 |
+
crop_image_path=str(crop_image_path),
|
| 293 |
+
element_index=element_index,
|
| 294 |
+
detector_label=str(det.get("label") or "text"),
|
| 295 |
+
detector_score=float(det.get("score") or 0.0),
|
| 296 |
+
detector_bbox_norm=bbox_norm,
|
| 297 |
+
crop_bbox_pixels=crop_px,
|
| 298 |
+
page_size=(page.width, page.height),
|
| 299 |
+
)
|
| 300 |
+
)
|
| 301 |
+
write_jsonl(crops_manifest_path, [crop.__dict__ for crop in crops])
|
| 302 |
+
print(f"Prepared {len(crops)} crops under {crops_dir}", flush=True)
|
| 303 |
+
return crops
|
| 304 |
+
|
| 305 |
+
|
| 306 |
+
def image_data_url(path: Path) -> str:
|
| 307 |
+
mime = mimetypes.guess_type(path.name)[0] or "image/jpeg"
|
| 308 |
+
data = base64.b64encode(path.read_bytes()).decode("ascii")
|
| 309 |
+
return f"data:{mime};base64,{data}"
|
| 310 |
+
|
| 311 |
+
|
| 312 |
+
def parse_model_detection_predictions(
|
| 313 |
+
args: argparse.Namespace,
|
| 314 |
+
pages: list[PageImage],
|
| 315 |
+
prediction_rows: list[dict[str, Any]],
|
| 316 |
+
) -> list[dict[str, Any]]:
|
| 317 |
+
page_by_key = {page.key: page for page in pages}
|
| 318 |
+
rows = []
|
| 319 |
+
for prediction in prediction_rows:
|
| 320 |
+
page = page_by_key.get(str(prediction.get("key")))
|
| 321 |
+
if page is None:
|
| 322 |
+
continue
|
| 323 |
+
detections = []
|
| 324 |
+
for match in BLOCK_RE.finditer(str(prediction.get("raw_out") or "")):
|
| 325 |
+
x1, y1, _body, x2, y2, cls = match.groups()
|
| 326 |
+
bbox_norm = model_bbox_to_image_norm(
|
| 327 |
+
(float(x1), float(y1), float(x2), float(y2)),
|
| 328 |
+
page.width,
|
| 329 |
+
page.height,
|
| 330 |
+
coordinate_frame=args.model_coordinate_frame,
|
| 331 |
+
canvas_width=args.canvas_width,
|
| 332 |
+
canvas_height=args.canvas_height,
|
| 333 |
+
pad_position=args.pad_position,
|
| 334 |
+
)
|
| 335 |
+
if bbox_norm[2] <= bbox_norm[0] or bbox_norm[3] <= bbox_norm[1]:
|
| 336 |
+
continue
|
| 337 |
+
detections.append(
|
| 338 |
+
{
|
| 339 |
+
"label": normalize_detector_label(cls),
|
| 340 |
+
"score": 1.0,
|
| 341 |
+
"bbox_norm": list(bbox_norm),
|
| 342 |
+
"model_class": cls,
|
| 343 |
+
"model_order": len(detections),
|
| 344 |
+
}
|
| 345 |
+
)
|
| 346 |
+
rows.append(
|
| 347 |
+
{
|
| 348 |
+
"key": page.key,
|
| 349 |
+
"image_name": page.image_name,
|
| 350 |
+
"image_path": str(page.image_path),
|
| 351 |
+
"page_size": [page.width, page.height],
|
| 352 |
+
"detector_backend": "model-detection",
|
| 353 |
+
"model_coordinate_frame": args.model_coordinate_frame,
|
| 354 |
+
"detections": detections,
|
| 355 |
+
"detection_error": prediction.get("error"),
|
| 356 |
+
}
|
| 357 |
+
)
|
| 358 |
+
return filter_detection_rows(args, rows)
|
| 359 |
+
|
| 360 |
+
|
| 361 |
+
def model_canvas_norm_to_image_norm(
|
| 362 |
+
bbox: tuple[float, float, float, float],
|
| 363 |
+
width: int,
|
| 364 |
+
height: int,
|
| 365 |
+
*,
|
| 366 |
+
canvas_width: float,
|
| 367 |
+
canvas_height: float,
|
| 368 |
+
pad_position: str,
|
| 369 |
+
) -> tuple[float, float, float, float]:
|
| 370 |
+
x1, y1, x2, y2 = bbox
|
| 371 |
+
max_coord = max(abs(x1), abs(y1), abs(x2), abs(y2))
|
| 372 |
+
if max_coord > 1.5 and max_coord <= 1000.0:
|
| 373 |
+
norm = (x1 / 1000.0, y1 / 1000.0, x2 / 1000.0, y2 / 1000.0)
|
| 374 |
+
elif max_coord > 1000.0:
|
| 375 |
+
norm = (
|
| 376 |
+
x1 / max(1, width),
|
| 377 |
+
y1 / max(1, height),
|
| 378 |
+
x2 / max(1, width),
|
| 379 |
+
y2 / max(1, height),
|
| 380 |
+
)
|
| 381 |
+
else:
|
| 382 |
+
norm = (x1, y1, x2, y2)
|
| 383 |
+
scale = min(1.0, canvas_width / max(1, width), canvas_height / max(1, height))
|
| 384 |
+
resized_width = width * scale
|
| 385 |
+
resized_height = height * scale
|
| 386 |
+
if pad_position == "top-left":
|
| 387 |
+
pad_x = 0.0
|
| 388 |
+
pad_y = 0.0
|
| 389 |
+
elif pad_position == "bottom-right":
|
| 390 |
+
pad_x = canvas_width - resized_width
|
| 391 |
+
pad_y = canvas_height - resized_height
|
| 392 |
+
else:
|
| 393 |
+
pad_x = (canvas_width - resized_width) / 2.0
|
| 394 |
+
pad_y = (canvas_height - resized_height) / 2.0
|
| 395 |
+
out = (
|
| 396 |
+
(norm[0] * canvas_width - pad_x) / max(1.0, resized_width),
|
| 397 |
+
(norm[1] * canvas_height - pad_y) / max(1.0, resized_height),
|
| 398 |
+
(norm[2] * canvas_width - pad_x) / max(1.0, resized_width),
|
| 399 |
+
(norm[3] * canvas_height - pad_y) / max(1.0, resized_height),
|
| 400 |
+
)
|
| 401 |
+
left, right = sorted((clamp(out[0]), clamp(out[2])))
|
| 402 |
+
top, bottom = sorted((clamp(out[1]), clamp(out[3])))
|
| 403 |
+
return left, top, right, bottom
|
| 404 |
+
|
| 405 |
+
|
| 406 |
+
def model_bbox_to_image_norm(
|
| 407 |
+
bbox: tuple[float, float, float, float],
|
| 408 |
+
width: int,
|
| 409 |
+
height: int,
|
| 410 |
+
*,
|
| 411 |
+
coordinate_frame: str,
|
| 412 |
+
canvas_width: float,
|
| 413 |
+
canvas_height: float,
|
| 414 |
+
pad_position: str,
|
| 415 |
+
) -> tuple[float, float, float, float]:
|
| 416 |
+
if coordinate_frame in {
|
| 417 |
+
"content_normalized",
|
| 418 |
+
"original_normalized",
|
| 419 |
+
"navit_content",
|
| 420 |
+
}:
|
| 421 |
+
x1, y1, x2, y2 = bbox
|
| 422 |
+
max_coord = max(abs(x1), abs(y1), abs(x2), abs(y2))
|
| 423 |
+
if max_coord > 1.5 and max_coord <= 1000.0:
|
| 424 |
+
x1, y1, x2, y2 = (value / 1000.0 for value in (x1, y1, x2, y2))
|
| 425 |
+
elif max_coord > 1000.0:
|
| 426 |
+
x1, x2 = x1 / max(1, width), x2 / max(1, width)
|
| 427 |
+
y1, y2 = y1 / max(1, height), y2 / max(1, height)
|
| 428 |
+
left, right = sorted((clamp(x1), clamp(x2)))
|
| 429 |
+
top, bottom = sorted((clamp(y1), clamp(y2)))
|
| 430 |
+
return left, top, right, bottom
|
| 431 |
+
return model_canvas_norm_to_image_norm(
|
| 432 |
+
bbox,
|
| 433 |
+
width,
|
| 434 |
+
height,
|
| 435 |
+
canvas_width=canvas_width,
|
| 436 |
+
canvas_height=canvas_height,
|
| 437 |
+
pad_position=pad_position,
|
| 438 |
+
)
|
nemotron_parse/engine/header_inference_probe.py
ADDED
|
@@ -0,0 +1,84 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
| 2 |
+
# SPDX-License-Identifier: OpenMDW-1.1
|
| 3 |
+
|
| 4 |
+
"""Infer leading header rows from predicted table spans.
|
| 5 |
+
No benchmark labels, source-image metadata, or dataset IDs are accepted as input.
|
| 6 |
+
"""
|
| 7 |
+
import re
|
| 8 |
+
from bs4 import BeautifulSoup
|
| 9 |
+
|
| 10 |
+
|
| 11 |
+
def infer_header_rows(raw):
|
| 12 |
+
soup = BeautifulSoup(raw, "html.parser")
|
| 13 |
+
diagnostics = []
|
| 14 |
+
for table in soup.find_all("table"):
|
| 15 |
+
if table.find_parent("table"):
|
| 16 |
+
continue
|
| 17 |
+
rows = [r for r in table.find_all("tr") if r.find_parent("table") == table]
|
| 18 |
+
if not rows:
|
| 19 |
+
continue
|
| 20 |
+
# Respect existing explicit cell roles; this probe is only for td-only output.
|
| 21 |
+
if table.find("th"):
|
| 22 |
+
diagnostics.append({"changed": False, "reason": "existing_th"})
|
| 23 |
+
continue
|
| 24 |
+
start = 0
|
| 25 |
+
while start < min(3, len(rows) - 1):
|
| 26 |
+
cells = rows[start].find_all(["td", "th"], recursive=False)
|
| 27 |
+
if (
|
| 28 |
+
len(cells) == 1
|
| 29 |
+
and int(cells[0].get("colspan", 1)) > 1
|
| 30 |
+
and int(cells[0].get("rowspan", 1)) == 1
|
| 31 |
+
):
|
| 32 |
+
start += 1
|
| 33 |
+
else:
|
| 34 |
+
break
|
| 35 |
+
cells = rows[start].find_all(["td", "th"], recursive=False)
|
| 36 |
+
if not cells:
|
| 37 |
+
continue
|
| 38 |
+
numeric = lambda s: bool(re.fullmatch(r"[\s\d.,%$€£()+−-]+", s))
|
| 39 |
+
texts = [
|
| 40 |
+
c.get_text(" ", strip=True) for c in cells if c.get_text(" ", strip=True)
|
| 41 |
+
]
|
| 42 |
+
if not texts or sum(numeric(t) for t in texts) / len(texts) > 0.5:
|
| 43 |
+
diagnostics.append({"changed": False, "reason": "numeric_first_row"})
|
| 44 |
+
continue
|
| 45 |
+
end = start + 1
|
| 46 |
+
rs = max(int(c.get("rowspan", 1)) for c in cells)
|
| 47 |
+
cs = max(int(c.get("colspan", 1)) for c in cells)
|
| 48 |
+
if 1 < rs <= 4:
|
| 49 |
+
end = min(len(rows), start + rs)
|
| 50 |
+
elif rs > 4:
|
| 51 |
+
diagnostics.append({"changed": False, "reason": "long_rowspan"})
|
| 52 |
+
continue
|
| 53 |
+
elif cs > 1 and start + 1 < len(rows):
|
| 54 |
+
nxt = [
|
| 55 |
+
c
|
| 56 |
+
for c in rows[start + 1].find_all(["td", "th"], recursive=False)
|
| 57 |
+
if c.get_text(" ", strip=True)
|
| 58 |
+
]
|
| 59 |
+
# A group heading's colspan does not establish the next row's role.
|
| 60 |
+
# Require explicit header styling on every nonempty next-row cell;
|
| 61 |
+
# text-heavy data rows and year-valued data must not become headers.
|
| 62 |
+
if (
|
| 63 |
+
len(nxt) >= 2
|
| 64 |
+
and all(c.find(["b", "strong"]) for c in nxt)
|
| 65 |
+
and sum(numeric(c.get_text(" ", strip=True)) for c in nxt) / len(nxt)
|
| 66 |
+
< 0.5
|
| 67 |
+
):
|
| 68 |
+
end = start + 2
|
| 69 |
+
# Only probe multi-level headers. Single-row tables remain byte-for-byte as supplied.
|
| 70 |
+
if end - start <= 1:
|
| 71 |
+
diagnostics.append({"changed": False, "reason": "single_level"})
|
| 72 |
+
continue
|
| 73 |
+
for row in rows[start:end]:
|
| 74 |
+
for c in row.find_all(["td", "th"], recursive=False):
|
| 75 |
+
c.name = "th"
|
| 76 |
+
diagnostics.append(
|
| 77 |
+
{
|
| 78 |
+
"changed": True,
|
| 79 |
+
"start": start,
|
| 80 |
+
"end": end,
|
| 81 |
+
"rule": "bounded_span_header_topology",
|
| 82 |
+
}
|
| 83 |
+
)
|
| 84 |
+
return str(soup), diagnostics
|
nemotron_parse/engine/markdown.py
ADDED
|
@@ -0,0 +1,61 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
| 2 |
+
# SPDX-License-Identifier: OpenMDW-1.1
|
| 3 |
+
|
| 4 |
+
from __future__ import annotations
|
| 5 |
+
|
| 6 |
+
|
| 7 |
+
import re
|
| 8 |
+
|
| 9 |
+
|
| 10 |
+
_RE_EXTRACT_CLASS_BBOX = re.compile(
|
| 11 |
+
r"<x_(\d+(?:\.\d+)?)><y_(\d+(?:\.\d+)?)>(.*?)<x_(\d+(?:\.\d+)?)><y_(\d+(?:\.\d+)?)><class_([^>]+)>",
|
| 12 |
+
re.DOTALL,
|
| 13 |
+
)
|
| 14 |
+
|
| 15 |
+
|
| 16 |
+
def extract_classes_bboxes(
|
| 17 |
+
text: str,
|
| 18 |
+
) -> tuple[list[str], list[tuple[float, float, float, float]], list[str]]:
|
| 19 |
+
classes: list[str] = []
|
| 20 |
+
bboxes: list[tuple[float, float, float, float]] = []
|
| 21 |
+
texts: list[str] = []
|
| 22 |
+
for match in _RE_EXTRACT_CLASS_BBOX.finditer(text or ""):
|
| 23 |
+
x1, y1, body, x2, y2, cls = match.groups()
|
| 24 |
+
classes.append(cls)
|
| 25 |
+
bboxes.append((float(x1), float(y1), float(x2), float(y2)))
|
| 26 |
+
texts.append(body)
|
| 27 |
+
return classes, bboxes, texts
|
| 28 |
+
|
| 29 |
+
|
| 30 |
+
def convert_latex_tables_in_markdown(markdown_text: str) -> str:
|
| 31 |
+
if "\\begin{tabular}" not in markdown_text:
|
| 32 |
+
return markdown_text
|
| 33 |
+
try:
|
| 34 |
+
from .convert.latex2html import latex_table_to_html
|
| 35 |
+
|
| 36 |
+
return latex_table_to_html(markdown_text)
|
| 37 |
+
except (ValueError, IndexError):
|
| 38 |
+
return markdown_text
|
| 39 |
+
|
| 40 |
+
|
| 41 |
+
def mmd_to_markdown(prediction: str, text_transform=None) -> str:
|
| 42 |
+
classes, _boxes, texts = extract_classes_bboxes(prediction or "")
|
| 43 |
+
if not texts:
|
| 44 |
+
cleaned = re.sub(r"<x_[^>]+>|<y_[^>]+>|<class_[^>]+>", "", prediction or "")
|
| 45 |
+
return convert_latex_tables_in_markdown(cleaned.replace("<tbc>", "").strip())
|
| 46 |
+
|
| 47 |
+
blocks = []
|
| 48 |
+
for cls, text in zip(classes, texts):
|
| 49 |
+
if text_transform is not None:
|
| 50 |
+
text = text_transform(cls, text)
|
| 51 |
+
text = text.replace("<tbc>", "").strip()
|
| 52 |
+
if not text:
|
| 53 |
+
continue
|
| 54 |
+
if cls in {"Table", "Chart"}:
|
| 55 |
+
text = convert_latex_tables_in_markdown(text)
|
| 56 |
+
elif cls == "Title" and not text.startswith("#"):
|
| 57 |
+
text = f"# {text}"
|
| 58 |
+
elif cls == "Section-header" and not text.startswith("#"):
|
| 59 |
+
text = f"## {text}"
|
| 60 |
+
blocks.append(text)
|
| 61 |
+
return "\n\n".join(blocks).strip()
|
nemotron_parse/engine/merge.py
ADDED
|
@@ -0,0 +1,630 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
| 2 |
+
# SPDX-License-Identifier: OpenMDW-1.1
|
| 3 |
+
|
| 4 |
+
from __future__ import annotations
|
| 5 |
+
|
| 6 |
+
import argparse
|
| 7 |
+
|
| 8 |
+
|
| 9 |
+
import json
|
| 10 |
+
|
| 11 |
+
|
| 12 |
+
import re
|
| 13 |
+
|
| 14 |
+
|
| 15 |
+
from dataclasses import dataclass
|
| 16 |
+
|
| 17 |
+
|
| 18 |
+
from pathlib import Path
|
| 19 |
+
|
| 20 |
+
|
| 21 |
+
from typing import Any
|
| 22 |
+
|
| 23 |
+
|
| 24 |
+
BLOCK_RE = re.compile(
|
| 25 |
+
r"<x_(\d+(?:\.\d+)?)><y_(\d+(?:\.\d+)?)>(.*?)"
|
| 26 |
+
r"<x_(\d+(?:\.\d+)?)><y_(\d+(?:\.\d+)?)><class_([^>]+)>",
|
| 27 |
+
re.DOTALL,
|
| 28 |
+
)
|
| 29 |
+
|
| 30 |
+
|
| 31 |
+
TOKEN_RE = re.compile(r"</?s>|<pad>|<unk>|<image>|<\|[^>]+\|>")
|
| 32 |
+
|
| 33 |
+
|
| 34 |
+
@dataclass(frozen=True)
|
| 35 |
+
class Block:
|
| 36 |
+
x1: float
|
| 37 |
+
y1: float
|
| 38 |
+
body: str
|
| 39 |
+
x2: float
|
| 40 |
+
y2: float
|
| 41 |
+
cls: str
|
| 42 |
+
raw: str
|
| 43 |
+
detector_label: str = ""
|
| 44 |
+
origin: str = ""
|
| 45 |
+
detector_score: float = 0.0
|
| 46 |
+
|
| 47 |
+
@property
|
| 48 |
+
def bbox(self) -> tuple[float, float, float, float]:
|
| 49 |
+
left, right = sorted((self.x1, self.x2))
|
| 50 |
+
top, bottom = sorted((self.y1, self.y2))
|
| 51 |
+
return left, top, right, bottom
|
| 52 |
+
|
| 53 |
+
|
| 54 |
+
def load_jsonl(path: Path) -> list[dict[str, Any]]:
|
| 55 |
+
rows = []
|
| 56 |
+
with path.open(encoding="utf-8", errors="replace") as f:
|
| 57 |
+
for line_no, line in enumerate(f, 1):
|
| 58 |
+
if not line.strip():
|
| 59 |
+
continue
|
| 60 |
+
try:
|
| 61 |
+
rows.append(json.loads(line))
|
| 62 |
+
except json.JSONDecodeError as exc:
|
| 63 |
+
raise ValueError(f"{path}:{line_no}: {exc}") from exc
|
| 64 |
+
return rows
|
| 65 |
+
|
| 66 |
+
|
| 67 |
+
def row_aliases(row: dict[str, Any]) -> set[str]:
|
| 68 |
+
aliases = set()
|
| 69 |
+
for key in ("key", "image_name"):
|
| 70 |
+
value = row.get(key)
|
| 71 |
+
if value:
|
| 72 |
+
text = str(value)
|
| 73 |
+
aliases.add(text)
|
| 74 |
+
aliases.add(Path(text).stem)
|
| 75 |
+
aliases.add(Path(text).name)
|
| 76 |
+
path = row.get("image_path")
|
| 77 |
+
if path:
|
| 78 |
+
aliases.add(Path(str(path)).name)
|
| 79 |
+
aliases.add(Path(str(path)).stem)
|
| 80 |
+
aliases = {alias for alias in aliases if alias}
|
| 81 |
+
return aliases
|
| 82 |
+
|
| 83 |
+
|
| 84 |
+
def index_rows(rows: list[dict[str, Any]]) -> dict[str, dict[str, Any]]:
|
| 85 |
+
out = {}
|
| 86 |
+
for row in rows:
|
| 87 |
+
for alias in row_aliases(row):
|
| 88 |
+
out.setdefault(alias, row)
|
| 89 |
+
return out
|
| 90 |
+
|
| 91 |
+
|
| 92 |
+
def prediction_text(row: dict[str, Any] | None) -> str:
|
| 93 |
+
if not row:
|
| 94 |
+
return ""
|
| 95 |
+
for key in ("raw_out", "output_text", "prediction", "pred", "text"):
|
| 96 |
+
value = row.get(key)
|
| 97 |
+
if isinstance(value, str):
|
| 98 |
+
return value
|
| 99 |
+
return ""
|
| 100 |
+
|
| 101 |
+
|
| 102 |
+
def parse_blocks(raw: str) -> list[Block]:
|
| 103 |
+
blocks = []
|
| 104 |
+
for match in BLOCK_RE.finditer(raw or ""):
|
| 105 |
+
x1, y1, body, x2, y2, cls = match.groups()
|
| 106 |
+
blocks.append(
|
| 107 |
+
Block(float(x1), float(y1), body, float(x2), float(y2), cls, match.group(0))
|
| 108 |
+
)
|
| 109 |
+
return blocks
|
| 110 |
+
|
| 111 |
+
|
| 112 |
+
def block_with_raw(
|
| 113 |
+
bbox: tuple[float, float, float, float],
|
| 114 |
+
cls: str,
|
| 115 |
+
body: str,
|
| 116 |
+
*,
|
| 117 |
+
detector_label: str = "",
|
| 118 |
+
origin: str = "",
|
| 119 |
+
detector_score: float = 0.0,
|
| 120 |
+
) -> Block:
|
| 121 |
+
x1, y1, x2, y2 = bbox
|
| 122 |
+
raw = make_block(bbox, cls, body)
|
| 123 |
+
return Block(x1, y1, body, x2, y2, cls, raw, detector_label, origin, detector_score)
|
| 124 |
+
|
| 125 |
+
|
| 126 |
+
def rebuild_blocks(blocks: list[Block]) -> str:
|
| 127 |
+
return "\n\n".join(block.raw for block in blocks).strip()
|
| 128 |
+
|
| 129 |
+
|
| 130 |
+
def iou(
|
| 131 |
+
a: tuple[float, float, float, float], b: tuple[float, float, float, float]
|
| 132 |
+
) -> float:
|
| 133 |
+
ax1, ay1, ax2, ay2 = a
|
| 134 |
+
bx1, by1, bx2, by2 = b
|
| 135 |
+
ix1, iy1 = max(ax1, bx1), max(ay1, by1)
|
| 136 |
+
ix2, iy2 = min(ax2, bx2), min(ay2, by2)
|
| 137 |
+
iw, ih = max(0.0, ix2 - ix1), max(0.0, iy2 - iy1)
|
| 138 |
+
inter = iw * ih
|
| 139 |
+
if inter <= 0:
|
| 140 |
+
return 0.0
|
| 141 |
+
area_a = max(0.0, ax2 - ax1) * max(0.0, ay2 - ay1)
|
| 142 |
+
area_b = max(0.0, bx2 - bx1) * max(0.0, by2 - by1)
|
| 143 |
+
denom = area_a + area_b - inter
|
| 144 |
+
return inter / denom if denom > 0 else 0.0
|
| 145 |
+
|
| 146 |
+
|
| 147 |
+
def block_area(block: Block) -> float:
|
| 148 |
+
left, top, right, bottom = block.bbox
|
| 149 |
+
return max(0.0, right - left) * max(0.0, bottom - top)
|
| 150 |
+
|
| 151 |
+
|
| 152 |
+
def normalized_text(text: str) -> str:
|
| 153 |
+
return " ".join(re.sub(r"<[^>]+>", " ", text or "").lower().split())
|
| 154 |
+
|
| 155 |
+
|
| 156 |
+
def text_similarity(a: str, b: str) -> float:
|
| 157 |
+
a_tokens = set(normalized_text(a).split())
|
| 158 |
+
b_tokens = set(normalized_text(b).split())
|
| 159 |
+
if not a_tokens and not b_tokens:
|
| 160 |
+
return 1.0
|
| 161 |
+
if not a_tokens or not b_tokens:
|
| 162 |
+
return 0.0
|
| 163 |
+
return len(a_tokens & b_tokens) / len(a_tokens | b_tokens)
|
| 164 |
+
|
| 165 |
+
|
| 166 |
+
def clamp(value: float) -> float:
|
| 167 |
+
return min(1.0, max(0.0, float(value)))
|
| 168 |
+
|
| 169 |
+
|
| 170 |
+
def fmt_coord(value: float) -> str:
|
| 171 |
+
return f"{clamp(value):.6f}".rstrip("0").rstrip(".") or "0"
|
| 172 |
+
|
| 173 |
+
|
| 174 |
+
def make_block(bbox: tuple[float, float, float, float], cls: str, text: str) -> str:
|
| 175 |
+
x1, y1, x2, y2 = bbox
|
| 176 |
+
return f"<x_{fmt_coord(x1)}><y_{fmt_coord(y1)}>{text}<x_{fmt_coord(x2)}><y_{fmt_coord(y2)}><class_{cls}>"
|
| 177 |
+
|
| 178 |
+
|
| 179 |
+
def strip_model_tokens(text: str) -> str:
|
| 180 |
+
return TOKEN_RE.sub("", text or "").replace("<tbc>", "").strip()
|
| 181 |
+
|
| 182 |
+
|
| 183 |
+
def crop_markdown(raw_text: str) -> str:
|
| 184 |
+
blocks = []
|
| 185 |
+
for match in BLOCK_RE.finditer(raw_text or ""):
|
| 186 |
+
body = strip_model_tokens(match.group(3))
|
| 187 |
+
if body:
|
| 188 |
+
blocks.append(body)
|
| 189 |
+
if blocks:
|
| 190 |
+
return "\n\n".join(blocks).strip()
|
| 191 |
+
cleaned = re.sub(r"<x_[^>]+>|<y_[^>]+>|<class_[^>]+>", "", raw_text or "")
|
| 192 |
+
return strip_model_tokens(cleaned)
|
| 193 |
+
|
| 194 |
+
|
| 195 |
+
def normalize_detector_label(label: Any) -> str:
|
| 196 |
+
text = str(label).strip().lower()
|
| 197 |
+
text = (
|
| 198 |
+
text.replace("&", "and").replace("/", "_").replace("-", "_").replace(" ", "_")
|
| 199 |
+
)
|
| 200 |
+
text = re.sub(r"_+", "_", text).strip("_")
|
| 201 |
+
return text
|
| 202 |
+
|
| 203 |
+
|
| 204 |
+
def parse_label_filter(value: str) -> set[str]:
|
| 205 |
+
return {normalize_detector_label(item) for item in value.split(",") if item.strip()}
|
| 206 |
+
|
| 207 |
+
|
| 208 |
+
def class_for_detector(label: str, bbox_norm: tuple[float, float, float, float]) -> str:
|
| 209 |
+
label = normalize_detector_label(label)
|
| 210 |
+
if label == "table":
|
| 211 |
+
return "Table"
|
| 212 |
+
if label == "chart":
|
| 213 |
+
return "Chart"
|
| 214 |
+
if label == "infographic":
|
| 215 |
+
return "Picture"
|
| 216 |
+
if label == "title":
|
| 217 |
+
return "Title"
|
| 218 |
+
if label == "header_footer":
|
| 219 |
+
return (
|
| 220 |
+
"Page-header"
|
| 221 |
+
if ((bbox_norm[1] + bbox_norm[3]) / 2.0) < 0.5
|
| 222 |
+
else "Page-footer"
|
| 223 |
+
)
|
| 224 |
+
return "Text"
|
| 225 |
+
|
| 226 |
+
|
| 227 |
+
def model_canvas_norm_to_image_norm(
|
| 228 |
+
bbox: tuple[float, float, float, float],
|
| 229 |
+
width: int,
|
| 230 |
+
height: int,
|
| 231 |
+
*,
|
| 232 |
+
canvas_width: float,
|
| 233 |
+
canvas_height: float,
|
| 234 |
+
pad_position: str,
|
| 235 |
+
) -> tuple[float, float, float, float]:
|
| 236 |
+
x1, y1, x2, y2 = bbox
|
| 237 |
+
max_coord = max(abs(x1), abs(y1), abs(x2), abs(y2))
|
| 238 |
+
if max_coord > 1.5 and max_coord <= 1000.0:
|
| 239 |
+
norm = (x1 / 1000.0, y1 / 1000.0, x2 / 1000.0, y2 / 1000.0)
|
| 240 |
+
elif max_coord > 1000.0:
|
| 241 |
+
norm = (
|
| 242 |
+
x1 / max(1, width),
|
| 243 |
+
y1 / max(1, height),
|
| 244 |
+
x2 / max(1, width),
|
| 245 |
+
y2 / max(1, height),
|
| 246 |
+
)
|
| 247 |
+
else:
|
| 248 |
+
norm = (x1, y1, x2, y2)
|
| 249 |
+
scale = min(1.0, canvas_width / max(1, width), canvas_height / max(1, height))
|
| 250 |
+
resized_width = width * scale
|
| 251 |
+
resized_height = height * scale
|
| 252 |
+
if pad_position == "top-left":
|
| 253 |
+
pad_x = 0.0
|
| 254 |
+
pad_y = 0.0
|
| 255 |
+
elif pad_position == "bottom-right":
|
| 256 |
+
pad_x = canvas_width - resized_width
|
| 257 |
+
pad_y = canvas_height - resized_height
|
| 258 |
+
else:
|
| 259 |
+
pad_x = (canvas_width - resized_width) / 2.0
|
| 260 |
+
pad_y = (canvas_height - resized_height) / 2.0
|
| 261 |
+
out = (
|
| 262 |
+
(norm[0] * canvas_width - pad_x) / max(1.0, resized_width),
|
| 263 |
+
(norm[1] * canvas_height - pad_y) / max(1.0, resized_height),
|
| 264 |
+
(norm[2] * canvas_width - pad_x) / max(1.0, resized_width),
|
| 265 |
+
(norm[3] * canvas_height - pad_y) / max(1.0, resized_height),
|
| 266 |
+
)
|
| 267 |
+
left, right = sorted((clamp(out[0]), clamp(out[2])))
|
| 268 |
+
top, bottom = sorted((clamp(out[1]), clamp(out[3])))
|
| 269 |
+
return left, top, right, bottom
|
| 270 |
+
|
| 271 |
+
|
| 272 |
+
def image_norm_to_model_canvas_norm(
|
| 273 |
+
bbox: tuple[float, float, float, float],
|
| 274 |
+
width: int,
|
| 275 |
+
height: int,
|
| 276 |
+
*,
|
| 277 |
+
canvas_width: float,
|
| 278 |
+
canvas_height: float,
|
| 279 |
+
pad_position: str,
|
| 280 |
+
) -> tuple[float, float, float, float]:
|
| 281 |
+
x1, y1, x2, y2 = bbox
|
| 282 |
+
scale = min(1.0, canvas_width / max(1, width), canvas_height / max(1, height))
|
| 283 |
+
resized_width = width * scale
|
| 284 |
+
resized_height = height * scale
|
| 285 |
+
if pad_position == "top-left":
|
| 286 |
+
pad_x = 0.0
|
| 287 |
+
pad_y = 0.0
|
| 288 |
+
elif pad_position == "bottom-right":
|
| 289 |
+
pad_x = canvas_width - resized_width
|
| 290 |
+
pad_y = canvas_height - resized_height
|
| 291 |
+
else:
|
| 292 |
+
pad_x = (canvas_width - resized_width) / 2.0
|
| 293 |
+
pad_y = (canvas_height - resized_height) / 2.0
|
| 294 |
+
return (
|
| 295 |
+
clamp((x1 * resized_width + pad_x) / canvas_width),
|
| 296 |
+
clamp((y1 * resized_height + pad_y) / canvas_height),
|
| 297 |
+
clamp((x2 * resized_width + pad_x) / canvas_width),
|
| 298 |
+
clamp((y2 * resized_height + pad_y) / canvas_height),
|
| 299 |
+
)
|
| 300 |
+
|
| 301 |
+
|
| 302 |
+
def page_norm_to_output_frame(
|
| 303 |
+
bbox: tuple[float, float, float, float],
|
| 304 |
+
width: int,
|
| 305 |
+
height: int,
|
| 306 |
+
args: argparse.Namespace,
|
| 307 |
+
) -> tuple[float, float, float, float]:
|
| 308 |
+
"""Encode a page-normalized box in the direct prediction coordinate frame."""
|
| 309 |
+
frame = getattr(args, "output_coordinate_frame", None)
|
| 310 |
+
if frame in {"content_normalized", "original_normalized", "navit_content"}:
|
| 311 |
+
x1, y1, x2, y2 = bbox
|
| 312 |
+
left, right = sorted((clamp(x1), clamp(x2)))
|
| 313 |
+
top, bottom = sorted((clamp(y1), clamp(y2)))
|
| 314 |
+
return left, top, right, bottom
|
| 315 |
+
if frame not in {None, "padded_grid", "fixed_canvas", "navit_grid"}:
|
| 316 |
+
raise ValueError(f"Unsupported hybrid output coordinate frame: {frame!r}")
|
| 317 |
+
return image_norm_to_model_canvas_norm(
|
| 318 |
+
bbox,
|
| 319 |
+
width,
|
| 320 |
+
height,
|
| 321 |
+
canvas_width=args.canvas_width,
|
| 322 |
+
canvas_height=args.canvas_height,
|
| 323 |
+
pad_position=args.prediction_pad_position,
|
| 324 |
+
)
|
| 325 |
+
|
| 326 |
+
|
| 327 |
+
def crop_box_to_page_canvas(
|
| 328 |
+
crop_row: dict[str, Any],
|
| 329 |
+
crop_local_model_bbox: tuple[float, float, float, float],
|
| 330 |
+
args: argparse.Namespace,
|
| 331 |
+
) -> tuple[float, float, float, float]:
|
| 332 |
+
left, top, right, bottom = [int(v) for v in crop_row["crop_bbox_pixels"]]
|
| 333 |
+
crop_width = max(1, right - left)
|
| 334 |
+
crop_height = max(1, bottom - top)
|
| 335 |
+
page_width, page_height = [int(v) for v in crop_row["page_size"]]
|
| 336 |
+
crop_frame = getattr(args, "crop_coordinate_frame", "padded_grid")
|
| 337 |
+
if crop_frame in {"content_normalized", "original_normalized", "navit_content"}:
|
| 338 |
+
x1, y1, x2, y2 = crop_local_model_bbox
|
| 339 |
+
max_coord = max(abs(x1), abs(y1), abs(x2), abs(y2))
|
| 340 |
+
if max_coord > 1.5 and max_coord <= 1000.0:
|
| 341 |
+
x1, y1, x2, y2 = (value / 1000.0 for value in (x1, y1, x2, y2))
|
| 342 |
+
left_norm, right_norm = sorted((clamp(x1), clamp(x2)))
|
| 343 |
+
top_norm, bottom_norm = sorted((clamp(y1), clamp(y2)))
|
| 344 |
+
crop_norm = (left_norm, top_norm, right_norm, bottom_norm)
|
| 345 |
+
else:
|
| 346 |
+
crop_norm = model_canvas_norm_to_image_norm(
|
| 347 |
+
crop_local_model_bbox,
|
| 348 |
+
crop_width,
|
| 349 |
+
crop_height,
|
| 350 |
+
canvas_width=args.canvas_width,
|
| 351 |
+
canvas_height=args.canvas_height,
|
| 352 |
+
pad_position=args.crop_pad_position,
|
| 353 |
+
)
|
| 354 |
+
page_norm = (
|
| 355 |
+
(left + crop_norm[0] * crop_width) / page_width,
|
| 356 |
+
(top + crop_norm[1] * crop_height) / page_height,
|
| 357 |
+
(left + crop_norm[2] * crop_width) / page_width,
|
| 358 |
+
(top + crop_norm[3] * crop_height) / page_height,
|
| 359 |
+
)
|
| 360 |
+
return page_norm_to_output_frame(
|
| 361 |
+
page_norm,
|
| 362 |
+
page_width,
|
| 363 |
+
page_height,
|
| 364 |
+
args,
|
| 365 |
+
)
|
| 366 |
+
|
| 367 |
+
|
| 368 |
+
def page_bbox_to_page_canvas(
|
| 369 |
+
crop_row: dict[str, Any], args: argparse.Namespace
|
| 370 |
+
) -> tuple[float, float, float, float]:
|
| 371 |
+
page_width, page_height = [int(v) for v in crop_row["page_size"]]
|
| 372 |
+
bbox = tuple(float(v) for v in crop_row["detector_bbox_norm"])
|
| 373 |
+
return page_norm_to_output_frame(
|
| 374 |
+
bbox,
|
| 375 |
+
page_width,
|
| 376 |
+
page_height,
|
| 377 |
+
args,
|
| 378 |
+
)
|
| 379 |
+
|
| 380 |
+
|
| 381 |
+
def rebuild_crop_row_blocks(
|
| 382 |
+
crop_row: dict[str, Any], args: argparse.Namespace
|
| 383 |
+
) -> list[Block]:
|
| 384 |
+
detector_label = normalize_detector_label(crop_row.get("detector_label") or "")
|
| 385 |
+
detector_score = float(crop_row.get("detector_score") or 0.0)
|
| 386 |
+
raw = str(crop_row.get("raw_out") or "")
|
| 387 |
+
blocks = []
|
| 388 |
+
for match in BLOCK_RE.finditer(raw):
|
| 389 |
+
x1, y1, body, x2, y2, cls = match.groups()
|
| 390 |
+
text = strip_model_tokens(body)
|
| 391 |
+
if not text:
|
| 392 |
+
continue
|
| 393 |
+
bbox = crop_box_to_page_canvas(
|
| 394 |
+
crop_row, (float(x1), float(y1), float(x2), float(y2)), args
|
| 395 |
+
)
|
| 396 |
+
blocks.append(
|
| 397 |
+
block_with_raw(
|
| 398 |
+
bbox,
|
| 399 |
+
cls,
|
| 400 |
+
text,
|
| 401 |
+
detector_label=detector_label,
|
| 402 |
+
origin="crop_element",
|
| 403 |
+
detector_score=detector_score,
|
| 404 |
+
)
|
| 405 |
+
)
|
| 406 |
+
if blocks:
|
| 407 |
+
return blocks
|
| 408 |
+
text = crop_markdown(raw)
|
| 409 |
+
if not text:
|
| 410 |
+
return []
|
| 411 |
+
bbox = page_bbox_to_page_canvas(crop_row, args)
|
| 412 |
+
cls = class_for_detector(
|
| 413 |
+
detector_label, tuple(float(v) for v in crop_row["detector_bbox_norm"])
|
| 414 |
+
)
|
| 415 |
+
return [
|
| 416 |
+
block_with_raw(
|
| 417 |
+
bbox,
|
| 418 |
+
cls,
|
| 419 |
+
text,
|
| 420 |
+
detector_label=detector_label,
|
| 421 |
+
origin="crop_element_fallback",
|
| 422 |
+
detector_score=detector_score,
|
| 423 |
+
)
|
| 424 |
+
]
|
| 425 |
+
|
| 426 |
+
|
| 427 |
+
def parse_class_set(value: str) -> set[str]:
|
| 428 |
+
return {item.strip() for item in value.split(",") if item.strip()}
|
| 429 |
+
|
| 430 |
+
|
| 431 |
+
def crop_element_blocks_by_page(args: argparse.Namespace) -> dict[str, list[Block]]:
|
| 432 |
+
if not args.crop_element_predictions:
|
| 433 |
+
return {}
|
| 434 |
+
structural_classes = parse_class_set(args.structural_classes)
|
| 435 |
+
include_detector_labels = parse_label_filter(args.include_detector_labels)
|
| 436 |
+
by_page: dict[str, list[Block]] = {}
|
| 437 |
+
for path in args.crop_element_predictions:
|
| 438 |
+
for row in load_jsonl(path):
|
| 439 |
+
if row.get("error"):
|
| 440 |
+
continue
|
| 441 |
+
detector_label = normalize_detector_label(row.get("detector_label") or "")
|
| 442 |
+
if (
|
| 443 |
+
include_detector_labels
|
| 444 |
+
and detector_label not in include_detector_labels
|
| 445 |
+
):
|
| 446 |
+
continue
|
| 447 |
+
raw_text = prediction_text(row)
|
| 448 |
+
if (
|
| 449 |
+
args.max_crop_output_chars > 0
|
| 450 |
+
and len(raw_text) > args.max_crop_output_chars
|
| 451 |
+
):
|
| 452 |
+
continue
|
| 453 |
+
if (
|
| 454 |
+
args.max_crop_output_blocks > 0
|
| 455 |
+
and len(BLOCK_RE.findall(raw_text)) > args.max_crop_output_blocks
|
| 456 |
+
):
|
| 457 |
+
continue
|
| 458 |
+
page_key = str(row.get("page_key") or "")
|
| 459 |
+
if not page_key:
|
| 460 |
+
continue
|
| 461 |
+
for block in rebuild_crop_row_blocks(row, args):
|
| 462 |
+
if block.cls not in structural_classes:
|
| 463 |
+
continue
|
| 464 |
+
if (
|
| 465 |
+
args.min_structural_area > 0
|
| 466 |
+
and block_area(block) < args.min_structural_area
|
| 467 |
+
):
|
| 468 |
+
continue
|
| 469 |
+
by_page.setdefault(page_key, []).append(block)
|
| 470 |
+
for blocks in by_page.values():
|
| 471 |
+
blocks.sort(key=reading_order_key)
|
| 472 |
+
return by_page
|
| 473 |
+
|
| 474 |
+
|
| 475 |
+
def find_row(
|
| 476 |
+
index: dict[str, dict[str, Any]], label: dict[str, Any]
|
| 477 |
+
) -> dict[str, Any] | None:
|
| 478 |
+
for alias in row_aliases(label):
|
| 479 |
+
if alias in index:
|
| 480 |
+
return index[alias]
|
| 481 |
+
key = str(label.get("key") or "")
|
| 482 |
+
image_name = str(label.get("image_name") or "")
|
| 483 |
+
return index.get(key) or index.get(Path(image_name).stem) or index.get(image_name)
|
| 484 |
+
|
| 485 |
+
|
| 486 |
+
def clone_prediction(
|
| 487 |
+
label: dict[str, Any], source_row: dict[str, Any] | None, *, source_name: str
|
| 488 |
+
) -> dict[str, Any]:
|
| 489 |
+
raw_out = prediction_text(source_row)
|
| 490 |
+
row = {
|
| 491 |
+
"key": label["key"],
|
| 492 |
+
"image_name": label.get("image_name") or f"{label['key']}.jpg",
|
| 493 |
+
"image_path": (source_row or {}).get("image_path") or label.get("image_path"),
|
| 494 |
+
"category": label.get("category"),
|
| 495 |
+
"raw_out": raw_out,
|
| 496 |
+
"hybrid_source": source_name,
|
| 497 |
+
"error": (
|
| 498 |
+
(source_row or {}).get("error") if source_row else "missing_prediction"
|
| 499 |
+
),
|
| 500 |
+
}
|
| 501 |
+
for key in (
|
| 502 |
+
"coordinate_frame",
|
| 503 |
+
"image_preprocessing",
|
| 504 |
+
"model_padding",
|
| 505 |
+
"model_grid_size",
|
| 506 |
+
"resized_content_size",
|
| 507 |
+
):
|
| 508 |
+
if (source_row or {}).get(key) is not None:
|
| 509 |
+
row[key] = source_row[key]
|
| 510 |
+
return row
|
| 511 |
+
|
| 512 |
+
|
| 513 |
+
def structural_blocks(
|
| 514 |
+
args: argparse.Namespace, row: dict[str, Any] | None
|
| 515 |
+
) -> list[Block]:
|
| 516 |
+
structural_classes = parse_class_set(args.structural_classes)
|
| 517 |
+
blocks = [
|
| 518 |
+
block
|
| 519 |
+
for block in parse_blocks(prediction_text(row))
|
| 520 |
+
if block.cls in structural_classes
|
| 521 |
+
]
|
| 522 |
+
if args.min_structural_area > 0:
|
| 523 |
+
blocks = [
|
| 524 |
+
block for block in blocks if block_area(block) >= args.min_structural_area
|
| 525 |
+
]
|
| 526 |
+
return blocks
|
| 527 |
+
|
| 528 |
+
|
| 529 |
+
def candidate_structural_blocks(
|
| 530 |
+
args: argparse.Namespace,
|
| 531 |
+
label: dict[str, Any],
|
| 532 |
+
structural_row: dict[str, Any] | None,
|
| 533 |
+
element_blocks: dict[str, list[Block]],
|
| 534 |
+
) -> list[Block]:
|
| 535 |
+
if element_blocks:
|
| 536 |
+
return list(element_blocks.get(str(label.get("key") or ""), []))
|
| 537 |
+
return structural_blocks(args, structural_row)
|
| 538 |
+
|
| 539 |
+
|
| 540 |
+
def detector_label_counts(blocks: list[Block]) -> dict[str, int]:
|
| 541 |
+
counts: dict[str, int] = {}
|
| 542 |
+
for block in blocks:
|
| 543 |
+
label = block.detector_label or "unknown"
|
| 544 |
+
counts[label] = counts.get(label, 0) + 1
|
| 545 |
+
return counts
|
| 546 |
+
|
| 547 |
+
|
| 548 |
+
def is_duplicate(
|
| 549 |
+
candidate: Block, existing: list[Block], args: argparse.Namespace
|
| 550 |
+
) -> bool:
|
| 551 |
+
for block in existing:
|
| 552 |
+
if block.cls != candidate.cls:
|
| 553 |
+
continue
|
| 554 |
+
if iou(block.bbox, candidate.bbox) < args.dedupe_overlap_threshold:
|
| 555 |
+
continue
|
| 556 |
+
if text_similarity(block.body, candidate.body) >= args.dedupe_text_threshold:
|
| 557 |
+
return True
|
| 558 |
+
return False
|
| 559 |
+
|
| 560 |
+
|
| 561 |
+
def reading_order_key(block: Block) -> tuple[float, float, float, float]:
|
| 562 |
+
return block.y1, block.x1, block.y2, block.x2
|
| 563 |
+
|
| 564 |
+
|
| 565 |
+
def insert_additions_preserving_direct_order(
|
| 566 |
+
direct_ordered: list[Block], additions: list[Block]
|
| 567 |
+
) -> list[Block]:
|
| 568 |
+
"""Keep the direct model's block sequence stable, inserting new blocks by bbox."""
|
| 569 |
+
if not additions:
|
| 570 |
+
return direct_ordered
|
| 571 |
+
additions_sorted = sorted(additions, key=reading_order_key)
|
| 572 |
+
merged = []
|
| 573 |
+
add_idx = 0
|
| 574 |
+
for block in direct_ordered:
|
| 575 |
+
block_key = reading_order_key(block)
|
| 576 |
+
while (
|
| 577 |
+
add_idx < len(additions_sorted)
|
| 578 |
+
and reading_order_key(additions_sorted[add_idx]) < block_key
|
| 579 |
+
):
|
| 580 |
+
merged.append(additions_sorted[add_idx])
|
| 581 |
+
add_idx += 1
|
| 582 |
+
merged.append(block)
|
| 583 |
+
merged.extend(additions_sorted[add_idx:])
|
| 584 |
+
return merged
|
| 585 |
+
|
| 586 |
+
|
| 587 |
+
def build_direct_plus_ordered_structural_addonly(
|
| 588 |
+
args: argparse.Namespace,
|
| 589 |
+
labels: list[dict[str, Any]],
|
| 590 |
+
indexes: dict[str, dict[str, dict[str, Any]]],
|
| 591 |
+
element_blocks: dict[str, list[Block]],
|
| 592 |
+
) -> list[dict[str, Any]]:
|
| 593 |
+
structural_index = indexes.get(args.structural_source)
|
| 594 |
+
if structural_index is None and not element_blocks:
|
| 595 |
+
raise SystemExit(
|
| 596 |
+
f"Missing --{args.structural_source}-predictions for structural merge"
|
| 597 |
+
)
|
| 598 |
+
rows = []
|
| 599 |
+
for label in labels:
|
| 600 |
+
direct = find_row(indexes["direct"], label)
|
| 601 |
+
structural = (
|
| 602 |
+
find_row(structural_index, label) if structural_index is not None else None
|
| 603 |
+
)
|
| 604 |
+
direct_blocks = parse_blocks(prediction_text(direct))
|
| 605 |
+
added = []
|
| 606 |
+
skipped_duplicates = 0
|
| 607 |
+
for block in candidate_structural_blocks(
|
| 608 |
+
args, label, structural, element_blocks
|
| 609 |
+
):
|
| 610 |
+
if is_duplicate(block, direct_blocks + added, args):
|
| 611 |
+
skipped_duplicates += 1
|
| 612 |
+
continue
|
| 613 |
+
added.append(block)
|
| 614 |
+
merged_blocks = insert_additions_preserving_direct_order(direct_blocks, added)
|
| 615 |
+
row = clone_prediction(
|
| 616 |
+
label,
|
| 617 |
+
direct,
|
| 618 |
+
source_name=f"direct-plus-{args.structural_source}-ordered-structural-addonly",
|
| 619 |
+
)
|
| 620 |
+
row["raw_out"] = (
|
| 621 |
+
rebuild_blocks(merged_blocks)
|
| 622 |
+
or prediction_text(direct)
|
| 623 |
+
or prediction_text(structural)
|
| 624 |
+
)
|
| 625 |
+
row["hybrid_kept_direct_blocks"] = len(direct_blocks)
|
| 626 |
+
row["hybrid_added_structural_blocks"] = len(added)
|
| 627 |
+
row["hybrid_skipped_duplicate_blocks"] = skipped_duplicates
|
| 628 |
+
row["hybrid_added_detector_labels"] = detector_label_counts(added)
|
| 629 |
+
rows.append(row)
|
| 630 |
+
return rows
|
nemotron_parse/engine/policies.py
ADDED
|
@@ -0,0 +1,338 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
| 2 |
+
# SPDX-License-Identifier: OpenMDW-1.1
|
| 3 |
+
|
| 4 |
+
"""Fixed single-checkpoint output conversion; no ensemble routing or benchmark inputs."""
|
| 5 |
+
|
| 6 |
+
import re
|
| 7 |
+
|
| 8 |
+
from functools import lru_cache
|
| 9 |
+
|
| 10 |
+
BLOCK = re.compile(
|
| 11 |
+
r"(<x_[^>]+><y_[^>]+>)(.*?)(<x_[^>]+><y_[^>]+><class_([^>]+)>)", re.S
|
| 12 |
+
)
|
| 13 |
+
|
| 14 |
+
MDTABLE = re.compile(
|
| 15 |
+
r"(?m)^\|[^\n]*\|[ \t]*\n\|[ \t:|\-]+\|[ \t]*\n(?:\|[^\n]*\|[ \t]*(?:\n|$))+"
|
| 16 |
+
)
|
| 17 |
+
|
| 18 |
+
|
| 19 |
+
def promote_bold_header_row(markdown):
|
| 20 |
+
"""Promote a mostly bold first row to header cells, preserving its styling.
|
| 21 |
+
|
| 22 |
+
Existing header cells take precedence. This rule uses predicted markup and
|
| 23 |
+
table shape; it contains no language vocabulary and never adds bold text.
|
| 24 |
+
"""
|
| 25 |
+
from bs4 import BeautifulSoup
|
| 26 |
+
|
| 27 |
+
def transform(match):
|
| 28 |
+
original = match[0]
|
| 29 |
+
soup = BeautifulSoup(original, "html.parser")
|
| 30 |
+
table = soup.find("table")
|
| 31 |
+
if table.find("th"):
|
| 32 |
+
return original
|
| 33 |
+
rows = [
|
| 34 |
+
row for row in table.find_all("tr") if row.find_parent("table") == table
|
| 35 |
+
]
|
| 36 |
+
if len(rows) < 2:
|
| 37 |
+
return original
|
| 38 |
+
cells = rows[0].find_all(["td", "th"], recursive=False)
|
| 39 |
+
next_cells = rows[1].find_all(["td", "th"], recursive=False)
|
| 40 |
+
texts = [cell.get_text(" ", strip=True) for cell in cells]
|
| 41 |
+
numeric = lambda text: bool(re.fullmatch(r"[\s\d.,%$€£()+−-]+", text))
|
| 42 |
+
eligible = (
|
| 43 |
+
2 <= len(cells) <= 20
|
| 44 |
+
and len(next_cells) >= 2
|
| 45 |
+
and all(len(text) <= 100 for text in texts)
|
| 46 |
+
and sum(bool(text) and not numeric(text) for text in texts)
|
| 47 |
+
>= len(cells) * 0.6
|
| 48 |
+
)
|
| 49 |
+
mostly_bold = (
|
| 50 |
+
sum(bool(cell.find(["b", "strong"])) for cell in cells) >= len(cells) * 0.6
|
| 51 |
+
)
|
| 52 |
+
if not eligible or not mostly_bold:
|
| 53 |
+
return original
|
| 54 |
+
before = [
|
| 55 |
+
(cell.get_text(), cell.get("rowspan"), cell.get("colspan"))
|
| 56 |
+
for cell in table.find_all(["td", "th"])
|
| 57 |
+
]
|
| 58 |
+
for cell in cells:
|
| 59 |
+
cell.name = "th"
|
| 60 |
+
after = [
|
| 61 |
+
(cell.get_text(), cell.get("rowspan"), cell.get("colspan"))
|
| 62 |
+
for cell in table.find_all(["td", "th"])
|
| 63 |
+
]
|
| 64 |
+
assert before == after, "Header conversion changed cell content or spans"
|
| 65 |
+
return str(table)
|
| 66 |
+
|
| 67 |
+
return re.sub(r"<table\b[^>]*>.*?</table>", transform, markdown, flags=re.S | re.I)
|
| 68 |
+
|
| 69 |
+
|
| 70 |
+
def chart_render(raw, panels=False, compound=False):
|
| 71 |
+
"""Render explicit panels and existing numeric value/percent pairs."""
|
| 72 |
+
from html import escape
|
| 73 |
+
|
| 74 |
+
number = r"[+−-]?\d[\d,. ]*(?:[kKmMbB])?"
|
| 75 |
+
pair = re.compile(r"^\s*(" + number + r")\s*\(\s*(" + number + r"\s*%)\s*\)\s*$")
|
| 76 |
+
|
| 77 |
+
def chart(m):
|
| 78 |
+
body = m[2]
|
| 79 |
+
title = re.search(r"(?m)^title:[ \t]*([^\n]*)$", body)
|
| 80 |
+
title = title[1].strip() if title else ""
|
| 81 |
+
panel_names = {
|
| 82 |
+
p[1].strip(): p[2].strip()
|
| 83 |
+
for p in re.finditer(
|
| 84 |
+
r"(?m)^panel:\s*([^;\n]+)\s*;[^;\n]+;\s*([^\n]+)", body
|
| 85 |
+
)
|
| 86 |
+
}
|
| 87 |
+
matches = list(MDTABLE.finditer(body))
|
| 88 |
+
changes = {}
|
| 89 |
+
for table in matches:
|
| 90 |
+
if r"\|" in table[0] or "`" in table[0]:
|
| 91 |
+
continue
|
| 92 |
+
rows = [
|
| 93 |
+
[c.strip() for c in line.strip().strip("|").split("|")]
|
| 94 |
+
for line in table[0].strip().splitlines()
|
| 95 |
+
]
|
| 96 |
+
if (
|
| 97 |
+
len(rows) < 3
|
| 98 |
+
or any(len(row) != len(rows[0]) for row in rows)
|
| 99 |
+
or any(not re.fullmatch(r":?-+:?", c.replace(" ", "")) for c in rows[1])
|
| 100 |
+
):
|
| 101 |
+
continue
|
| 102 |
+
marker = re.search(r"table\[([^\]\n]+)\]:\s*$", body[: table.start()])
|
| 103 |
+
panel = panel_names.get(marker[1].strip(), "") if marker else ""
|
| 104 |
+
split = {}
|
| 105 |
+
if compound:
|
| 106 |
+
for c in range(1, len(rows[0])):
|
| 107 |
+
nonempty = [row[c] for row in rows[2:] if row[c]]
|
| 108 |
+
if len(nonempty) >= 2 and all(pair.fullmatch(v) for v in nonempty):
|
| 109 |
+
split[c] = True
|
| 110 |
+
if not split and not (panels and panel):
|
| 111 |
+
continue
|
| 112 |
+
caption = " — ".join(v for v in [title, panel if panels else ""] if v)
|
| 113 |
+
out = ["<table>"]
|
| 114 |
+
if caption:
|
| 115 |
+
out.append("<caption>" + escape(caption) + "</caption>")
|
| 116 |
+
out.append("<thead><tr>")
|
| 117 |
+
for c, label in enumerate(rows[0]):
|
| 118 |
+
span = (
|
| 119 |
+
' colspan="2"' if c in split else (' rowspan="2"' if split else "")
|
| 120 |
+
)
|
| 121 |
+
out.append('<th scope="col"' + span + ">" + escape(label) + "</th>")
|
| 122 |
+
out.append("</tr>")
|
| 123 |
+
if split:
|
| 124 |
+
out.append(
|
| 125 |
+
"<tr>"
|
| 126 |
+
+ "".join(
|
| 127 |
+
'<th scope="col">Value</th><th scope="col">Percent</th>'
|
| 128 |
+
for c in split
|
| 129 |
+
)
|
| 130 |
+
+ "</tr>"
|
| 131 |
+
)
|
| 132 |
+
out.append("</thead><tbody>")
|
| 133 |
+
for row in rows[2:]:
|
| 134 |
+
out.append("<tr>")
|
| 135 |
+
for c, text in enumerate(row):
|
| 136 |
+
values = (
|
| 137 |
+
(pair.fullmatch(text).groups() if text else ("", ""))
|
| 138 |
+
if c in split
|
| 139 |
+
else (text,)
|
| 140 |
+
)
|
| 141 |
+
out.extend("<td>" + escape(v.strip()) + "</td>" for v in values)
|
| 142 |
+
out.append("</tr>")
|
| 143 |
+
out.append("</tbody></table>\n")
|
| 144 |
+
changes[table.start()] = "".join(out)
|
| 145 |
+
if not changes:
|
| 146 |
+
return m[0]
|
| 147 |
+
new = MDTABLE.sub(lambda t: changes.get(t.start(), t[0]), body)
|
| 148 |
+
if len(changes) == len(matches) and title:
|
| 149 |
+
new = re.sub(r"(?m)^title:[ \t]*[^\n]*\n?", "", new, count=1)
|
| 150 |
+
return m[1] + new + m[3]
|
| 151 |
+
|
| 152 |
+
return re.sub(r"(<chart\b[^>]*>)(.*?)(</chart>)", chart, raw, flags=re.S)
|
| 153 |
+
|
| 154 |
+
|
| 155 |
+
def native_grid(body, repair=False):
|
| 156 |
+
from .convert.rich_table_reader import split_top, argument
|
| 157 |
+
|
| 158 |
+
pattern = re.compile(
|
| 159 |
+
r"\\begin\{tabular\}\{([clrx]+)\}\n(.*?)\n\\end\{tabular\}", re.S
|
| 160 |
+
)
|
| 161 |
+
matches = list(pattern.finditer(body))
|
| 162 |
+
if len(matches) != 1:
|
| 163 |
+
return body, None
|
| 164 |
+
m = matches[0]
|
| 165 |
+
new = m[0]
|
| 166 |
+
changed = False
|
| 167 |
+
|
| 168 |
+
def cell(part):
|
| 169 |
+
text = part.strip()
|
| 170 |
+
rs = cs = 1
|
| 171 |
+
while text.startswith((r"\multicolumn", r"\multirow")):
|
| 172 |
+
command = "multicolumn" if text.startswith(r"\multicolumn") else "multirow"
|
| 173 |
+
pos = len(command) + 1
|
| 174 |
+
count, pos = argument(text, pos)
|
| 175 |
+
_, pos = argument(text, pos)
|
| 176 |
+
inner, pos = argument(text, pos)
|
| 177 |
+
if text[pos:].strip():
|
| 178 |
+
raise ValueError("trailing_span_content")
|
| 179 |
+
if command == "multicolumn":
|
| 180 |
+
cs = int(count)
|
| 181 |
+
else:
|
| 182 |
+
rs = int(count)
|
| 183 |
+
text = inner
|
| 184 |
+
header = False
|
| 185 |
+
if text.startswith(r"\thead{"):
|
| 186 |
+
inner, end = argument(text, len(r"\thead"))
|
| 187 |
+
header = True
|
| 188 |
+
if text[end:].strip():
|
| 189 |
+
raise ValueError("trailing_header_content")
|
| 190 |
+
text = inner
|
| 191 |
+
return rs, cs, text, header
|
| 192 |
+
|
| 193 |
+
try:
|
| 194 |
+
rows = split_top(m[2], r"\\")
|
| 195 |
+
if rows[-1].strip():
|
| 196 |
+
return body, {"valid": False, "reason": "unterminated_row"}
|
| 197 |
+
rows.pop()
|
| 198 |
+
parsed = [[cell(c) for c in split_top(row, "&")] for row in rows]
|
| 199 |
+
widths = [sum(c[1] for c in row) for row in parsed]
|
| 200 |
+
columns = len(m[1])
|
| 201 |
+
if (
|
| 202 |
+
repair
|
| 203 |
+
and widths
|
| 204 |
+
and len(set(widths)) == 1
|
| 205 |
+
and 1 <= widths[0] <= 512
|
| 206 |
+
and widths[0] != columns
|
| 207 |
+
):
|
| 208 |
+
columns = widths[0]
|
| 209 |
+
cols = (m[1] + "c" * columns)[:columns]
|
| 210 |
+
new = r"\begin{tabular}{" + cols + "}\n" + m[2] + "\n" + r"\end{tabular}"
|
| 211 |
+
changed = True
|
| 212 |
+
occupied = set()
|
| 213 |
+
headers = 0
|
| 214 |
+
if not rows or len(rows) > 4096 or columns > 512:
|
| 215 |
+
raise ValueError("grid_size")
|
| 216 |
+
for ri, row in enumerate(parsed):
|
| 217 |
+
ci = 0
|
| 218 |
+
for rs, cs, content, header in row:
|
| 219 |
+
if rs < 1 or cs < 1 or ri + rs > len(rows) or ci + cs > columns:
|
| 220 |
+
raise ValueError("invalid_span")
|
| 221 |
+
if (ri, ci) in occupied:
|
| 222 |
+
empty = re.sub(r"\s+", "", content) in ("", r"\mbox{}")
|
| 223 |
+
if (
|
| 224 |
+
not empty
|
| 225 |
+
or header
|
| 226 |
+
or rs != 1
|
| 227 |
+
or any((ri, cc) not in occupied for cc in range(ci, ci + cs))
|
| 228 |
+
):
|
| 229 |
+
raise ValueError("invalid_placeholder")
|
| 230 |
+
else:
|
| 231 |
+
if any(
|
| 232 |
+
(rr, cc) in occupied
|
| 233 |
+
for rr in range(ri, ri + rs)
|
| 234 |
+
for cc in range(ci, ci + cs)
|
| 235 |
+
):
|
| 236 |
+
raise ValueError("overlap")
|
| 237 |
+
headers += int(header)
|
| 238 |
+
for rr in range(ri, ri + rs):
|
| 239 |
+
for cc in range(ci, ci + cs):
|
| 240 |
+
occupied.add((rr, cc))
|
| 241 |
+
ci += cs
|
| 242 |
+
if ci != columns:
|
| 243 |
+
raise ValueError("row_width")
|
| 244 |
+
header_shape = []
|
| 245 |
+
for row in parsed:
|
| 246 |
+
if not any(c[3] for c in row):
|
| 247 |
+
break
|
| 248 |
+
header_shape.append(tuple((c[0], c[1], c[3]) for c in row))
|
| 249 |
+
status = {
|
| 250 |
+
"valid": True,
|
| 251 |
+
"columns": columns,
|
| 252 |
+
"rows": len(rows),
|
| 253 |
+
"header_cells": headers,
|
| 254 |
+
"header_shape": tuple(header_shape),
|
| 255 |
+
"changed_column_declaration": changed,
|
| 256 |
+
}
|
| 257 |
+
except (ValueError, TypeError, IndexError, KeyError):
|
| 258 |
+
return body, {"valid": False, "reason": "inconsistent_native_grid"}
|
| 259 |
+
return body[: m.start()] + new + body[m.end() :], status
|
| 260 |
+
|
| 261 |
+
|
| 262 |
+
native_grid = lru_cache(maxsize=32768)(native_grid)
|
| 263 |
+
|
| 264 |
+
|
| 265 |
+
def replace_guarded(
|
| 266 |
+
raw,
|
| 267 |
+
candidates,
|
| 268 |
+
merger,
|
| 269 |
+
similarity=0.65,
|
| 270 |
+
guard="valid",
|
| 271 |
+
repair=False,
|
| 272 |
+
min_rows_ratio=0.0,
|
| 273 |
+
iou_threshold=0.5,
|
| 274 |
+
protect_headers=False,
|
| 275 |
+
emit_repair=True,
|
| 276 |
+
):
|
| 277 |
+
blocks = merger.parse_blocks(raw)
|
| 278 |
+
pairs = []
|
| 279 |
+
for i, b in enumerate(blocks):
|
| 280 |
+
if b.cls != "Table":
|
| 281 |
+
continue
|
| 282 |
+
_, oldstatus = native_grid(b.body, repair)
|
| 283 |
+
for j, c in enumerate(candidates):
|
| 284 |
+
if c.cls != "Table" or not c.body.strip():
|
| 285 |
+
continue
|
| 286 |
+
overlap = merger.iou(b.bbox, c.bbox)
|
| 287 |
+
if overlap < iou_threshold:
|
| 288 |
+
continue
|
| 289 |
+
body, status = native_grid(c.body, repair)
|
| 290 |
+
if status and not status["valid"]:
|
| 291 |
+
continue
|
| 292 |
+
if c.body.count(r"\begin{tabular}") != c.body.count(r"\end{tabular}"):
|
| 293 |
+
continue
|
| 294 |
+
if (
|
| 295 |
+
guard == "columns"
|
| 296 |
+
and status
|
| 297 |
+
and oldstatus
|
| 298 |
+
and oldstatus["valid"]
|
| 299 |
+
and status["columns"] != oldstatus["columns"]
|
| 300 |
+
):
|
| 301 |
+
continue
|
| 302 |
+
if (
|
| 303 |
+
min_rows_ratio
|
| 304 |
+
and status
|
| 305 |
+
and oldstatus
|
| 306 |
+
and oldstatus["valid"]
|
| 307 |
+
and status["rows"] < oldstatus["rows"] * min_rows_ratio
|
| 308 |
+
):
|
| 309 |
+
continue
|
| 310 |
+
if (
|
| 311 |
+
protect_headers
|
| 312 |
+
and status
|
| 313 |
+
and oldstatus
|
| 314 |
+
and oldstatus["valid"]
|
| 315 |
+
and oldstatus["header_shape"]
|
| 316 |
+
and status["header_shape"] != oldstatus["header_shape"]
|
| 317 |
+
):
|
| 318 |
+
continue
|
| 319 |
+
if not emit_repair:
|
| 320 |
+
body = c.body
|
| 321 |
+
sim = merger.text_similarity(b.body, body)
|
| 322 |
+
if sim < similarity:
|
| 323 |
+
continue
|
| 324 |
+
pairs.append((overlap, sim, i, j, body))
|
| 325 |
+
chosen = {}
|
| 326 |
+
used = set()
|
| 327 |
+
for overlap, sim, i, j, body in sorted(pairs, reverse=True):
|
| 328 |
+
if i in chosen or j in used:
|
| 329 |
+
continue
|
| 330 |
+
chosen[i] = body
|
| 331 |
+
used.add(j)
|
| 332 |
+
index = iter(range(len(blocks)))
|
| 333 |
+
|
| 334 |
+
def f(m):
|
| 335 |
+
i = next(index)
|
| 336 |
+
return m[1] + chosen[i] + m[3] if i in chosen else m[0]
|
| 337 |
+
|
| 338 |
+
return BLOCK.sub(f, raw)
|