Upload folder using huggingface_hub

README.md

@@ -0,0 +1,324 @@
+---
+license: other
+---
+# Model Overview
+
+### Description:
+NVIDIA Nemotron Parse v1.2 is designed to understand document semantics and extract text and tables elements with spatial grounding. Given an image, NVIDIA Nemotron Parse v1.2 produces structured annotations, including formatted text, bounding-boxes and the corresponding semantic classes, ordered according to the document's reading flow. It overcomes the shortcomings of traditional OCR technologies that struggle with complex document layouts with structural variability, and helps transform unstructured documents into actionable and machine-usable representations. This has several downstream benefits such as increasing the availability of training-data for Large Language Models (LLMs), improving the accuracy of extractor, curator, retriever and AI agentic applications, and enhancing document understanding pipelines.<br>
+
+This model is ready for commercial use. <br>
+
+## Quick Start
+
+### Install dependencies in your environment
+
+You can use a public image _nvcr.io/nvidia/pytorch:25.03-py3_ with the following library versions installed on top:
+
+```bash
+pip install accelerate==1.12.0
+pip install albumentations==2.0.8
+pip install transformers==4.51.3
+pip install timm==1.0.22
+```
+
+### Usage example
+
+```python
+import torch
+from PIL import Image, ImageDraw
+from transformers import AutoModel, AutoProcessor, AutoTokenizer, AutoConfig, AutoImageProcessor, GenerationConfig
+from postprocessing import extract_classes_bboxes, transform_bbox_to_original, postprocess_text
+
+# Load model and processor
+model_path = "nvidia/NVIDIA-Nemotron-Parse-v1.2"  # Or use a local path
+device = "cuda:0"
+
+model = AutoModel.from_pretrained(
+    model_path,
+    trust_remote_code=True,
+    torch_dtype=torch.bfloat16
+).to(device).eval()
+tokenizer = AutoTokenizer.from_pretrained(model_path)
+processor = AutoProcessor.from_pretrained(model_path, trust_remote_code=True)
+
+# Load image
+image = Image.open("path/to/your/image.jpg")
+task_prompt = "</s><s><predict_bbox><predict_classes><output_markdown><predict_no_text_in_pic>"
+# task_prompt = "</s><s><predict_bbox><predict_classes><output_markdown><predict_text_in_pic>"
+
+# Process image
+inputs = processor(images=[image], text=task_prompt, return_tensors="pt", add_special_tokens=False).to(device)
+
+generation_config = GenerationConfig.from_pretrained(model_path, trust_remote_code=True)
+# Generate text
+outputs = model.generate(**inputs,  generation_config=generation_config)
+
+# Decode the generated text
+generated_text = processor.batch_decode(outputs, skip_special_tokens=True)[0]
+```
+### Postprocessing
+
+```python
+from PIL import Image, ImageDraw
+from postprocessing import extract_classes_bboxes, transform_bbox_to_original, postprocess_text
+
+classes, bboxes, texts = extract_classes_bboxes(generated_text)
+bboxes = [transform_bbox_to_original(bbox, image.width, image.height) for bbox in bboxes]
+
+# Specify output formats for postprocessing
+table_format = 'latex' # latex | HTML | markdown
+text_format = 'markdown' # markdown | plain
+blank_text_in_figures = False # remove text inside 'Picture' class
+texts = [postprocess_text(text, cls = cls, table_format=table_format, text_format=text_format, blank_text_in_figures=blank_text_in_figures) for text, cls in zip(texts, classes)]
+
+for cl, bb, txt in zip(classes, bboxes, texts):
+    print(cl, ': ', txt)
+
+draw = ImageDraw.Draw(image)
+for bbox in bboxes:
+  draw.rectangle((bbox[0], bbox[1], bbox[2], bbox[3]), outline="red")
+```
+
+## Inference with VLLM
+Nemotron-Parse-v1.2 is [available in vllm main](https://github.com/vllm-project/vllm/pull/30864) and can be found in [vllm/vllm-openai:v0.14.1 docker image](https://hub.docker.com/layers/vllm/vllm-openai/v0.14.1/images/sha256-8e67731819426f7df194e5a0dfd6649d3aa3474f80c44f75b1e8711e76f8030a). 
+
+Note: when running on A100/A10 we recommend running vllm serve with _--attention-backend=TRITON_ATTN_
+
+You will need to install the following dependencies on top, and then follow the VLLM Inference example below:
+```bash
+pip install albumentations timm open_clip_torch
+```
+
+### VLLM Inference example
+
+#### Option 1: end-to-end python inference
+```python
+from vllm import LLM, SamplingParams
+from PIL import Image
+
+def main():
+    sampling_params = SamplingParams(
+        temperature=0,
+        top_k=1,
+        repetition_penalty=1.1,
+        max_tokens=9000,
+        skip_special_tokens=False,
+    )
+    
+    llm = LLM(
+        model="nvidia/NVIDIA-Nemotron-Parse-v1.2",
+        max_num_seqs=64,
+        limit_mm_per_prompt={"image": 1},
+        dtype="bfloat16",
+        trust_remote_code=True,
+    )
+    
+    image = Image.open("<YOUR-IMAGE-PATH>")
+    
+    prompts = [
+        {  # Implicit prompt
+            "prompt": "</s><s><predict_bbox><predict_classes><output_markdown><predict_no_text_in_pic>",
+            "multi_modal_data": {
+                "image": image
+            },
+        },
+        {  # Explicit encoder/decoder prompt
+            "encoder_prompt": {
+                "prompt": "",
+                "multi_modal_data": {
+                    "image": image
+                },
+            },
+            "decoder_prompt": "</s><s><predict_bbox><predict_classes><output_markdown><predict_no_text_in_pic>",
+        },
+    ]
+    
+    outputs = llm.generate(prompts, sampling_params)
+    
+    for output in outputs:
+        prompt = output.prompt
+        generated_text = output.outputs[0].text
+        print(f"Decoder prompt: {prompt!r}, Generated text: {generated_text!r}")
+
+if __name__ == "__main__":
+    main()
+```
+
+#### Option 2: vllm serve
+
+Alternatively, you can start a vllm server as:
+```bash
+vllm serve nvidia/NVIDIA-Nemotron-Parse-v1.2 \
+    --dtype bfloat16 \
+    --max-num-seqs 8 \
+    --limit-mm-per-prompt '{"image": 1}' \
+    --trust-remote-code \
+    --port 8000 \
+    --chat-template chat_template.jinja
+```
+with *chat_template.jinja* provided in this repository. Then, you can run inference as:
+
+```python
+import base64
+from openai import OpenAI
+
+client = OpenAI(
+    base_url="http://localhost:8000/v1",
+)
+
+# Read and base64-encode the image
+with open(<your-image-path>, "rb") as f:
+    img_b64 = base64.b64encode(f.read()).decode("utf-8")
+prompt_text = "</s><s><predict_bbox><predict_classes><output_markdown><predict_no_text_in_pic>"
+
+resp = client.chat.completions.create(
+    model="nvidia/NVIDIA-Nemotron-Parse-v1.2",
+    messages=[
+        {
+            "role": "user",
+            "content": [
+                {
+                    "type": "text",
+                    "text": prompt_text,
+                },
+                {
+                    "type": "image_url",
+                    "image_url": {
+                        "url": f"data:image/png;base64,{img_b64}",
+                    },
+                },
+            ],
+        }
+    ],
+    max_tokens=9000,
+    temperature=0.0,
+    extra_body={
+        "repetition_penalty": 1.1,
+        "top_k": 1,
+        "skip_special_tokens": False,
+    },
+)
+print(resp.choices[0].message.content)
+```
+
+*Note:* we recommend using the default prompt that extracts bounding boxes, classes, and text in markdown formatting for all use cases (`</s><s><predict_bbox><predict_classes><output_markdown><predict_no_text_in_pic>` or `</s><s><predict_bbox><predict_classes><output_markdown><predict_text_in_pic>`). If necessary, optionally the prompt that omits text extraction and only outputs bounding boxes and classes could be used: `</s><s><predict_bbox><predict_classes><output_no_text><predict_no_text_in_pic>`.
+
+#### Logits processors
+With Nemotron-Parse-v1.2 we share 2 logits processors available in logitsprocessors/ dir for vllm and in hf_logits_processor.py for the python model.
+NemotronParseRepetitionStopProcessor - when used during generation, detects repeating n-grams and forces the model to close the <x_><y_> block when detecting potential hallucination.
+NemotronParseTableInsertionLogitsProcessor - forces every block to follow a table structure (useful if, e.g., you are running the model on table image crops)
+
+Please refer to the example_with_processor.py for example usage with python model. With vllm, you can provide these as arguments to vllm serve, after exporting logitsprocs/ to PYTHONPATH, e.g.: 
+```
+vllm serve nvidia/NVIDIA-Nemotron-Parse-v1.2 \
+  --dtype bfloat16 \
+  --max-num-seqs 4 \
+  --limit-mm-per-prompt '{"image": 1}' \
+  --attention-backend=TRITON_ATTN \
+  --trust-remote-code \
+  --logits-processors nemotron_parse_vllm_logitprocs:NemotronParseTableInsertionLogitsProcessor \
+  --port 8000
+```
+
+An example of inference with vllm openai server is available in vllm_example.py
+
+### License/Terms of Use
+
+GOVERNING TERMS: [Product-Specific Terms for NVIDIA AI Products](https://www.nvidia.com/en-us/agreements/enterprise-software/product-specific-terms-for-ai-products/). Use of this model is governed by the [NVIDIA Community Model License](https://www.nvidia.com/en-us/agreements/enterprise-software/nvidia-community-models-license/). Use of the tokenizer included in this model is governed by the [CC-BY-4.0 license](https://creativecommons.org/licenses/by/4.0/).<br>
+
+
+### Deployment Geography:
+Global<br>
+
+### Use Case: <br>
+NVIDIA Nemotron Parse v1.2 will be capable of comprehensive text understanding and document structure understanding. It will be used in retriever and curator solutions. Its text extraction datasets and capabilities will help with LLM and VLM training, as well as improve run-time inference accuracy of VLMs. The NVIDIA Nemotron Parse v1.2 model will perform text extraction from PDF and PPT documents. The NVIDIA Nemotron Parse v1.2 can classify the objects (title, section, caption, index, footnote, lists, tables, bibliography, image) in a given document, and provide bounding boxes with coordinates.<br>
+
+### Release Date:  <br>
+Hugging Face [02/17/2026] via [[URL](https://huggingface.co/nvidia/NVIDIA-Nemotron-Parse-v1.2)] <br>
+
+## References(s):
+* https://huggingface.co/docs/transformers/en/model_doc/mbart  <br>
+
+## Model Architecture:
+**Architecture Type:** Transformer-based vision-encoder-decoder model<br>
+
+**Network Architecture:** <br>
+* Vision Encoder: ViT-H model (https://huggingface.co/nvidia/C-RADIO)<br>
+* Adapter Layer: 1D convolutions & norms to compress dimensionality and sequence length of the latent space (1280 tokens to 320 tokens)<br>
+* Decoder: mBart [1] 10 blocks<br>
+* Tokenizer: Use of the tokenizer included in this model is governed by the [CC-BY-4.0 license](https://creativecommons.org/licenses/by/4.0/)<br>
+* Number of Parameters: < 1B<br>
+
+## Input: <br>
+* Input Type: Image, Text<br>
+* Input Type(s): Red, Green, Blue (RGB) + Prompt (String)
+* Input Parameters: Two-Dimensional (2D), One-Dimensional (1D)
+* Other Properties Related to Input:
+  - Max Input Resolution (Width, Height): 1664, 2048
+  - Min Input Resolution (Width, Height): 1024, 1280
+* Channel Count: 3
+
+## Output: <br>
+* Output Type: Text<br>
+* Output Format: String
+* Output Parameters: One-Dimensional (1D)
+* Other Properties Related to Output: Nemotron-parse output format is a string which encodes text content (formatted or not) as well as bounding boxes and class attributes.<br>
+
+Our AI models are designed and/or optimized to run on NVIDIA GPU-accelerated systems. By leveraging NVIDIA’s hardware (e.g. GPU cores) and software frameworks (e.g., CUDA libraries), the model achieves faster training and inference times compared to CPU-only solutions. <br>
+
+## Software Integration:
+**Runtime Engine(s):**
+* TensorRT-LLM <br>
+* vLLM <br>
+
+
+**Supported Hardware Microarchitecture Compatibility:** <br>
+* NVIDIA Ampere <br>
+* NVIDIA Blackwell  <br>
+* NVIDIA Hopper <br>
+* NVIDIA Turing <br>
+
+
+**Supported Operating System(s):**
+* [Linux] <br>
+
+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 are essential to mitigate risks, meet technical and functional requirements, and ensure compliance with safety and ethical standards before deployment. <br>
+
+
+## Model Version(s):
+Nemotron Parse 1.2 <br>
+
+
+## Training, Testing, and Evaluation Datasets:
+
+### Training Dataset
+
+** Image Training Data Size <br>
+* [1 Million to 1 Billion Images] <br>
+
+** Text Training Data Size <br>
+* [1 Billion to 10 Trillion Tokens] <br>
+
+** Data Collection Method by dataset <br>
+* Hybrid: Automated, Human, Synthetic <br>
+
+** Labeling Method by dataset <br>
+* Hybrid: Automated, Human, Synthetic <br>
+
+**Properties (Quantity, Dataset Descriptions, Sensor(s)):** The training set contains millions of image–text items, aggregated across many large document and table datasets totaling several terabytes of data. The data consists of document-page and table images paired with OCR text, bounding boxes, and layout labels, drawn from real-world sources (scientific papers, PDFs, Wikipedia pages) as well as fully synthetic tables and word/character renderings. Modalities are primarily images plus associated text and structural annotations; content spans public-domain resources, and synthetic data. Images are obtained by rendering digital documents or generating synthetic layouts, and annotations come from OCR/layout models, third-party OCR services, and human labeling. <br>
+
+
+# Inference:
+**Acceleration Engine:** Tensor(RT)-LLM, vLLM <br>
+**Test Hardware:** <br>
+* H100 <br>
+* A100 <br>
+
+## Ethical Considerations:
+NVIDIA believes Trustworthy AI is a shared responsibility and we have established policies and practices to enable development for a wide array of AI applications.  When downloaded or used in accordance with our terms of service, developers should work with their internal model team to ensure this model meets requirements for the relevant industry and use case and addresses unforeseen product misuse. <br>
+
+Please make sure you have proper rights and permissions for all input image and video content; if image or video includes people, personal health information, or intellectual property, the image or video generated will not blur or maintain proportions of image subjects included. <br>
+
+For more detailed information on ethical considerations for this model, please see the [Model Card++ Bias, Explainability, Safety & Security, and Privacy Subcards](https://huggingface.co/nvidia/NVIDIA-Nemotron-Parse-v1.2/blob/main/bias.md,https://huggingface.co/nvidia/NVIDIA-Nemotron-Parse-v1.2/blob/main/explainability.md,https://huggingface.co/nvidia/NVIDIA-Nemotron-Parse-v1.2/blob/main/safety.md,https://huggingface.co/nvidia/NVIDIA-Nemotron-Parse-v1.2/blob/main/privacy.md)
+Please report model quality, risk, security vulnerabilities or NVIDIA AI Concerns [here](https://app.intigriti.com/programs/nvidia/nvidiavdp/detail).
+<br>

bias.md

@@ -0,0 +1,6 @@
+## Bias
+| Field | Response |
+| :---- | :---- |
+| Participation considerations from adversely impacted groups in model design and testing: | None |
+| Measures taken to mitigate against unwanted bias: | None|
+| Bias Metric (If Measured): | None |

chat_template.jinja

@@ -0,0 +1 @@
+{% for message in messages %}{{ message['content'] }}{% endfor %}

config.json

@@ -0,0 +1,402 @@
+{
+  "architectures": [
+    "NemotronParseForConditionalGeneration"
+  ],
+  "auto_map": {
+    "AutoConfig": "hf_nemotron_parse_config.NemotronParseConfig",
+    "AutoModel": "hf_nemotron_parse_modeling.NemotronParseForConditionalGeneration",
+    "AutoImageProcessor": "hf_nemotron_parse_processor.NemotronParseImageProcessor",
+    "AutoProcessor": "hf_nemotron_parse_processor.NemotronParseProcessor"
+  },
+  "bos_token_id": 0,
+  "decoder": {
+    "_attn_implementation": "sdpa",
+    "_name_or_path": "",
+    "activation_dropout": 0.0,
+    "activation_function": "gelu",
+    "add_cross_attention": true,
+    "add_final_layer_norm": true,
+    "architectures": null,
+    "attention_dropout": 0.0,
+    "bad_words_ids": null,
+    "begin_suppress_tokens": null,
+    "bos_token_id": 0,
+    "chunk_size_feed_forward": 0,
+    "classifier_dropout": 0.0,
+    "cross_attention_hidden_size": null,
+    "d_model": 1024,
+    "decoder_attention_heads": 16,
+    "decoder_ffn_dim": 4096,
+    "decoder_layerdrop": 0.0,
+    "decoder_layers": 10,
+    "decoder_start_token_id": null,
+    "diversity_penalty": 0.0,
+    "do_sample": false,
+    "dropout": 0.1,
+    "early_stopping": false,
+    "encoder_attention_heads": 16,
+    "encoder_ffn_dim": 4096,
+    "encoder_layerdrop": 0.0,
+    "encoder_layers": 12,
+    "encoder_no_repeat_ngram_size": 0,
+    "eos_token_id": 2,
+    "exponential_decay_length_penalty": null,
+    "finetuning_task": null,
+    "forced_bos_token_id": null,
+    "forced_eos_token_id": 2,
+    "hidden_size": 1024,
+    "id2label": {
+      "0": "LABEL_0",
+      "1": "LABEL_1",
+      "2": "LABEL_2"
+    },
+    "init_std": 0.02,
+    "is_decoder": true,
+    "is_encoder_decoder": false,
+    "label2id": {
+      "LABEL_0": 0,
+      "LABEL_1": 1,
+      "LABEL_2": 2
+    },
+    "length_penalty": 1.0,
+    "max_length": 20,
+    "min_length": 0,
+    "model_type": "nemotron_parse_text",
+    "no_repeat_ngram_size": 0,
+    "num_beam_groups": 1,
+    "num_beams": 1,
+    "num_hidden_layers": 12,
+    "num_return_sequences": 1,
+    "output_attentions": false,
+    "output_hidden_states": false,
+    "output_scores": false,
+    "pad_token_id": 1,
+    "prefix": null,
+    "problem_type": null,
+    "pruned_heads": {},
+    "remove_invalid_values": false,
+    "repetition_penalty": 1.0,
+    "return_dict": true,
+    "return_dict_in_generate": false,
+    "scale_embedding": true,
+    "sep_token_id": null,
+    "suppress_tokens": null,
+    "task_specific_params": null,
+    "temperature": 1.0,
+    "tf_legacy_loss": false,
+    "tie_encoder_decoder": false,
+    "tie_word_embeddings": false,
+    "tokenizer_class": null,
+    "top_k": 50,
+    "top_p": 1.0,
+    "torch_dtype": "bfloat16",
+    "torchscript": false,
+    "transformers_version": "4.51.3",
+    "typical_p": 1.0,
+    "use_bfloat16": true,
+    "use_cache": true,
+    "vocab_size": 52352
+  },
+  "decoder_start_token_id": 2,
+  "encoder": {
+    "attn_implementation": "eager",
+    "_name_or_path": "nvidia/C-RADIOv2-H",
+    "adaptor_configs": {},
+    "adaptor_names": null,
+    "add_cross_attention": false,
+    "architectures": [
+      "RADIOModel"
+    ],
+    "args": {
+      "aa": null,
+      "amp": true,
+      "amp_dtype": "bfloat16",
+      "amp_impl": "native",
+      "aug_repeats": 0,
+      "aug_splits": 0,
+      "bn_eps": null,
+      "bn_momentum": null,
+      "cache_dir": null,
+      "channels_last": false,
+      "checkpoint_hist": 10,
+      "chk_keep_forever": 100,
+      "class_map": "",
+      "clip_grad": null,
+      "clip_mode": "norm",
+      "cls_token_per_teacher": true,
+      "coco_annotations_file": "/datasets/coco2017-adlsa/annotations/captions_val2017.json",
+      "coco_image_dir": "/datasets/coco2017-adlsa/val2017",
+      "color_jitter": 0.4,
+      "cooldown_epochs": 0,
+      "cpe_max_size": 2048,
+      "crd_loss": false,
+      "crd_loss_weight": 0.8,
+      "crop_pct": null,
+      "cutmix": 0.0,
+      "cutmix_minmax": null,
+      "dataset_download": false,
+      "debug_full_knn": false,
+      "decay_epochs": 90,
+      "decay_milestones": [
+        90,
+        180,
+        270
+      ],
+      "decay_rate": 0.1,
+      "depchain": true,
+      "dist_bn": "reduce",
+      "dist_norm_weight": 0.0,
+      "distributed": true,
+      "drop": 0.0,
+      "drop_block": null,
+      "drop_connect": null,
+      "drop_path": null,
+      "dtype": "bfloat16",
+      "epoch_repeats": 0.0,
+      "eval": false,
+      "eval_metric": "knn_top1",
+      "eval_teacher": false,
+      "eval_teacher_only": false,
+      "eval_throughput": false,
+      "fast_norm": false,
+      "fd_loss_fn": "MSE",
+      "feature_normalization": "SHIP_NORM",
+      "feature_summarizer": "cls_token",
+      "feature_upscale_factor": null,
+      "force_new_wandb_id": false,
+      "force_spectral_reparam": true,
+      "freeze_bn": false,
+      "fsdp": false,
+      "fuser": "",
+      "gp": null,
+      "grad_accum_steps": 1,
+      "grad_checkpointing": false,
+      "head_init_bias": null,
+      "head_init_scale": null,
+      "head_warmup": 5,
+      "head_weight_decay": 0.001,
+      "hflip": 0.5,
+      "img_size": null,
+      "in_chans": null,
+      "initial_checkpoint": null,
+      "input_size": null,
+      "interpolation": "",
+      "layer_decay": null,
+      "local_rank": 0,
+      "log_interval": 50,
+      "log_mlflow": false,
+      "log_wandb": true,
+      "loss_auto_balance": false,
+      "lr_base": 0.1,
+      "lr_base_scale": "",
+      "lr_base_size": 256,
+      "lr_cycle_decay": 0.5,
+      "lr_cycle_limit": 1,
+      "lr_cycle_mul": 1.0,
+      "lr_k_decay": 1.0,
+      "lr_noise": null,
+      "lr_noise_pct": 0.67,
+      "lr_noise_std": 1.0,
+      "mean": null,
+      "mesa": false,
+      "min_lr": 0,
+      "mixup": 0.0,
+      "mixup_mode": "batch",
+      "mixup_off_epoch": 0,
+      "mixup_prob": 1.0,
+      "mixup_switch_prob": 0.5,
+      "mlp_hidden_size": 1520,
+      "mlp_num_inner": 3,
+      "mlp_version": "v2",
+      "model": "vit_huge_patch16_224",
+      "model_kwargs": {},
+      "model_norm": false,
+      "momentum": 0.9,
+      "no_aug": false,
+      "no_ddp_bb": true,
+      "no_prefetcher": false,
+      "no_resume_opt": false,
+      "num_classes": null,
+      "opt_betas": null,
+      "opt_eps": null,
+      "patience_epochs": 10,
+      "pin_mem": false,
+      "prefetcher": true,
+      "pretrained": false,
+      "rank": 0,
+      "ratio": [
+        0.75,
+        1.3333333333333333
+      ],
+      "recount": 1,
+      "recovery_interval": 0,
+      "register_multiple": 8,
+      "remode": "pixel",
+      "reprob": 0.0,
+      "reset_loss_state": false,
+      "resplit": false,
+      "save_images": false,
+      "scale": [
+        0.5,
+        1.0
+      ],
+      "sched": "cosine",
+      "seed": 42,
+      "smoothing": 0.1,
+      "spectral_heads": false,
+      "spectral_reparam": false,
+      "split_bn": false,
+      "start_epoch": null,
+      "std": null,
+      "stream_teachers": true,
+      "sync_bn": false,
+      "synchronize_step": false,
+      "teachers": [
+        {
+          "fd_normalize": false,
+          "feature_distillation": true,
+          "input_size": 378,
+          "model": "ViT-H-14-378-quickgelu",
+          "name": "clip",
+          "pretrained": "dfn5b",
+          "type": "open_clip",
+          "use_summary": true
+        },
+        {
+          "fd_normalize": false,
+          "feature_distillation": true,
+          "input_size": 378,
+          "model": "ViT-SO400M-14-SigLIP-384",
+          "name": "siglip",
+          "pretrained": "webli",
+          "type": "open_clip",
+          "use_summary": true
+        },
+        {
+          "fd_normalize": false,
+          "feature_distillation": true,
+          "input_size": 378,
+          "model": "dinov2_vitg14_reg",
+          "name": "dino_v2",
+          "type": "dino_v2",
+          "use_summary": true
+        },
+        {
+          "fd_normalize": false,
+          "feature_distillation": true,
+          "input_size": 1024,
+          "model": "vit-h",
+          "name": "sam",
+          "type": "sam",
+          "use_summary": false
+        }
+      ],
+      "torchcompile": null,
+      "torchscript": false,
+      "train_interpolation": "random",
+      "train_split": "train",
+      "tta": 0,
+      "use_coco": false,
+      "use_multi_epochs_loader": false,
+      "val_ema_only": false,
+      "val_split": "val",
+      "vflip": 0.0,
+      "vitdet_version": 1,
+      "wandb_entity": "",
+      "wandb_job_type": "",
+      "wandb_name": "",
+      "wandb_project": "",
+      "warmup_lr": 1e-05,
+      "warmup_prefix": false,
+      "worker_seeding": "all",
+      "workers": 8,
+      "world_size": 256
+    },
+    "auto_map": {
+      "AutoConfig": "nvidia/C-RADIOv2-H--hf_model.RADIOConfig",
+      "AutoModel": "nvidia/C-RADIOv2-H--hf_model.RADIOModel"
+    },
+    "bad_words_ids": null,
+    "begin_suppress_tokens": null,
+    "bos_token_id": null,
+    "chunk_size_feed_forward": 0,
+    "cross_attention_hidden_size": null,
+    "decoder_start_token_id": null,
+    "diversity_penalty": 0.0,
+    "do_sample": false,
+    "early_stopping": false,
+    "encoder_no_repeat_ngram_size": 0,
+    "eos_token_id": null,
+    "exponential_decay_length_penalty": null,
+    "feature_normalizer_config": null,
+    "finetuning_task": null,
+    "forced_bos_token_id": null,
+    "forced_eos_token_id": null,
+    "id2label": {
+      "0": "LABEL_0",
+      "1": "LABEL_1"
+    },
+    "inter_feature_normalizer_config": null,
+    "is_decoder": false,
+    "is_encoder_decoder": false,
+    "label2id": {
+      "LABEL_0": 0,
+      "LABEL_1": 1
+    },
+    "length_penalty": 1.0,
+    "max_length": 20,
+    "max_resolution": 2048,
+    "min_length": 0,
+    "model_type": "",
+    "no_repeat_ngram_size": 0,
+    "num_beam_groups": 1,
+    "num_beams": 1,
+    "num_return_sequences": 1,
+    "output_attentions": false,
+    "output_hidden_states": false,
+    "output_scores": false,
+    "pad_token_id": null,
+    "patch_size": 16,
+    "preferred_resolution": [
+      768,
+      768
+    ],
+    "prefix": null,
+    "problem_type": null,
+    "pruned_heads": {},
+    "remove_invalid_values": false,
+    "repetition_penalty": 1.0,
+    "return_dict": true,
+    "return_dict_in_generate": false,
+    "sep_token_id": null,
+    "suppress_tokens": null,
+    "task_specific_params": null,
+    "temperature": 1.0,
+    "tf_legacy_loss": false,
+    "tie_encoder_decoder": false,
+    "tie_word_embeddings": true,
+    "tokenizer_class": null,
+    "top_k": 50,
+    "top_p": 1.0,
+    "torch_dtype": "bfloat16",
+    "torchscript": false,
+    "transformers_version": "4.51.3",
+    "typical_p": 1.0,
+    "use_bfloat16": false,
+    "version": "radio_v2.5-h",
+    "vitdet_window_size": null
+  },
+  "eos_token_id": 2,
+  "image_size": [
+    2048,
+    1664
+  ],
+  "is_encoder_decoder": true,
+  "max_sequence_length": 9000,
+  "model_type": "nemotron_parse",
+  "pad_token_id": 1,
+  "tie_word_embeddings": false,
+  "torch_dtype": "bfloat16",
+  "transformers_version": "4.51.3",
+  "vocab_size": 52329
+}

example_with_processor.py

@@ -0,0 +1,96 @@
+"""
+Example usage of LogitsProcessors for document parsing.
+
+This example shows how to use:
+- TableInsertionLogitsProcessor: Force \begin{tabular} at the start of every object
+- RepetitionStopProcessor: Detect hallucination/repetition and force coordinate tokens
+"""
+
+import torch
+from PIL import Image, ImageDraw
+from transformers import AutoModel, AutoProcessor, AutoTokenizer, GenerationConfig
+from postprocessing import extract_classes_bboxes, transform_bbox_to_original, postprocess_text
+from hf_logits_processor import TableInsertionLogitsProcessor, RepetitionStopProcessor
+
+# Load model and processor
+model_path = "nvidia/NVIDIA-Nemotron-Parse-v1.2" or use a local path
+device = "cuda:0"
+
+model = AutoModel.from_pretrained(
+    model_path,
+    trust_remote_code=True,
+    torch_dtype=torch.bfloat16
+).to(device).eval()
+tokenizer = AutoTokenizer.from_pretrained(model_path)
+processor = AutoProcessor.from_pretrained(model_path, trust_remote_code=True)
+
+# Load image
+image = Image.open('example.png').convert("RGB")
+task_prompt = "</s><s><predict_bbox><predict_classes><output_markdown><predict_no_text_in_pic>"
+
+# Process image
+inputs = processor(images=[image], text=task_prompt, return_tensors="pt", add_special_tokens=False).to(device)
+
+generation_config = GenerationConfig.from_pretrained(model_path, trust_remote_code=True)
+
+# Create the table processor - inserts \begin{tabular} after every <x_...><y_...> that starts a new object
+table_processor = TableInsertionLogitsProcessor(
+    tokenizer=tokenizer,
+    table_prefix="\\begin{tabular}"
+)
+
+# Create the repetition stop processor - detects hallucination and forces <x_...> tokens
+repetition_processor = RepetitionStopProcessor(
+    tokenizer=tokenizer,
+    max_repetitions=10,  # Force stop after any pattern repeats 10+ times
+    ngram_sizes=[3, 4, 5, 6],  # Check these n-gram sizes for repetition
+    window_size=500  # Only check the last 500 tokens
+)
+
+# Generate with both logits processors
+outputs = model.generate(
+    **inputs,
+    generation_config=generation_config,
+    logits_processor=[table_processor, repetition_processor]
+)
+
+# Reset processor states for next generation (important for batch processing)
+table_processor.reset()
+repetition_processor.reset()
+
+# Decode and process the generated text
+generated_text = processor.batch_decode(outputs, skip_special_tokens=True)[0]
+print(outputs)
+print('--------------------------------')
+print("Generated text:", generated_text)
+print('--------------------------------')
+classes, bboxes, texts = extract_classes_bboxes(generated_text)
+bboxes = [transform_bbox_to_original(bbox, image.width, image.height) for bbox in bboxes]
+
+# Specify output formats for postprocessing
+table_format = 'HTML'  # latex | HTML | markdown
+text_format = 'markdown'  # markdown | plain
+blank_text_in_figures = False  # remove text inside 'Picture' class
+
+texts = [
+    postprocess_text(
+        text, 
+        cls=cls, 
+        table_format=table_format, 
+        text_format=text_format, 
+        blank_text_in_figures=blank_text_in_figures
+    ) 
+    for text, cls in zip(texts, classes)
+]
+
+for cl, bb, txt in zip(classes, bboxes, texts):
+    print(cl, ': ', txt)
+
+# OPTIONAL - Draw bounding boxes
+draw = ImageDraw.Draw(image)
+for bbox in bboxes:
+    draw.rectangle((bbox[0], bbox[1], (max(bbox[0], bbox[2])), (max(bbox[1], bbox[3]))), outline="red", width=2)
+
+# Save or display the image
+image.save("output_with_boxes.jpg")
+# image.show()

explainability.md

@@ -0,0 +1,15 @@
+## Explainability
+| Field | Response |
+| :---- | :---- |
+| Intended Task/Domain: | Document Understanding & Text Extraction|
+| Model Type: | Transformer-based vision-encoder-decoder model|
+| Intended Users: | Generative AI creators working with conversational AI models and image content. |
+| Output: | Text |
+| Describe how the model works: | Generates text by predicting the next word or token based on the context provided in the input sequence using multiple self-attention layers. |
+| Name the adversely impacted groups this has been tested to deliver comparable outcomes regardless of: | Not Applicable |
+| Technical Limitations & Mitigation: | The model demonstrates weakness to alignment-breaking attacks. Users are advised to deploy language model guardrails alongside this model to prevent potentially harmful outputs. The Model may generate answers that are inaccurate, omit key information, or include irrelevant or redundant text. |
+| Verified to have met prescribed NVIDIA quality standards: | Yes |
+| Performance Metrics: | Accuracy, Throughput, and User-side throughput |
+| Potential Known Risks: | The model was optimized explicitly for instruction following and as such is more susceptible to prompt injection and jailbreaking in various forms as a result of its instruction tuning. This means that the model should be paired with additional rails or system filtering to limit exposure to instructions from malicious sources  \-- either directly or indirectly by retrieval (e.g. via visiting a website)  \-- as they may yield outputs that can lead to harmful, system-level outcomes up to and including remote code execution in agentic systems when effective security controls including guardrails are not in place. The model was trained on data that contains toxic language and societal biases originally crawled from the internet. Therefore, the model may amplify those biases and return toxic responses especially when prompted with toxic prompts. The model may generate answers that may be inaccurate, omit key information, or include irrelevant or redundant text producing socially unacceptable or undesirable text, even if the prompt itself does not include anything explicitly offensive. |
+| Licensing: | GOVERNING TERMS: The NIM container is governed by the [NVIDIA Software License Agreement](https://www.nvidia.com/en-us/agreements/enterprise-software/nvidia-software-license-agreement/) and [Product-Specific Terms for NVIDIA AI Products](https://www.nvidia.com/en-us/agreements/enterprise-software/product-specific-terms-for-ai-products/). Use of this model is governed by the [NVIDIA Community Model License](https://www.nvidia.com/en-us/agreements/enterprise-software/nvidia-community-models-license/). Use of the tokenizer included in this model is governed by the [CC-BY-4.0 license](https://creativecommons.org/licenses/by/4.0/) |
+

generation_config.json

@@ -0,0 +1,13 @@
+{
+  "_from_model_config": true,
+  "bos_token_id": 0,
+  "decoder_start_token_id": 2,
+  "eos_token_id": 2,
+  "forced_eos_token_id": 2,
+  "pad_token_id": 1,
+  "transformers_version": "4.51.3",
+  "max_new_tokens": 9000,
+  "do_sample": false,
+  "num_beams": 1,
+  "repetition_penalty": 1.1
+}

hf_logits_processor.py

@@ -0,0 +1,349 @@
+"""
+LogitsProcessors for document parsing models.
+
+Contains:
+- TableInsertionLogitsProcessor: Forces table structure insertion after coordinates
+- RepetitionStopProcessor: Detects repetition/hallucination and forces coordinate tokens
+"""
+
+import torch
+from collections import Counter
+from transformers import LogitsProcessor
+from typing import Set, List, Tuple
+
+
+class TableInsertionLogitsProcessor(LogitsProcessor):
+    """
+    A LogitsProcessor that inserts \begin{tabular} tokens after coordinate pairs
+    that mark the START of an object (not the END coordinates).
+    
+    Pattern: <x_start><y_start>CONTENT<x_end><y_end><class_...>
+    
+    This processor triggers after the first <x_...><y_...> pair (start coords),
+    but not after the second pair (end coords).
+    
+    Args:
+        tokenizer: The tokenizer used for encoding/decoding
+        table_prefix: The string to insert (default: "\\begin{tabular}")
+    """
+    
+    def __init__(
+        self, 
+        tokenizer, 
+        table_prefix: str = "\\begin{tabular}"
+    ):
+        self.tokenizer = tokenizer
+        self.table_prefix = table_prefix
+        
+        # Tokenize the table prefix to get the sequence of tokens to force
+        self.table_prefix_ids = tokenizer.encode(
+            table_prefix, 
+            add_special_tokens=False
+        )
+        
+        # Build sets of token IDs for detection
+        self._build_token_sets()
+        
+        # State tracking for forced insertion
+        self._insertion_position = {}  # batch_idx -> position in table_prefix_ids
+        self._insertion_active = {}    # batch_idx -> bool
+        
+        # State tracking for coordinate pairs
+        # False = expecting START coordinates (should trigger)
+        # True = expecting END coordinates (should NOT trigger)
+        self._expecting_end_coords = {}  # batch_idx -> bool
+    
+    def _build_token_sets(self):
+        """Build sets of token IDs for x_, y_, class_, and special tokens."""
+        vocab = self.tokenizer.get_vocab()
+        
+        # Collect all x_ coordinate tokens
+        self.x_token_ids: Set[int] = set()
+        for token, token_id in vocab.items():
+            if token.startswith("<x_") and token.endswith(">"):
+                self.x_token_ids.add(token_id)
+        
+        # Collect all y_ coordinate tokens
+        self.y_token_ids: Set[int] = set()
+        for token, token_id in vocab.items():
+            if token.startswith("<y_") and token.endswith(">"):
+                self.y_token_ids.add(token_id)
+        
+        # Collect all class tokens
+        self.class_token_ids: Set[int] = set()
+        for token, token_id in vocab.items():
+            if token.startswith("<class_") and token.endswith(">"):
+                self.class_token_ids.add(token_id)
+    
+    def _is_xy_pair(self, input_ids: torch.Tensor, batch_idx: int) -> bool:
+        """Check if the sequence ends with <x_...><y_...>."""
+        seq = input_ids[batch_idx].tolist()
+        
+        if len(seq) < 2:
+            return False
+        
+        return seq[-1] in self.y_token_ids and seq[-2] in self.x_token_ids
+    
+    def _last_token_is_class(self, input_ids: torch.Tensor, batch_idx: int) -> bool:
+        """Check if the last token is a <class_...> token."""
+        seq = input_ids[batch_idx].tolist()
+        
+        if len(seq) < 1:
+            return False
+        
+        return seq[-1] in self.class_token_ids
+    
+    def reset(self):
+        """Reset the processor state for a new generation."""
+        self._insertion_position = {}
+        self._insertion_active = {}
+        self._expecting_end_coords = {}
+    
+    def __call__(
+        self, 
+        input_ids: torch.LongTensor, 
+        scores: torch.FloatTensor
+    ) -> torch.FloatTensor:
+        """
+        Process logits to force table prefix insertion when appropriate.
+        
+        Args:
+            input_ids: (batch_size, seq_len) - tokens generated so far
+            scores: (batch_size, vocab_size) - logits for next token
+            
+        Returns:
+            Modified scores with forced tokens where appropriate
+        """
+        batch_size = input_ids.shape[0]
+        
+        for batch_idx in range(batch_size):
+            # Check if we're currently in an insertion sequence
+            if self._insertion_active.get(batch_idx, False):
+                pos = self._insertion_position[batch_idx]
+                
+                if pos < len(self.table_prefix_ids):
+                    # Force the next token in the sequence
+                    forced_token_id = self.table_prefix_ids[pos]
+                    scores[batch_idx] = torch.full_like(scores[batch_idx], float('-inf'))
+                    scores[batch_idx, forced_token_id] = 0.0
+                    
+                    # Advance position
+                    self._insertion_position[batch_idx] = pos + 1
+                else:
+                    # Finished inserting, deactivate
+                    self._insertion_active[batch_idx] = False
+                continue
+            
+            # Check if we just saw a <class_...> token - reset state to expect START coords
+            if self._last_token_is_class(input_ids, batch_idx):
+                self._expecting_end_coords[batch_idx] = False
+            
+            # Check if we have an <x_...><y_...> pair
+            if self._is_xy_pair(input_ids, batch_idx):
+                expecting_end = self._expecting_end_coords.get(batch_idx, False)
+                
+                if not expecting_end:
+                    # This is START coordinates - trigger insertion!
+                    self._insertion_active[batch_idx] = True
+                    self._insertion_position[batch_idx] = 0
+                    
+                    # After START coords, we expect END coords next
+                    self._expecting_end_coords[batch_idx] = True
+                    
+                    # Force the first token
+                    if len(self.table_prefix_ids) > 0:
+                        forced_token_id = self.table_prefix_ids[0]
+                        scores[batch_idx] = torch.full_like(scores[batch_idx], float('-inf'))
+                        scores[batch_idx, forced_token_id] = 0.0
+                        self._insertion_position[batch_idx] = 1
+                else:
+                    # This is END coordinates - don't trigger, but reset for next object
+                    # The <class_...> token will come next, which resets expecting_end_coords
+                    pass
+        
+        return scores
+
+
+class RepetitionStopProcessor(LogitsProcessor):
+    """
+    A LogitsProcessor that detects CONSECUTIVE repetition/hallucination in generated 
+    content and forces the model to output <x_...> tokens to close the current object.
+    
+    This detects patterns like:
+    - <x_0.0><x_0.0><x_0.0>... (same token repeating consecutively)
+    - "ABC ABC ABC ABC..." (same phrase repeating consecutively)
+    
+    But NOT patterns like:
+    - "& 1 & 2 & 3 & 4..." (delimiters with different content between them)
+    
+    When consecutive repetition exceeds the threshold, the processor forces an <x_...> 
+    token to start coordinates, effectively ending the current content. After triggering, 
+    it enters a cooldown until a <class_...> token is seen (indicating the object was 
+    closed properly).
+    
+    Args:
+        tokenizer: The tokenizer used for encoding/decoding
+        max_repetitions: Max consecutive repetitions before forcing stop (default: 10)
+        ngram_sizes: List of n-gram sizes to check for repetition (default: [1, 2, 3, 4, 5])
+        window_size: Size of the sliding window to check for repetitions (default: 2 * max_ngram * (max_repetitions + 1))
+    """
+    
+    def __init__(
+        self,
+        tokenizer,
+        max_repetitions: int = 10,
+        ngram_sizes: List[int] = None,
+        window_size: int | None = None
+    ):
+        self.tokenizer = tokenizer
+        self.max_repetitions = max_repetitions
+        # Check various n-gram sizes for consecutive repetition
+        # n=1 catches single token repetition like <x_0.0><x_0.0><x_0.0>...
+        # n=2,3,4,5 catch phrase repetition like "ABC ABC ABC..."
+        self.ngram_sizes = ngram_sizes if ngram_sizes is not None else [1, 2, 3, 4, 5]
+        max_ngram = max((n for n in self.ngram_sizes if n > 0), default=1)
+        default_window = 2 * max_ngram * (self.max_repetitions + 1)
+        self.window_size = int(window_size) if window_size is not None else int(default_window)
+        
+        # Build set of token IDs
+        self._build_token_sets()
+        
+        # State tracking - once triggered, stay in cooldown until we see <class_...>
+        self._in_cooldown = {}  # batch_idx -> bool
+        # Track the start of the "current object segment" (tokens after the last <class_...>)
+        # We only look for repetition inside this segment to avoid old repetition re-triggering immediately.
+        self._segment_start = {}  # batch_idx -> int (seq index)
+    
+    def _build_token_sets(self):
+        """Build sets of token IDs for detection."""
+        vocab = self.tokenizer.get_vocab()
+        
+        # Collect all x_ coordinate tokens
+        self.x_token_ids: Set[int] = set()
+        self.x_token_list: List[int] = []
+        for token, token_id in vocab.items():
+            if token.startswith("<x_") and token.endswith(">"):
+                self.x_token_ids.add(token_id)
+                self.x_token_list.append(token_id)
+        
+        # Sort and pick a middle x token as default
+        self.x_token_list.sort()
+        self.default_x_token = self.x_token_list[len(self.x_token_list) // 2] if self.x_token_list else None
+        
+        # Collect all class tokens (used to reset cooldown)
+        self.class_token_ids: Set[int] = set()
+        for token, token_id in vocab.items():
+            if token.startswith("<class_") and token.endswith(">"):
+                self.class_token_ids.add(token_id)
+    
+    def _count_consecutive_repetitions(self, seq: List[int], n: int) -> int:
+        """
+        Count the maximum number of times an n-gram repeats CONSECUTIVELY.
+        
+        For example, if seq contains "ABC ABC ABC ABC DEF ABC":
+        - The n-gram "ABC" appears 4 times consecutively at the start
+        - Returns 4 (the max consecutive count)
+        
+        Returns:
+            Maximum consecutive repetition count for any n-gram
+        """
+        if len(seq) < n:
+            return 0
+        
+        max_consecutive = 1
+        current_consecutive = 1
+        
+        # Slide through the sequence comparing adjacent n-grams
+        prev_ngram = tuple(seq[0:n])
+        
+        i = n
+        while i <= len(seq) - n:
+            current_ngram = tuple(seq[i:i + n])
+            
+            if current_ngram == prev_ngram:
+                current_consecutive += 1
+                max_consecutive = max(max_consecutive, current_consecutive)
+                i += n  # Skip by n to check the next occurrence
+            else:
+                current_consecutive = 1
+                prev_ngram = current_ngram
+                i += 1  # Move by 1 to find new patterns
+        
+        return max_consecutive
+    
+    def _has_excessive_repetition(self, seq: List[int]) -> bool:
+        """
+        Check if the sequence has excessive CONSECUTIVE repetition of any n-gram.
+        
+        This detects hallucination patterns like:
+        - <x_0.0><x_0.0><x_0.0><x_0.0>... (same token repeating)
+        - "ABC ABC ABC ABC..." (same phrase repeating)
+        
+        But NOT patterns like:
+        - "& 1 & 2 & 3 & 4..." (delimiters with different content)
+        
+        Returns:
+            True if any n-gram repeats consecutively more than max_repetitions times
+        """
+        # Only check the recent window
+        check_seq = seq[-self.window_size:] if len(seq) > self.window_size else seq
+        
+        for n in self.ngram_sizes:
+            consecutive_count = self._count_consecutive_repetitions(check_seq, n)
+            if consecutive_count > self.max_repetitions:
+                return True
+        
+        return False
+    
+    def reset(self):
+        """Reset the processor state for a new generation."""
+        self._in_cooldown = {}
+        self._segment_start = {}
+    
+    def __call__(
+        self,
+        input_ids: torch.LongTensor,
+        scores: torch.FloatTensor
+    ) -> torch.FloatTensor:
+        """
+        Process logits to force <x_...> token when repetition is detected.
+        
+        Args:
+            input_ids: (batch_size, seq_len) - tokens generated so far
+            scores: (batch_size, vocab_size) - logits for next token
+            
+        Returns:
+            Modified scores with forced <x_...> tokens where repetition detected
+        """
+        batch_size = input_ids.shape[0]
+        
+        for batch_idx in range(batch_size):
+            seq = input_ids[batch_idx].tolist()
+            seq_len = len(seq)
+            
+            # Check if we just saw a class token - reset cooldown and start a new segment
+            if seq_len > 0 and seq[-1] in self.class_token_ids:
+                self._in_cooldown[batch_idx] = False
+                self._segment_start[batch_idx] = seq_len
+            
+            # If in cooldown, don't check for repetition (let model generate naturally)
+            if self._in_cooldown.get(batch_idx, False):
+                continue
+            
+            # Check if repetition threshold is exceeded
+            segment_start = self._segment_start.get(batch_idx, 0)
+            segment_seq = seq[segment_start:]
+            if self._has_excessive_repetition(segment_seq):
+                # Enter cooldown to avoid repeated triggering
+                self._in_cooldown[batch_idx] = True
+                
+                # Force one of the <x_...> tokens to start closing the object
+                # Keep the model's original preferences for which x_ coordinate to use
+                original_scores = scores[batch_idx].clone()
+                scores[batch_idx] = torch.full_like(scores[batch_idx], float('-inf'))
+                
+                # Restore original logits for x_ tokens only
+                for x_token_id in self.x_token_ids:
+                    scores[batch_idx, x_token_id] = original_scores[x_token_id]
+        
+        return scores

hf_nemotron_parse_config.py

@@ -0,0 +1,136 @@
+from os import truncate
+from quopri import decodestring
+from transformers import PretrainedConfig
+from typing import List, Optional
+
+from transformers.dynamic_module_utils import get_class_from_dynamic_module
+
+class NemotronParseTextConfig(PretrainedConfig):
+    """
+    Configuration class for NemotronParse text decoder (mBART-based).
+    """
+    model_type = "nemotron_parse_text"
+
+    def __init__(
+        self,
+        vocab_size: int = 250027,
+        d_model: int = 1024,
+        encoder_layers: int = 12,
+        decoder_layers: int = 12,
+        encoder_attention_heads: int = 16,
+        decoder_attention_heads: int = 16,
+        decoder_ffn_dim: int = 4096,
+        encoder_ffn_dim: int = 4096,
+        activation_function: str = "gelu",
+        dropout: float = 0.1,
+        attention_dropout: float = 0.0,
+        activation_dropout: float = 0.0,
+        classifier_dropout: float = 0.0,
+        init_std: float = 0.02,
+        encoder_layerdrop: float = 0.0,
+        decoder_layerdrop: float = 0.0,
+        scale_embedding: bool = False,
+        use_cache: bool = True,
+        num_labels: int = 3,
+        forced_eos_token_id: int = 2,
+        add_cross_attention: bool = True,  # Enable cross-attention for vision-encoder-decoder
+        is_decoder: bool = True,  # This is a decoder
+        max_sequence_length: int = 9000,
+        **kwargs
+    ):
+        super().__init__(**kwargs)
+        self.vocab_size = vocab_size
+        self.d_model = d_model
+        self.encoder_layers = encoder_layers
+        self.decoder_layers = decoder_layers
+        self.encoder_attention_heads = encoder_attention_heads
+        self.decoder_attention_heads = decoder_attention_heads
+        self.decoder_ffn_dim = decoder_ffn_dim
+        self.encoder_ffn_dim = encoder_ffn_dim
+        self.activation_function = activation_function
+        self.dropout = dropout
+        self.attention_dropout = attention_dropout
+        self.activation_dropout = activation_dropout
+        self.classifier_dropout = classifier_dropout
+        self.init_std = init_std
+        self.encoder_layerdrop = encoder_layerdrop
+        self.decoder_layerdrop = decoder_layerdrop
+        self.scale_embedding = scale_embedding
+        self.use_cache = use_cache
+        self.num_labels = num_labels
+        self.add_cross_attention = add_cross_attention
+        self.is_decoder = is_decoder
+        
+        # Add hidden_size as alias for d_model (for compatibility)
+        self.hidden_size = self.d_model
+        self.forced_eos_token_id = forced_eos_token_id
+        self.num_attention_heads = self.encoder_attention_heads
+
+        self.max_sequence_length = max_sequence_length
+
+
+class NemotronParseConfig(PretrainedConfig):
+    """
+    Configuration class for NemotronParse model.
+    
+    This configuration class is used to store the configuration of a [`NemotronParseForConditionalGeneration`] model.
+    It is used to instantiate an NemotronParse model according to the specified arguments, defining the vision and text model configs.
+    """
+    model_type = "nemotron_parse"
+    is_composition = True
+    max_sequence_length = 9000
+    
+    def __init__(
+        self,
+        encoder: Optional[dict] = None,
+        decoder: Optional[dict] = None,
+        tie_word_embeddings: bool = False,
+        decoder_start_token_id: int = 2,
+        pad_token_id: int = 1,
+        eos_token_id: int = 2,
+        bos_token_id: int = 0,
+        image_size: List[int] = [2048, 1664],
+        is_encoder_decoder: bool = True,
+        max_sequence_length: int = 9000,
+        **kwargs
+    ):
+        super().__init__(
+            tie_word_embeddings=tie_word_embeddings,
+            decoder_start_token_id=decoder_start_token_id,
+            pad_token_id=pad_token_id,
+            eos_token_id=eos_token_id,
+            bos_token_id=bos_token_id,
+            max_sequence_length=max_sequence_length,
+            **kwargs
+        )
+        
+        
+        if decoder is None:
+            decoder = {}
+            
+        if encoder is not None:
+            assert "auto_map" in encoder and "AutoConfig" in encoder["auto_map"]
+            vision_auto_config = get_class_from_dynamic_module(*encoder["auto_map"]["AutoConfig"].split("--")[::-1])
+            self.encoder = vision_auto_config(**encoder)
+        else:
+            self.encoder = PretrainedConfig()
+
+        decoder["max_sequence_length"] = max_sequence_length
+        self.decoder = NemotronParseTextConfig(**decoder)
+        self.image_size = image_size
+        
+        # Initialize vocab size from text config
+        self.vocab_size = self.decoder.vocab_size
+        self.is_encoder_decoder = is_encoder_decoder
+        self.max_sequence_length = max_sequence_length
+        
+    def to_dict(self):
+        """
+        Serializes this instance to a Python dictionary. Override the default [`~PretrainedConfig.to_dict`].
+        """
+        output = super().to_dict()
+        output["encoder"] = self.encoder.to_dict()
+        output["decoder"] = self.decoder.to_dict()
+        output["model_type"] = self.model_type
+        output["is_encoder_decoder"] = self.is_encoder_decoder
+        return output 

hf_nemotron_parse_modeling.py

@@ -0,0 +1,585 @@
+import math
+import torch
+import torch.nn as nn
+from torch.nn import CrossEntropyLoss
+from transformers import PreTrainedModel, GenerationMixin
+from transformers.models.vision_encoder_decoder.modeling_vision_encoder_decoder import VisionEncoderDecoderModel
+from transformers.models.vision_encoder_decoder.configuration_vision_encoder_decoder import VisionEncoderDecoderConfig
+from transformers.modeling_outputs import Seq2SeqLMOutput
+from transformers.models.mbart.modeling_mbart import MBartPreTrainedModel, MBartConfig, MBartScaledWordEmbedding, MBartDecoderLayer, BaseModelOutputWithPastAndCrossAttentions
+from transformers.models.donut.modeling_donut_swin import DonutSwinModelOutput
+from einops import rearrange
+from typing import Optional, List, Union, Tuple
+import warnings
+from transformers.modeling_outputs import BaseModelOutput
+from transformers.models.encoder_decoder.modeling_encoder_decoder import shift_tokens_right
+from .hf_nemotron_parse_config import NemotronParseConfig
+from transformers import AutoModel
+import time
+from transformers.modeling_attn_mask_utils import (
+    _prepare_4d_attention_mask,
+    _prepare_4d_attention_mask_for_sdpa,
+    _prepare_4d_causal_attention_mask,
+    _prepare_4d_causal_attention_mask_for_sdpa,
+)
+
+
+class NemotronParseDecoder(MBartPreTrainedModel):
+    """
+    Transformer decoder consisting of *config.decoder_layers* layers. Each layer is a [`MBartDecoderLayer`]
+
+    Args:
+        config: MBartConfig
+        embed_tokens (nn.Embedding): output embedding
+    """
+
+    def __init__(self, config: MBartConfig, embed_tokens: Optional[nn.Embedding] = None):
+        super().__init__(config)
+        self.dropout = config.dropout
+        self.layerdrop = config.decoder_layerdrop
+        self.padding_idx = config.pad_token_id
+        embed_scale = math.sqrt(config.d_model) if config.scale_embedding else 1.0
+
+        self.embed_tokens = MBartScaledWordEmbedding(
+            config.vocab_size, config.d_model, self.padding_idx, embed_scale=embed_scale
+        )
+
+        if embed_tokens is not None:
+            self.embed_tokens.weight = embed_tokens.weight
+
+        self.layers = nn.ModuleList([MBartDecoderLayer(config) for _ in range(config.decoder_layers)])
+        self.config = config
+
+        self.layernorm_embedding = nn.LayerNorm(config.d_model)
+        self.layer_norm = nn.LayerNorm(config.d_model)
+
+        self.gradient_checkpointing = False
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    def get_input_embeddings(self):
+        return self.embed_tokens
+
+    def set_input_embeddings(self, value):
+        self.embed_tokens = value
+
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        encoder_hidden_states: Optional[torch.FloatTensor] = None,
+        encoder_attention_mask: Optional[torch.LongTensor] = None,
+        head_mask: Optional[torch.Tensor] = None,
+        cross_attn_head_mask: Optional[torch.Tensor] = None,
+        past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+    ) -> Union[Tuple, BaseModelOutputWithPastAndCrossAttentions]:
+        r"""
+        Args:
+            input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`):
+                Indices of input sequence tokens in the vocabulary. Padding will be ignored by default should you
+                provide it.
+
+                Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and
+                [`PreTrainedTokenizer.__call__`] for details.
+
+                [What are input IDs?](../glossary#input-ids)
+            attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*):
+                Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:
+
+                - 1 for tokens that are **not masked**,
+                - 0 for tokens that are **masked**.
+
+                [What are attention masks?](../glossary#attention-mask)
+            encoder_hidden_states (`torch.FloatTensor` of shape `(batch_size, encoder_sequence_length, hidden_size)`, *optional*):
+                Sequence of hidden-states at the output of the last layer of the encoder. Used in the cross-attention
+                of the decoder.
+            encoder_attention_mask (`torch.LongTensor` of shape `(batch_size, encoder_sequence_length)`, *optional*):
+                Mask to avoid performing cross-attention on padding tokens indices of encoder input_ids. Mask values
+                selected in `[0, 1]`:
+
+                - 1 for tokens that are **not masked**,
+                - 0 for tokens that are **masked**.
+
+                [What are attention masks?](../glossary#attention-mask)
+            head_mask (`torch.Tensor` of shape `(decoder_layers, decoder_attention_heads)`, *optional*):
+                Mask to nullify selected heads of the attention modules. Mask values selected in `[0, 1]`:
+
+                - 1 indicates the head is **not masked**,
+                - 0 indicates the head is **masked**.
+
+            cross_attn_head_mask (`torch.Tensor` of shape `(decoder_layers, decoder_attention_heads)`, *optional*):
+                Mask to nullify selected heads of the cross-attention modules in the decoder to avoid performing
+                cross-attention on hidden heads. Mask values selected in `[0, 1]`:
+
+                - 1 indicates the head is **not masked**,
+                - 0 indicates the head is **masked**.
+
+            past_key_values (`tuple(tuple(torch.FloatTensor))`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`):
+                Tuple of `tuple(torch.FloatTensor)` of length `config.n_layers`, with each tuple having 2 tensors of
+                shape `(batch_size, num_heads, sequence_length, embed_size_per_head)`) and 2 additional tensors of
+                shape `(batch_size, num_heads, encoder_sequence_length, embed_size_per_head)`.
+
+                Contains pre-computed hidden-states (key and values in the self-attention blocks and in the
+                cross-attention blocks) that can be used (see `past_key_values` input) to speed up sequential decoding.
+
+                If `past_key_values` are used, the user can optionally input only the last `decoder_input_ids` (those
+                that don't have their past key value states given to this model) of shape `(batch_size, 1)` instead of
+                all `decoder_input_ids` of shape `(batch_size, sequence_length)`.
+            inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*):
+                Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation.
+                This is useful if you want more control over how to convert `input_ids` indices into associated vectors
+                than the model's internal embedding lookup matrix.
+            output_attentions (`bool`, *optional*):
+                Whether or not to return the attentions tensors of all attention layers. See `attentions` under
+                returned tensors for more detail.
+            output_hidden_states (`bool`, *optional*):
+                Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors
+                for more detail.
+            return_dict (`bool`, *optional*):
+                Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
+        """
+        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+        output_hidden_states = (
+            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        )
+        use_cache = use_cache if use_cache is not None else self.config.use_cache
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        # retrieve input_ids and inputs_embeds
+        if input_ids is not None and inputs_embeds is not None:
+            raise ValueError("You cannot specify both decoder_input_ids and decoder_inputs_embeds at the same time")
+        elif input_ids is not None:
+            input = input_ids
+            input_shape = input.size()
+            input_ids = input_ids.view(-1, input_shape[-1])
+        elif inputs_embeds is not None:
+            input_shape = inputs_embeds.size()[:-1]
+            input = inputs_embeds[:, :, -1]
+        else:
+            raise ValueError("You have to specify either decoder_input_ids or decoder_inputs_embeds")
+
+        # past_key_values_length
+        past_key_values_length = past_key_values[0][0].shape[2] if past_key_values is not None else 0
+
+        if inputs_embeds is None:
+            inputs_embeds = self.embed_tokens(input_ids)
+
+        if self.config._attn_implementation == "flash_attention_2":
+            # 2d mask is passed through the layers
+            attention_mask = attention_mask if (attention_mask is not None and 0 in attention_mask) else None
+        elif self.config._attn_implementation == "sdpa" and not output_attentions and cross_attn_head_mask is None:
+            # output_attentions=True & cross_attn_head_mask can not be supported when using SDPA, and we fall back on
+            # the manual implementation that requires a 4D causal mask in all cases.
+            attention_mask = _prepare_4d_causal_attention_mask_for_sdpa(
+                attention_mask,
+                input_shape,
+                inputs_embeds,
+                past_key_values_length,
+            )
+        else:
+            # 4d mask is passed through the layers
+            attention_mask = _prepare_4d_causal_attention_mask(
+                attention_mask, input_shape, inputs_embeds, past_key_values_length
+            )
+
+        # expand encoder attention mask
+        if encoder_hidden_states is not None and encoder_attention_mask is not None:
+            if self.config._attn_implementation == "flash_attention_2":
+                encoder_attention_mask = encoder_attention_mask if 0 in encoder_attention_mask else None
+            elif self.config._attn_implementation == "sdpa" and cross_attn_head_mask is None and not output_attentions:
+                # output_attentions=True & cross_attn_head_mask can not be supported when using SDPA, and we fall back on
+                # the manual implementation that requires a 4D causal mask in all cases.
+                # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len]
+                encoder_attention_mask = _prepare_4d_attention_mask_for_sdpa(
+                    encoder_attention_mask,
+                    inputs_embeds.dtype,
+                    tgt_len=input_shape[-1],
+                )
+            else:
+                # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len]
+                encoder_attention_mask = _prepare_4d_attention_mask(
+                    encoder_attention_mask, inputs_embeds.dtype, tgt_len=input_shape[-1]
+                )
+        hidden_states = inputs_embeds
+        hidden_states = self.layernorm_embedding(hidden_states)
+
+        hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)
+
+        if self.gradient_checkpointing and self.training:
+            if use_cache:
+                logger.warning_once(
+                    "`use_cache=True` is incompatible with gradient checkpointing`. Setting `use_cache=False`..."
+                )
+                use_cache = False
+
+        # decoder layers
+        all_hidden_states = () if output_hidden_states else None
+        all_self_attns = () if output_attentions else None
+        all_cross_attentions = () if (output_attentions and encoder_hidden_states is not None) else None
+        next_decoder_cache = () if use_cache else None
+
+        # check if head_mask/cross_attn_head_mask has a correct number of layers specified if desired
+        for attn_mask, mask_name in zip([head_mask, cross_attn_head_mask], ["head_mask", "cross_attn_head_mask"]):
+            if attn_mask is not None:
+                if attn_mask.size()[0] != len(self.layers):
+                    raise ValueError(
+                        f"The `{mask_name}` should be specified for {len(self.layers)} layers, but it is for"
+                        f" {attn_mask.size()[0]}."
+                    )
+        for idx, decoder_layer in enumerate(self.layers):
+            # add LayerDrop (see https://arxiv.org/abs/1909.11556 for description)
+            if output_hidden_states:
+                all_hidden_states += (hidden_states,)
+            if self.training:
+                dropout_probability = torch.rand([])
+                if dropout_probability < self.layerdrop:
+                    continue
+
+            past_key_value = past_key_values[idx] if past_key_values is not None else None
+
+            if self.gradient_checkpointing and self.training:
+                layer_outputs = self._gradient_checkpointing_func(
+                    decoder_layer.__call__,
+                    hidden_states,
+                    attention_mask,
+                    encoder_hidden_states,
+                    encoder_attention_mask,
+                    head_mask[idx] if head_mask is not None else None,
+                    cross_attn_head_mask[idx] if cross_attn_head_mask is not None else None,
+                    None,
+                    output_attentions,
+                    use_cache,
+                )
+            else:
+                layer_outputs = decoder_layer(
+                    hidden_states,
+                    attention_mask=attention_mask,
+                    encoder_hidden_states=encoder_hidden_states,
+                    encoder_attention_mask=encoder_attention_mask,
+                    layer_head_mask=(head_mask[idx] if head_mask is not None else None),
+                    cross_attn_layer_head_mask=(
+                        cross_attn_head_mask[idx] if cross_attn_head_mask is not None else None
+                    ),
+                    past_key_value=past_key_value,
+                    output_attentions=output_attentions,
+                    use_cache=use_cache,
+                )
+            hidden_states = layer_outputs[0]
+
+            if use_cache:
+                next_decoder_cache += (layer_outputs[3 if output_attentions else 1],)
+
+            if output_attentions:
+                all_self_attns += (layer_outputs[1],)
+
+                if encoder_hidden_states is not None:
+                    all_cross_attentions += (layer_outputs[2],)
+
+        hidden_states = self.layer_norm(hidden_states)
+
+        # add hidden states from the last decoder layer
+        if output_hidden_states:
+            all_hidden_states += (hidden_states,)
+
+        next_cache = next_decoder_cache if use_cache else None
+        if not return_dict:
+            return tuple(
+                v
+                for v in [hidden_states, next_cache, all_hidden_states, all_self_attns, all_cross_attentions]
+                if v is not None
+            )
+        return BaseModelOutputWithPastAndCrossAttentions(
+            last_hidden_state=hidden_states,
+            past_key_values=next_cache,
+            hidden_states=all_hidden_states,
+            attentions=all_self_attns,
+            cross_attentions=all_cross_attentions,
+        )
+
+
+class RadioWithNeck(nn.Module):
+    """Vision encoder using RADIO model with custom neck."""
+    
+    def __init__(self, config):
+        super().__init__()
+        self.config = config
+        
+        self.model_encoder = AutoModel.from_config(config, trust_remote_code=True)
+        
+        # Neck components
+        last_hidden_state = 1024
+        self.conv1 = nn.Conv1d(1280, last_hidden_state, 1)
+        self.layer_norm1 = nn.LayerNorm(last_hidden_state, eps=1e-06, elementwise_affine=True)
+        self.conv2 = nn.Conv2d(last_hidden_state, last_hidden_state, kernel_size=(1,4), stride=(1,4), padding=0, bias=False)
+        self.layer_norm2 = nn.LayerNorm(last_hidden_state, eps=1e-06, elementwise_affine=True)
+        self.sum_proj = nn.Linear(3840, last_hidden_state)
+        self.layer_norm3 = nn.LayerNorm(last_hidden_state, eps=1e-06, elementwise_affine=True)
+        
+    def forward(self, pixel_values, output_attentions=False, output_hidden_states=False, return_dict=False, **kwargs):
+        radio_output = self.model_encoder(pixel_values)
+        summary, feature = radio_output
+        
+        
+        output = self.conv1(feature.permute(0,2,1)).permute(0,2,1)
+        output = self.layer_norm1(output)
+        
+        patch_size = self.config.patch_size
+        output = rearrange(output, 'b (h w) d -> b d h w',
+                    h=pixel_values.shape[-2] // patch_size,
+                    w=pixel_values.shape[-1] // patch_size)
+        
+        output = self.conv2(output)
+        output = rearrange(output, 'b d h w -> b (h w) d')
+        output = self.layer_norm2(output)
+        summary = self.layer_norm3(self.sum_proj(summary))
+        output = torch.cat((output, summary.unsqueeze(1)), dim=1)
+        
+        return DonutSwinModelOutput(last_hidden_state=output)
+
+
+class NemotronParsePreTrainedModel(PreTrainedModel):
+    """
+    An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained models.
+    """
+    config_class = NemotronParseConfig
+    base_model_prefix = "vision_encoder_decoder"  # Use VisionEncoderDecoder prefix
+    main_input_name = "pixel_values"
+    supports_gradient_checkpointing = True
+    _no_split_modules = ["RadioWithNeck", "MBartDecoder"]
+    _skip_keys_device_placement = "past_key_values"
+
+    def _init_weights(self, module):
+        """Initialize the weights"""
+        if isinstance(module, nn.Linear):
+            module.weight.data.normal_(mean=0.0, std=self.config.decoder.init_std)
+            if module.bias is not None:
+                module.bias.data.zero_()
+        elif isinstance(module, nn.Embedding):
+            module.weight.data.normal_(mean=0.0, std=self.config.decoder.init_std)
+            if module.padding_idx is not None:
+                module.weight.data[module.padding_idx].zero_()
+
+# Based on transformers.models.encoder_decoder.modeling_encoder_decoder
+class NemotronParseForConditionalGeneration(NemotronParsePreTrainedModel, GenerationMixin):
+    """
+    NemotronParse model for conditional generation tasks.
+    
+    This model combines a RADIO-based vision encoder with an mBART-based text decoder.
+    """
+    
+    def __init__(self, config: NemotronParseConfig):
+        super().__init__(config)
+        
+        self.encoder = RadioWithNeck(config.encoder)
+        self.encoder.main_input_name = 'pixel_values'
+        self.encoder = self.encoder.to(config.encoder.torch_dtype)
+        
+        self.decoder = NemotronParseDecoder(config.decoder)
+        self.decoder = self.decoder.to(config.decoder.torch_dtype)
+
+        self.lm_head = nn.Linear(config.decoder.d_model, config.decoder.vocab_size, bias=False, dtype=config.decoder.torch_dtype)
+ 
+        # Extra heads
+        num_extra_heads = getattr(config, 'num_extra_heads', 0)
+        self.decoder.extra_heads = nn.ModuleList([
+            nn.Linear(config.decoder.d_model, config.decoder.d_model) 
+            for _ in range(num_extra_heads)
+        ])
+        self.decoder.extra_proj = nn.ModuleList([
+            nn.Linear(config.decoder.d_model, config.decoder.d_model)
+            for _ in range(num_extra_heads)
+        ])
+        
+        # Class token index for loss weighting
+        self.class_token_indx_start = getattr(config, 'class_token_start_idx', 50000)
+        
+        self.post_init()
+
+    def get_encoder(self):
+        return self.encoder
+
+    def get_decoder(self):
+        return self.decoder
+
+    def get_output_embeddings(self):
+        return self.lm_head
+
+    def set_output_embeddings(self, new_embeddings):
+        self.lm_head = new_embeddings
+    
+    def get_input_embeddings(self):
+        return self.decoder.get_input_embeddings()
+
+    def forward(
+        self,
+        pixel_values: Optional[torch.FloatTensor] = None,
+        decoder_input_ids: Optional[torch.LongTensor] = None,
+        decoder_attention_mask: Optional[torch.BoolTensor] = None,
+        encoder_outputs: Optional[Tuple[torch.FloatTensor]] = None,
+        past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None,
+        decoder_inputs_embeds: Optional[torch.FloatTensor] = None,
+        labels: Optional[torch.LongTensor] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+        __subflavors__: Optional[str] = None,
+        __keys__: Optional[List[str]] = None,
+        return_sample_losses: Optional[torch.FloatTensor] = None,
+        **kwargs,
+    ) -> Union[Tuple[torch.FloatTensor], Seq2SeqLMOutput]:
+        
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        kwargs_encoder = {argument: value for argument, value in kwargs.items() if not argument.startswith("decoder_")}
+
+        kwargs_decoder = {
+            argument[len("decoder_") :]: value for argument, value in kwargs.items() if argument.startswith("decoder_")
+        }
+
+        if encoder_outputs is None:
+            if pixel_values is None:
+                raise ValueError("You have to specify pixel_values")
+
+            encoder_outputs = self.encoder(
+                pixel_values,
+                output_attentions=output_attentions,
+                output_hidden_states=output_hidden_states,
+                return_dict=return_dict,
+                **kwargs_encoder,
+            )
+
+        elif isinstance(encoder_outputs, tuple):
+            encoder_outputs = BaseModelOutput(*encoder_outputs)
+
+        encoder_hidden_states = encoder_outputs[0]
+
+        encoder_attention_mask = None
+
+        if (labels is not None) and (decoder_input_ids is None and decoder_inputs_embeds is None):
+            decoder_input_ids = shift_tokens_right(
+                labels, self.config.pad_token_id, self.config.decoder_start_token_id
+            )
+
+        output_hidden_states = True
+
+        decoder_outputs = self.decoder(
+            input_ids=decoder_input_ids,
+            attention_mask=decoder_attention_mask,
+            encoder_hidden_states=encoder_hidden_states,
+            encoder_attention_mask=encoder_attention_mask,
+            inputs_embeds=decoder_inputs_embeds,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            use_cache=use_cache,
+            past_key_values=past_key_values,
+            return_dict=return_dict,
+            **kwargs_decoder,
+        )
+        loss = None
+
+        if labels is not None:
+            main_logits = self.lm_head(decoder_outputs.last_hidden_state)
+            logits = [main_logits]
+            decoder_inputs_embeds = decoder_outputs.inputs_embeds
+            for iii, head in enumerate(self.decoder.extra_heads):
+
+                decoder_input_embeds_shift = self.decoder.extra_proj[iii](torch.cat((decoder_inputs_embeds[:,1:,:], torch.zeros_like(decoder_inputs_embeds[:,0,:].unsqueeze(1))), axis=1))
+                hidden = head(decoder_outputs['hidden_states'][-1] + decoder_input_embeds_shift)
+                logits.append(self.lm_head(hidden))  # Use main lm_head, NOT decoder.lm_head
+
+            logits = torch.stack(logits, dim=-2)
+            loss_fct = CrossEntropyLoss(reduction="none")
+
+            losses_per_head = []
+            tokens_per_head = []
+            for head_num in range(len(self.decoder.extra_heads)+1):
+                logits_head = logits[:,:,head_num,:]
+                labels_head = torch.cat(
+                    (labels[:, head_num:], torch.full_like(labels[:, :head_num], -100)),
+                    1
+                )
+                loss_full = loss_fct(logits_head.permute(0, 2, 1), labels_head)
+                loss_full[labels_head >= self.class_token_indx_start] *= 10
+                losses_per_head.append(loss_full.sum(1))
+                tokens_per_head.append((labels_head != -100).sum(1))
+
+            losses_per_sample = torch.stack(losses_per_head, dim=1).sum(1)
+            tokens_per_sample = torch.stack(tokens_per_head, dim=1).sum(1)
+            loss = losses_per_sample.sum() / (tokens_per_sample.sum() + 1e-6)
+            if return_sample_losses is not None:
+                return_sample_losses.copy_(losses_per_sample.detach() / (tokens_per_sample + 1e-6))
+
+        if not return_dict:
+            if loss is not None:
+                return (loss,) + decoder_outputs + encoder_outputs
+            else:
+                return decoder_outputs + encoder_outputs
+        output_logits = self.lm_head(decoder_outputs.last_hidden_state)
+        return Seq2SeqLMOutput(
+            loss=loss,
+            logits=output_logits, 
+            past_key_values=decoder_outputs.past_key_values,
+            decoder_hidden_states=decoder_outputs.hidden_states,
+            decoder_attentions=decoder_outputs.attentions,
+            cross_attentions=decoder_outputs.cross_attentions,
+            encoder_last_hidden_state=encoder_outputs.last_hidden_state,
+            encoder_hidden_states=encoder_outputs.hidden_states,
+            encoder_attentions=encoder_outputs.attentions,
+        )
+
+    def prepare_decoder_input_ids_from_labels(self, labels: torch.Tensor):
+        return shift_tokens_right(labels, self.config.pad_token_id, self.config.decoder_start_token_id)
+
+
+    def resize_token_embeddings(self, new_num_tokens: Optional[int] = None, pad_to_multiple_of: Optional[int] = None):
+        """Resize token embeddings and update lm_head accordingly."""
+        # Resize decoder embeddings
+        new_embeddings = self.decoder.resize_token_embeddings(new_num_tokens, pad_to_multiple_of)
+        
+        # Update lm_head to match new vocab size
+        if new_embeddings is not None:
+            old_vocab_size, hidden_size = self.lm_head.weight.shape
+            new_vocab_size = new_embeddings.num_embeddings
+            
+            if old_vocab_size != new_vocab_size:
+                print(f"Resizing lm_head from {old_vocab_size} to {new_vocab_size} tokens")
+                new_lm_head = nn.Linear(hidden_size, new_vocab_size, bias=False, device=self.lm_head.weight.device, dtype=self.lm_head.weight.dtype)
+                
+                # Copy old weights to new lm_head
+                num_tokens_to_copy = min(old_vocab_size, new_vocab_size)
+                new_lm_head.weight.data[:num_tokens_to_copy] = self.lm_head.weight.data[:num_tokens_to_copy]
+                
+                # Update reference
+                self.lm_head = new_lm_head
+                # DO NOT update decoder.lm_head - keep them separate
+        
+        return new_embeddings
+
+    def _reorder_cache(self, past_key_values, beam_idx):
+        # apply decoder cache reordering here
+        return self.decoder._reorder_cache(past_key_values, beam_idx)
+
+
+# Copied from transformers.models.encoder_decoder.modeling_encoder_decoder.shift_tokens_right
+def shift_tokens_right(input_ids: torch.Tensor, pad_token_id: int, decoder_start_token_id: int):
+    """
+    Shift input ids one token to the right.
+    """
+    shifted_input_ids = input_ids.new_zeros(input_ids.shape)
+    shifted_input_ids[:, 1:] = input_ids[:, :-1].clone()
+    if decoder_start_token_id is None:
+        raise ValueError("Make sure to set the decoder_start_token_id attribute of the model's configuration.")
+    shifted_input_ids[:, 0] = decoder_start_token_id
+
+    if pad_token_id is None:
+        raise ValueError("Make sure to set the pad_token_id attribute of the model's configuration.")
+    # replace possible -100 values in labels by `pad_token_id`
+    shifted_input_ids.masked_fill_(shifted_input_ids == -100, pad_token_id)
+
+    return shifted_input_ids

hf_nemotron_parse_processor.py

@@ -0,0 +1,396 @@
+import numpy as np
+from PIL import Image
+from typing import List, Optional, Union, Dict, Any
+import torch
+from torchvision import transforms as T
+import albumentations as A
+import cv2
+import json
+
+from transformers import ProcessorMixin, BaseImageProcessor, ImageProcessingMixin
+from transformers.tokenization_utils_base import BatchEncoding
+from transformers.image_utils import ChannelDimension, ImageInput, PILImageResampling, infer_channel_dimension_format
+from transformers.utils import TensorType
+
+
+class NemotronParseImageProcessor(BaseImageProcessor, ImageProcessingMixin):
+    """
+    Image processor for NemotronParse model.
+    
+    This processor inherits from BaseImageProcessor to be compatible with transformers AutoImageProcessor.
+    """
+    
+    model_input_names = ["pixel_values"]
+    
+    def __init__(
+        self,
+        final_size: tuple = (2048, 1664),
+        **kwargs,
+    ):
+        clean_kwargs = {}
+        for k, v in kwargs.items():
+            if not k.startswith('_') and k not in ['transform', 'torch_transform']:
+                clean_kwargs[k] = v
+        
+        if 'size' in clean_kwargs:
+            size_config = clean_kwargs.pop('size')
+            if isinstance(size_config, dict):
+                if 'longest_edge' in size_config:
+                    longest_edge = size_config['longest_edge']
+                    if isinstance(longest_edge, (list, tuple)):
+                        final_size = tuple(int(x) for x in longest_edge)
+                    else:
+                        final_size = (int(longest_edge), int(longest_edge))
+                elif 'height' in size_config and 'width' in size_config:
+                    final_size = (int(size_config['height']), int(size_config['width']))
+        
+        super().__init__(**clean_kwargs)
+        
+        if isinstance(final_size, (list, tuple)) and len(final_size) >= 2:
+            self.final_size = (int(final_size[0]), int(final_size[1]))
+        elif isinstance(final_size, (int, float)):
+            self.final_size = (int(final_size), int(final_size))
+        else:
+            self.final_size = (2048, 1664)  # Default fallback
+        
+        self._create_transforms()
+    
+    def _create_transforms(self):
+        """Create transform objects (not serialized to JSON)."""
+        if isinstance(self.final_size, (list, tuple)):
+            self.target_height, self.target_width = int(self.final_size[0]), int(self.final_size[1])
+        else:
+            self.target_height = self.target_width = int(self.final_size)
+        
+        self.transform = A.Compose([
+            A.PadIfNeeded(
+                min_height=self.target_height, 
+                min_width=self.target_width, 
+                border_mode=cv2.BORDER_CONSTANT, 
+                value=[255, 255, 255],
+                p=1.0
+            ),
+        ])
+        
+        self.torch_transform = T.Compose([
+            T.ToTensor(),
+            # Note: Normalization is done within RADIO model
+        ])
+
+    def to_dict(self):
+        """Override to exclude non-serializable transforms."""
+        output = super().to_dict()
+        output.pop('transform', None)
+        output.pop('torch_transform', None)
+        return output
+    
+    @classmethod
+    def from_dict(cls, config_dict: dict, **kwargs):
+        """Override to recreate transforms after loading."""
+        config_dict = config_dict.copy()
+        config_dict.pop('transform', None)
+        config_dict.pop('torch_transform', None)
+        
+        # Clean any problematic entries
+        for key in list(config_dict.keys()):
+            if key.startswith('_') or config_dict[key] is None:
+                config_dict.pop(key, None)
+        
+        # Ensure numeric types are correct
+        if 'final_size' in config_dict:
+            final_size = config_dict['final_size']
+            if isinstance(final_size, (list, tuple)):
+                config_dict['final_size'] = tuple(int(x) for x in final_size)
+        
+        try:
+            return cls(**config_dict, **kwargs)
+        except Exception as e:
+            print(f"Warning: Error in from_dict: {e}")
+            print("Using default parameters...")
+            return cls(**kwargs)
+    
+    def save_pretrained(self, save_directory, **kwargs):
+        """Save image processor configuration."""
+        import os
+        import json
+        
+        os.makedirs(save_directory, exist_ok=True)
+        
+        # Save preprocessor config in standard HuggingFace format
+        config = {
+            "feature_extractor_type": "NemotronParseImageProcessor",
+            "image_processor_type": "NemotronParseImageProcessor", 
+            "processor_class": "NemotronParseImageProcessor",
+            "size": {
+                "height": self.final_size[0],
+                "width": self.final_size[1],
+                "longest_edge": self.final_size
+            },
+            "final_size": self.final_size,
+        }
+        
+        config_path = os.path.join(save_directory, "preprocessor_config.json")
+        with open(config_path, 'w') as f:
+            json.dump(config, f, indent=2)
+
+    def _resize_with_aspect_ratio(self, image: np.ndarray) -> np.ndarray:
+        """Resize image maintaining aspect ratio (exact replica of original LongestMaxSizeHW)."""
+        height, width = image.shape[:2]
+        max_size_height = self.target_height
+        max_size_width = self.target_width
+        
+        # Original LongestMaxSizeHW algorithm from custom_augmentations.py
+        aspect_ratio = width / height
+        new_height = height
+        new_width = width
+
+        if height > max_size_height:
+            new_height = max_size_height
+            new_width = int(new_height * aspect_ratio)
+
+        if new_width > max_size_width:
+            new_width = max_size_width
+            new_height = int(new_width / aspect_ratio)
+        
+        return cv2.resize(image, (new_width, new_height), interpolation=cv2.INTER_LINEAR)
+    
+    def _pad_to_size(self, image: np.ndarray) -> np.ndarray:
+        """Pad image to target size with white padding (matches A.PadIfNeeded behavior)."""
+        h, w = image.shape[:2]
+        min_height, min_width = self.target_height, self.target_width
+        
+        pad_h = max(0, min_height - h)
+        pad_w = max(0, min_width - w)
+        
+        if pad_h == 0 and pad_w == 0:
+            return image
+        
+        if len(image.shape) == 3:
+            padded = np.pad(
+                image, 
+                ((0, pad_h), (0, pad_w), (0, 0)), 
+                mode='constant', 
+                constant_values=255
+            )
+        else:
+            padded = np.pad(
+                image, 
+                ((0, pad_h), (0, pad_w)), 
+                mode='constant', 
+                constant_values=255
+            )
+        
+        return padded
+    
+    def preprocess(
+        self,
+        images: ImageInput,
+        return_tensors: Optional[Union[str, TensorType]] = None,
+        **kwargs,
+    ) -> Dict[str, torch.Tensor]:
+        """
+        Preprocess an image or batch of images for the NemotronParse model.
+        
+        Args:
+            images: Input image(s)
+            return_tensors: Type of tensors to return
+        """
+        
+        # Ensure images is a list
+        if not isinstance(images, list):
+            images = [images]
+
+        # Ensure images are RGB
+        for i in range(len(images)):
+            images[i] = images[i].convert('RGB')
+        
+        # Convert PIL images to numpy arrays if needed
+        processed_images = []
+        for image in images:
+            if isinstance(image, Image.Image):
+                image = np.asarray(image)
+            processed_images.append(image)
+        
+        # Apply NemotronParse-specific transforms
+        pixel_values = []
+        for image in processed_images:
+            processed_image = self._resize_with_aspect_ratio(image)
+            
+            if self.transform is not None:
+                transformed = self.transform(image=processed_image)
+                processed_image = transformed["image"]
+            else:
+                # Fallback: just pad to target size
+                processed_image = self._pad_to_size(processed_image)
+            
+            pixel_values_tensor = self.torch_transform(processed_image)
+            
+            if pixel_values_tensor.shape[0] == 1:
+                pixel_values_tensor = pixel_values_tensor.expand(3, -1, -1)
+                
+            pixel_values.append(pixel_values_tensor)
+                        
+        pixel_values = torch.stack(pixel_values)
+        
+        data = {"pixel_values": pixel_values}
+        
+        if return_tensors is not None:
+            data = self._convert_output_format(data, return_tensors)
+            
+        return data
+    
+    def _convert_output_format(self, data: Dict[str, torch.Tensor], return_tensors: Union[str, TensorType]) -> Dict:
+        """Convert output format based on return_tensors parameter."""
+        if return_tensors == "pt" or return_tensors == TensorType.PYTORCH:
+            return data
+        elif return_tensors == "np" or return_tensors == TensorType.NUMPY:
+            return {k: v.numpy() for k, v in data.items()}
+        else:
+            return data
+    
+    def __call__(self, images: Union[Image.Image, List[Image.Image]], **kwargs) -> Dict[str, torch.Tensor]:
+        """Process images for the model (backward compatibility)."""
+        return self.preprocess(images, **kwargs)
+
+
+class NemotronParseProcessor(ProcessorMixin):
+    
+    attributes = ["image_processor", "tokenizer"]
+    image_processor_class = "AutoImageProcessor"
+    tokenizer_class = ("PreTrainedTokenizer", "PreTrainedTokenizerFast")
+    
+    def __init__(self, image_processor=None, tokenizer=None, **kwargs):
+        if image_processor is None:
+            image_processor = NemotronParseImageProcessor(**kwargs)
+               
+        super().__init__(image_processor, tokenizer)
+            
+    
+    def __call__(
+        self,
+        images: Union[Image.Image, List[Image.Image]] = None,
+        text: Union[str, List[str]] = None,
+        add_special_tokens: bool = True,
+        padding: Union[bool, str] = False,
+        truncation: Union[bool, str] = False,
+        max_length: Optional[int] = None,
+        stride: int = 0,
+        pad_to_multiple_of: Optional[int] = None,
+        return_attention_mask: Optional[bool] = None,
+        return_overflowing_tokens: bool = False,
+        return_special_tokens_mask: bool = False,
+        return_offsets_mapping: bool = False,
+        return_token_type_ids: bool = False,
+        return_length: bool = False,
+        verbose: bool = True,
+        return_tensors: Optional[Union[str, "TensorType"]] = None,
+        **kwargs
+    ) -> BatchEncoding:
+        """
+        Main method to prepare for the model one or several text(s) and image(s).
+        """
+        
+        # Process images
+        if images is not None:
+            image_inputs = self.image_processor(images, **kwargs)
+        else:
+            image_inputs = {}
+        
+        # Process text
+        if text is not None:
+            text_inputs = self.tokenizer(
+                text,
+                add_special_tokens=add_special_tokens,
+                padding=padding,
+                truncation=truncation,
+                max_length=max_length,
+                stride=stride,
+                pad_to_multiple_of=pad_to_multiple_of,
+                return_attention_mask=return_attention_mask,
+                return_overflowing_tokens=return_overflowing_tokens,
+                return_special_tokens_mask=return_special_tokens_mask,
+                return_offsets_mapping=return_offsets_mapping,
+                return_token_type_ids=return_token_type_ids,
+                return_length=return_length,
+                verbose=verbose,
+                return_tensors=return_tensors,
+                **kwargs,
+            )
+        else:
+            text_inputs = {}
+        
+        # Combine inputs
+        return BatchEncoding({**image_inputs, **text_inputs})
+    
+    def decode(self, *args, **kwargs):
+        """Decode token ids to strings."""
+        return self.tokenizer.decode(*args, **kwargs)
+    
+    def batch_decode(self, *args, **kwargs):
+        """Batch decode token ids to strings."""
+        return self.tokenizer.batch_decode(*args, **kwargs)
+    
+    def post_process_generation(self, sequences, fix_markdown=False):
+        """Post-process generated sequences."""
+        if hasattr(self.tokenizer, 'post_process_generation'):
+            return self.tokenizer.post_process_generation(sequences, fix_markdown=fix_markdown)
+        else:
+            # Fallback processing
+            if isinstance(sequences, str):
+                sequences = [sequences]
+            
+            processed = []
+            for seq in sequences:
+                # Basic cleaning
+                seq = seq.replace('<s>', '').replace('</s>', '').strip()
+                processed.append(seq)
+            
+            return processed[0] if len(processed) == 1 else processed
+    
+    @classmethod
+    def from_pretrained(cls, pretrained_model_name_or_path, **kwargs):
+        """
+        Load processor from pretrained model.
+        
+        This method is compatible with AutoProcessor.from_pretrained().
+        """
+        # Explicitly load subcomponents via Auto* to ensure remote auto_map is honored.
+        from transformers import AutoImageProcessor, AutoTokenizer
+        trust_remote_code = kwargs.get("trust_remote_code", None)
+        revision = kwargs.get("revision", None)
+        token = kwargs.get("token", None)
+        image_processor = AutoImageProcessor.from_pretrained(
+            pretrained_model_name_or_path,
+            trust_remote_code=trust_remote_code,
+            revision=revision,
+            token=token,
+        )
+        tokenizer = AutoTokenizer.from_pretrained(
+            pretrained_model_name_or_path,
+            trust_remote_code=trust_remote_code,
+            revision=revision,
+            token=token,
+        )
+        return cls(image_processor=image_processor, tokenizer=tokenizer)
+    
+    def save_pretrained(self, save_directory, **kwargs):
+        """
+        Save processor to directory.
+        
+        This method is compatible with AutoProcessor/AutoImageProcessor loading.
+        """
+        import os
+        os.makedirs(save_directory, exist_ok=True)
+        
+        # Save tokenizer with proper configuration for AutoTokenizer
+        print("Saving tokenizer for AutoTokenizer compatibility...")
+        self.tokenizer.save_pretrained(save_directory, **kwargs)
+        
+        # Save image processor  
+        print("Saving image processor...")
+        self.image_processor.save_pretrained(save_directory, **kwargs)
+        
+        # Use the parent class's save_pretrained method for processor config
+        super().save_pretrained(save_directory, **kwargs)
+        print(f"NemotronParseProcessor saved to {save_directory}")
+        print(f"AutoTokenizer.from_pretrained('{save_directory}') should now work!")

latex2html.py

@@ -0,0 +1,448 @@
+import re
+from bs4 import BeautifulSoup
+
+def skip_whitespace(text, i):
+    """Advance index i past any whitespace."""
+    while i < len(text) and text[i].isspace():
+        i += 1
+    return i
+
+def parse_braced_argument(text, i):
+    """
+    Given text and an index i that should point at an opening '{',
+    return a tuple (argument_content, new_index) where argument_content is the full
+    string inside the balanced braces and new_index is the position just after the matching '}'.
+    """
+    if i >= len(text) or text[i] != '{':
+        raise ValueError("Expected '{' at position {}".format(i))
+    i += 1  # skip the opening brace
+    start = i
+    level = 1
+    while i < len(text) and level > 0:
+        if text[i] == '{':
+            level += 1
+        elif text[i] == '}':
+            level -= 1
+        i += 1
+    if level != 0:
+        raise ValueError("Unbalanced braces starting at position {}".format(start-1))
+    # The argument content is from start to i-1 (excluding the closing brace)
+    return text[start:i-1], i
+
+def parse_command(text, i):
+    """
+    Parse a \multirow or \multicolumn command starting at index i.
+    This function assumes the command has exactly three braced arguments.
+
+    It processes each argument recursively. For the third argument, after recursive processing,
+    it replaces any unescaped & with \&.
+
+    Returns a tuple (command_text, new_index) where command_text is the reconstructed command.
+    """
+    # Determine which command we have.
+    if text.startswith(r"\multirow", i):
+        command_name = r"\multirow"
+        i += len(r"\multirow")
+    elif text.startswith(r"\multicolumn", i):
+        command_name = r"\multicolumn"
+        i += len(r"\multicolumn")
+    else:
+        raise ValueError("Expected \\multirow or \\multicolumn at position {}".format(i))
+
+    # Skip whitespace between the command name and the first argument.
+    i = skip_whitespace(text, i)
+    args = []
+    # Expect exactly three arguments
+    for arg_index in range(3):
+        if i >= len(text) or text[i] != '{':
+            raise ValueError("Expected '{' for argument {} at position {}".format(arg_index+1, i))
+        arg_content, i = parse_braced_argument(text, i)
+        # Process the content recursively to catch nested commands
+        processed_arg = clean_multi_cells(arg_content)
+        if arg_index == 2:
+            # For the cell text (third argument), replace any unescaped &
+            processed_arg = re.sub(r'(?<!\\)&', r'\\&', processed_arg)
+        args.append(processed_arg)
+        # Only skip whitespace between arguments, not after the last one.
+        if arg_index < 2:
+            i = skip_whitespace(text, i)
+    # Reconstruct the full command with its three arguments
+    command_text = f"{command_name}{{{args[0]}}}{{{args[1]}}}{{{args[2]}}}"
+    return command_text, i
+
+def clean_multi_cells(text):
+    """
+    Process an arbitrary LaTeX text string and look for occurrences of \multirow or \multicolumn commands.
+    When found, the command is parsed (handling nested braces and nested commands) and its third argument is fixed.
+
+    Returns the processed text.
+    """
+    result = []
+    i = 0
+    while i < len(text):
+        # Find next occurrence of either command.
+        idx_multi = text.find(r"\multirow", i)
+        idx_multiC = text.find(r"\multicolumn", i)
+
+        # Determine the next index among the two (if any)
+        if idx_multi == -1 and idx_multiC == -1:
+            result.append(text[i:])
+            break
+        if idx_multi == -1:
+            next_idx = idx_multiC
+        elif idx_multiC == -1:
+            next_idx = idx_multi
+        else:
+            next_idx = min(idx_multi, idx_multiC)
+
+        # Append text before the command (preserving any whitespace)
+        result.append(text[i:next_idx])
+        # Process the command starting at next_idx
+        command_text, new_index = parse_command(text, next_idx)
+        result.append(command_text)
+        i = new_index
+    return ''.join(result)
+
+def parse_brace(s, pos):
+    """
+    Given a string s and an index pos pointing to an opening '{',
+    returns a tuple (content, new_pos) where content is the string
+    between the matching braces (handling nested braces) and new_pos is
+    the index just after the closing '}'.
+    """
+    if pos >= len(s) or s[pos] != '{':
+        raise ValueError("Expected '{' at position %d" % pos)
+    pos += 1  # skip the opening brace
+    content = ""
+    depth = 1
+    while pos < len(s) and depth:
+        char = s[pos]
+        if char == '{':
+            depth += 1
+            content += char
+        elif char == '}':
+            depth -= 1
+            if depth:
+                content += char
+        else:
+            content += char
+        pos += 1
+    if depth != 0:
+        raise ValueError("Unmatched '{' in string.")
+    return content, pos
+
+def parse_command_merge(s, pos):
+    """
+    Parse a multirow or multicolumn command starting at s[pos]. If the content
+    of the command contains a nested command, then recursively parse the inner
+    command and merge its parameters with the outer ones. The merging is done
+    so that the outer multirow’s parameters (e.g. rowspan and width) are kept
+    while the inner command’s parameters (e.g. colspan, alignment) and its innermost
+    content are returned.
+
+    Returns a tuple (merged_dict, new_pos) where merged_dict is a dictionary
+    containing the combined parameters and new_pos is the updated index after
+    parsing the command.
+    """
+    if s.startswith(r"\multirow", pos):
+        newpos = pos + len(r"\multirow")
+        # Parse the three required arguments for multirow: rowspan, width, and content.
+        rowspan, newpos = parse_brace(s, newpos)
+        width, newpos = parse_brace(s, newpos)
+        content, newpos = parse_brace(s, newpos)
+        # Look for a nested command (either \multirow or \multicolumn) in the content.
+        index_mr = content.find(r"\multirow")
+        index_mc = content.find(r"\multicolumn")
+        if index_mr == -1 and index_mc == -1:
+            # No nested command found; return this command’s details.
+            return {"rowspan": rowspan.strip(), "width": width.strip(), "content": content.strip()}, newpos
+        else:
+            # At least one nested command is present. Pick the first occurrence.
+            indices = [i for i in (index_mr, index_mc) if i != -1]
+            first_index = min(indices)
+            # Parse the inner (nested) command from within the content.
+            inner, _ = parse_command_merge(content, first_index)
+            # Merge: keep the outer multirow’s parameters and add the inner ones.
+            merged = {"rowspan": rowspan.strip(), "width": width.strip()}
+            merged.update(inner)
+            return merged, newpos
+
+    elif s.startswith(r"\multicolumn", pos):
+        newpos = pos + len(r"\multicolumn")
+        # Parse the three arguments for multicolumn: colspan, alignment, and content.
+        colspan, newpos = parse_brace(s, newpos)
+        alignment, newpos = parse_brace(s, newpos)
+        content, newpos = parse_brace(s, newpos)
+        # Look for a nested command in the content.
+        index_mr = content.find(r"\multirow")
+        index_mc = content.find(r"\multicolumn")
+        if index_mr == -1 and index_mc == -1:
+            return {"colspan": colspan.strip(), "alignment": alignment.strip(), "content": content.strip()}, newpos
+        else:
+            indices = [i for i in (index_mr, index_mc) if i != -1]
+            first_index = min(indices)
+            inner, _ = parse_command_merge(content, first_index)
+            merged = {"colspan": colspan.strip(), "alignment": alignment.strip()}
+            merged.update(inner)
+            return merged, newpos
+
+    # Not a recognized command starting at pos.
+    return None, pos
+
+def extract_merged_commands(s):
+    """
+    Scan through the LaTeX string s and extract merged multirow/multicolumn commands.
+    For each command found, if there is nesting the parser merges the outer and inner
+    parameters so that the final result includes both the rowspan (or width) and the colspan
+    (or alignment) along with the innermost content.
+
+    Returns a list of dictionaries.
+    """
+    pos = 0
+    results = []
+    while pos < len(s):
+        if s[pos] == '\\':
+            res, newpos = parse_command_merge(s, pos)
+            if res is not None:
+                results.append(res)
+                pos = newpos
+                continue
+        pos += 1
+    return results
+
+def remove_tags(html, tags_to_remove):
+    soup = BeautifulSoup(html, "html.parser")
+    # Loop through the tags to remove
+    for tag_name in tags_to_remove:
+        for tag in soup.find_all(tag_name):
+            # Move the children of the tag to the parent tag
+            tag.unwrap()  # This removes the tag but keeps its contents
+    # Return the modified HTML as a string
+    return str(soup)
+
+def convert_th_to_td(html):
+    """Replace all th tags with td tags
+    """
+    soup = BeautifulSoup(html)
+    for th_tag in soup.find_all('th'):
+        th_tag.name = 'td'
+    return str(soup)
+
+def replace_italic(text):
+    pattern = re.compile(r'(?<!\\)_(.*?)(?<!\\)_')
+
+    def italic_replacer(match):
+        # Get the text inside the underscores.
+        content = match.group(1)
+        # Remove the escape (backslash) from any escaped underscores inside.
+        content = content.replace(r'\_', '_')
+        return f"<i>{content}</i>"
+
+    # Replace all occurrences of the pattern using the replacer function.
+    return pattern.sub(italic_replacer, text)
+
+
+def replace_bold(text):
+    pattern = re.compile(r'(?<!\\)\*\*(.*?)(?<!\\)\*\*')
+
+    def bold_replacer(match):
+        content = match.group(1)
+        # Unescape any escaped asterisks within the captured text.
+        content = content.replace(r'\*', '*')
+        return f"<b>{content}</b>"
+
+    return pattern.sub(bold_replacer, text)
+
+def latex_table_to_html(latex_str, add_head_body = False):
+    # Pattern to match the entire tabular environment
+    table_pattern = r'\\begin{tabular}{([^}]*)}\s*(.*?)\\end{tabular}'
+
+    def process_cell(cell):
+        # Clean up cell content
+        cell = cell.strip()
+
+        out = extract_merged_commands(cell)
+        if len(out) > 0:
+            cell = process_cell(out[0]["content"])["content"]
+            rowspan = int(out[0].get("rowspan", "1"))
+            colspan = int(out[0].get("colspan", "1"))
+            return {
+                "content": cell,
+                "colspan": colspan,
+                "rowspan": rowspan
+            }
+
+        # Replace latex and markdown formatting with HTML tags
+        cell = re.sub(r'\$([^$]*)\$', r'\1', cell)  # Remove math mode
+        cell = re.sub(r'\\textbf{([^}]*)}', r'<b>\1</b>', cell)  # Convert latex bold
+        cell = re.sub(r'\\textit{([^}]*)}', r'<i>\1</i>', cell)  # Convert latex italic
+        cell = replace_italic(cell)
+        cell = replace_bold(cell)
+        cell = cell.replace("\\$", "$").replace("\\%", "%").replace("\\newline", "\n").replace("\\textless", "<").replace("\\textgreater", ">").replace("\\*", "*").replace("\\_", "_").replace("\\backslash", "\\")
+
+        # Replace \& with & in the cell text
+        cell = cell.replace(r'\&', '&')
+        cell = cell.replace('<tbc>', '')
+        # Preserve newlines for downstream row-splitting; clean other tokens
+        cell = cell.replace('\\unknown', '').replace('\\<|unk|\\>', '').replace('<u>', '<underline>').replace('</u>', '</underline>')
+        return {
+            'content': cell,
+            'colspan': 1,
+            'rowspan': 1
+        }
+
+    def split_row(input_string):
+        # Use a regular expression to split on '&' that is not preceded by a backslash
+        return re.split(r'(?<!\\)&', input_string)
+
+    def convert_table(match):
+        # Extract table content
+        format_spec, content = match.groups()
+
+        # Start building HTML table
+        html = ['<table>']
+
+        # Track cells for multirow
+        multirow_tracker = set()
+
+        # Process rows
+        rows = re.split(r'\\\\', content)
+        current_row = 0
+        
+        for row in rows:
+            if not row.strip():
+                continue
+
+            row = row.strip()
+
+            # Skip \hline
+            if '\\hline' in row:
+                row = row.replace('\\hline', '')
+                if not row.strip():
+                    continue
+
+            row = clean_multi_cells(row)
+
+            # Process cells
+            cells = split_row(row)
+            processed_cells = [process_cell(cell) for cell in cells]
+
+            # Build per-cell line lists splitting on newline or <br> tokens
+            def split_lines(text):
+                parts = re.split(r'(?:\n|<br\s*/?>)+', text)
+                return parts if parts is not None else ['']
+
+            line_lists = [split_lines(cell['content']) for cell in processed_cells]
+            max_lines = max(len(lst) for lst in line_lists) if line_lists else 1
+
+            # Emit one or more rows based on max_lines
+            for line_idx in range(max_lines):
+                if add_head_body:
+                    if current_row == 0:
+                        html.append(' <thead>')
+                    if current_row == 1:
+                        html.append(' <tbody>')
+                html.append('  <tr>')
+                current_col = 0
+
+                for col_idx, cell in enumerate(processed_cells):
+                    content_segment = line_lists[col_idx][line_idx] if line_idx < len(line_lists[col_idx]) else ''
+
+                    attrs = []
+                    if cell['colspan'] > 1:
+                        attrs.append(f'colspan="{cell["colspan"]}"')
+                    # Only apply original rowspan to the first emitted line
+                    if cell['rowspan'] > 1 and line_idx == 0:
+                        attrs.append(f'rowspan="{cell["rowspan"]}"')
+                        for r in range(current_row + 1, current_row + cell['rowspan']):
+                            for c in range(current_col, current_col + cell['colspan']):
+                                multirow_tracker.add((r, c))
+
+                    # If this position is covered by a prior rowspan, skip rendering a duplicate cell
+                    if cell['rowspan'] > 1 and line_idx > 0:
+                        current_col += cell['colspan']
+                        continue
+
+                    if (current_row, current_col) in multirow_tracker and content_segment == '' and cell["colspan"] == 1 and cell["rowspan"] == 1:
+                        current_col += cell['colspan']
+                        continue
+
+                    attr_str = ' ' + ' '.join(attrs) if attrs else ''
+                    cell_tag = 'td'
+                    html.append(f'    <{cell_tag}{attr_str}>{content_segment}</{cell_tag}>')
+                    current_col += cell['colspan']
+
+                if add_head_body and current_row == 0:
+                    html.append(' </thead>')
+                html.append('  </tr>')
+                current_row += 1
+        if add_head_body:
+            html.append(' </tbody>')
+        html.append('</table>')
+        return '\n'.join(html)
+
+    # Convert all tabular environments in the input
+    return re.sub(table_pattern, convert_table, latex_str, flags=re.DOTALL)
+def convert_single_table(table):
+    """
+    Convert a single HTML table to Markdown format.
+    
+    Args:
+        table: BeautifulSoup table element
+        
+    Returns:
+        str: Markdown table string
+    """
+    markdown_lines = []
+    rows = table.find_all('tr')
+    
+    for i, row in enumerate(rows):
+        cells = row.find_all(['td', 'th'])
+        if not cells:
+            continue
+        
+        # Convert cells to text, handling nested elements
+        row_data = []
+        for cell in cells:
+            # Get text content, handling nested elements
+            cell_text = cell.get_text(separator=' ', strip=True)
+            # Escape pipe characters
+            cell_text = cell_text.replace('|', '\\|')
+            row_data.append(cell_text)
+        
+        # Add row to markdown
+        markdown_lines.append('| ' + ' | '.join(row_data) + ' |')
+        
+        # Add separator after header row
+        if i == 0:
+            separator = '| ' + ' | '.join(['---'] * len(cells)) + ' |'
+            markdown_lines.append(separator)
+    
+    return '\n'.join(markdown_lines)
+def convert_html_tables_to_markdown(html_content):
+    """
+    Find all HTML tables and convert them to Markdown while preserving all other content.
+    
+    Args:
+        html_content (str): HTML content that may contain tables
+        
+    Returns:
+        str: HTML content with tables converted to Markdown
+    """
+    soup = BeautifulSoup(html_content, 'html.parser')
+    
+    # Find all tables
+    tables = soup.find_all('table')
+    
+    if not tables:
+        return html_content  # Return original content unchanged
+    
+    # Convert each table to markdown and replace it
+    for table in tables:
+        markdown_table = convert_single_table(table)
+        
+        # Create a new element to replace the table
+        replacement = soup.new_string('\n' + markdown_table + '\n')
+        table.replace_with(replacement)
+    
+    return str(soup)

logitsprocs/nemotron_parse_vllm_logitprocs.py

@@ -0,0 +1,347 @@
+"""
+Custom vLLM v1 logits processors for NVIDIA Nemotron Parse.
+
+Tuning via environment variables:
+  - NEMOTRON_PARSE_TABLE_PREFIX (default: \\begin{tabular})
+  - NEMOTRON_PARSE_REP_MAX      (default: 10)
+  - NEMOTRON_PARSE_REP_WINDOW   (default: 2 * max_ngram * (rep_max + 1))
+  - NEMOTRON_PARSE_REP_NGRAMS   (default: 1,2,3,4,5)
+"""
+
+from __future__ import annotations
+
+import os
+import warnings
+from dataclasses import dataclass
+from typing import Dict, List, Set
+
+import torch
+from transformers import AutoTokenizer
+
+from vllm.v1.sample.logits_processor.builtin import process_dict_updates
+from vllm.v1.sample.logits_processor.interface import BatchUpdate, LogitsProcessor
+
+
+def _strip_trailing_negative_token_ids(token_ids: List[int]) -> List[int]:
+    """vLLM v1 keeps a trailing -1 placeholder in output_token_ids."""
+    i = len(token_ids)
+    while i > 0 and token_ids[i - 1] < 0:
+        i -= 1
+    return token_ids[:i]
+
+
+def _env_int(name: str, default: int) -> int:
+    val = os.environ.get(name, "").strip()
+    if not val:
+        return default
+    try:
+        return int(val)
+    except ValueError:
+        return default
+
+
+def _env_csv_ints(name: str, default: List[int]) -> List[int]:
+    val = os.environ.get(name, "").strip()
+    if not val:
+        return default
+    out: List[int] = []
+    for part in val.split(","):
+        part = part.strip()
+        if not part:
+            continue
+        try:
+            out.append(int(part))
+        except ValueError:
+            return default
+    return out or default
+
+
+def _warn_if_unexpected_arch(vllm_config) -> None:
+    try:
+        arches = getattr(vllm_config.model_config, "architectures", None) or []
+    except Exception:
+        arches = []
+    if any(a == "NemotronParseForConditionalGeneration" for a in arches):
+        return
+    warnings.warn(
+        "Nemotron-Parse logits processors enabled for a model whose "
+        f"architectures={arches!r}. These processors assume Nemotron-Parse-style "
+        "special tokens like <x_...>, <y_...>, and <class_...>; behavior may be "
+        "unsupported for other models.",
+        RuntimeWarning,
+        stacklevel=2,
+    )
+
+
+def _load_hf_tokenizer(vllm_config):
+    model_cfg = vllm_config.model_config
+    tok_name = getattr(model_cfg, "tokenizer", None) or getattr(model_cfg, "model", None)
+    trust_rc = bool(getattr(model_cfg, "trust_remote_code", False))
+    revision = getattr(model_cfg, "tokenizer_revision", None) or getattr(model_cfg, "revision", None)
+    token = getattr(model_cfg, "hf_token", None)
+    kwargs = {"trust_remote_code": trust_rc}
+    if revision:
+        kwargs["revision"] = revision
+    if token:
+        # HF uses `token=` in modern versions.
+        kwargs["token"] = token
+    return AutoTokenizer.from_pretrained(tok_name, **kwargs)
+
+
+def _build_token_sets(tokenizer) -> tuple[Set[int], Set[int], Set[int]]:
+    """Build token-id sets for <x_...>, <y_...>, and <class_...>."""
+
+    def _scan(items) -> tuple[Set[int], Set[int], Set[int]]:
+        x_ids: Set[int] = set()
+        y_ids: Set[int] = set()
+        class_ids: Set[int] = set()
+        for token, tid in items:
+            if token.startswith("<x_") and token.endswith(">") and token.count(">") == 1:
+                x_ids.add(tid)
+            elif token.startswith("<y_") and token.endswith(">") and token.count(">") == 1:
+                y_ids.add(tid)
+            elif token.startswith("<class_") and token.endswith(">") and token.count(">") == 1:
+                class_ids.add(tid)
+        return x_ids, y_ids, class_ids
+
+    added_vocab = {}
+    if hasattr(tokenizer, "get_added_vocab"):
+        try:
+            added_vocab = tokenizer.get_added_vocab() or {}
+        except Exception:
+            added_vocab = {}
+
+    x_ids, y_ids, class_ids = _scan(added_vocab.items())
+
+    if len(x_ids) < 100 or len(y_ids) < 100 or not class_ids:
+        vocab = tokenizer.get_vocab()
+        x2, y2, c2 = _scan(vocab.items())
+        x_ids, y_ids, class_ids = x2, y2, c2
+
+    return x_ids, y_ids, class_ids
+
+
+@dataclass
+class _TableReqState:
+    output_ids: List[int]
+    insertion_active: bool = False
+    insertion_pos: int = 0
+    expecting_end_coords: bool = False
+    _last_class_pos: int = -1
+
+
+class NemotronParseTableInsertionLogitsProcessor(LogitsProcessor):
+    """Force a table prefix right after each START <x_...><y_...> coordinate pair."""
+
+    def __init__(self, vllm_config, device: torch.device, is_pin_memory: bool) -> None:
+        self.enabled = True
+        self.req_states: Dict[int, _TableReqState] = {}
+
+        self.table_prefix = os.environ.get("NEMOTRON_PARSE_TABLE_PREFIX", r"\begin{tabular}")
+        self.table_prefix_ids: List[int] = []
+        self.x_token_ids: Set[int] = set()
+        self.y_token_ids: Set[int] = set()
+        self.class_token_ids: Set[int] = set()
+
+        if self.enabled:
+            _warn_if_unexpected_arch(vllm_config)
+            tok = _load_hf_tokenizer(vllm_config)
+            self.table_prefix_ids = tok.encode(self.table_prefix, add_special_tokens=False)
+            self.x_token_ids, self.y_token_ids, self.class_token_ids = _build_token_sets(tok)
+
+        if not self.x_token_ids or not self.y_token_ids or not self.table_prefix_ids:
+            self.enabled = False
+
+    @classmethod
+    def validate_params(cls, sampling_params):
+        return None
+
+    def is_argmax_invariant(self) -> bool:
+        return False
+
+    def update_state(self, batch_update: BatchUpdate | None) -> None:
+        if not self.enabled:
+            return
+
+        def _new_state(params, prompt_ids, output_ids):
+            return _TableReqState(output_ids=output_ids)
+
+        process_dict_updates(self.req_states, batch_update, _new_state)
+
+        for st in self.req_states.values():
+            out = st.output_ids
+            if not out:
+                continue
+
+            view = _strip_trailing_negative_token_ids(out)
+            if not view:
+                continue
+
+            # If we ended an object (a newly-generated <class_...>), next xy should be START.
+            if self.class_token_ids:
+                last_class_pos = -1
+                for i in range(len(view) - 1, -1, -1):
+                    if view[i] in self.class_token_ids:
+                        last_class_pos = i
+                        break
+                if last_class_pos != -1 and last_class_pos != st._last_class_pos:
+                    st._last_class_pos = last_class_pos
+                    st.expecting_end_coords = False
+
+            if st.insertion_active:
+                continue
+
+            if len(view) >= 2 and view[-2] in self.x_token_ids and view[-1] in self.y_token_ids:
+                if not st.expecting_end_coords:
+                    st.insertion_active = True
+                    st.insertion_pos = 0
+                    st.expecting_end_coords = True
+                else:
+                    pass
+
+    def apply(self, logits: torch.Tensor) -> torch.Tensor:
+        if not self.enabled or not self.req_states:
+            return logits
+
+        for req_idx, st in self.req_states.items():
+            if not st.insertion_active:
+                continue
+            pos = st.insertion_pos
+            if pos >= len(self.table_prefix_ids):
+                st.insertion_active = False
+                continue
+
+            forced_tid = self.table_prefix_ids[pos]
+            logits[req_idx].fill_(-float("inf"))
+            logits[req_idx, forced_tid] = 0.0
+
+            st.insertion_pos = pos + 1
+            if st.insertion_pos >= len(self.table_prefix_ids):
+                st.insertion_active = False
+
+        return logits
+
+
+@dataclass
+class _RepReqState:
+    output_ids: List[int]
+    in_cooldown: bool = False
+    segment_start: int = 0
+    _last_class_pos: int = -1
+
+
+class NemotronParseRepetitionStopLogitsProcessor(LogitsProcessor):
+    """Force an <x_...> token when consecutive repetition exceeds a threshold."""
+
+    def __init__(self, vllm_config, device: torch.device, is_pin_memory: bool) -> None:
+        self.enabled = True
+        self.req_states: Dict[int, _RepReqState] = {}
+
+        self.max_repetitions = _env_int("NEMOTRON_PARSE_REP_MAX", 10)
+        self.ngram_sizes = _env_csv_ints("NEMOTRON_PARSE_REP_NGRAMS", [1, 2, 3, 4, 5])
+        max_ngram = max((n for n in self.ngram_sizes if n > 0), default=1)
+        default_window = 2 * max_ngram * (self.max_repetitions + 1)
+        self.window_size = _env_int("NEMOTRON_PARSE_REP_WINDOW", default_window)
+
+        self.x_token_ids: Set[int] = set()
+        self.class_token_ids: Set[int] = set()
+
+        if self.enabled:
+            _warn_if_unexpected_arch(vllm_config)
+            tok = _load_hf_tokenizer(vllm_config)
+            x_ids, _, class_ids = _build_token_sets(tok)
+            self.x_token_ids = x_ids
+            self.class_token_ids = class_ids
+
+        if not self.x_token_ids:
+            self.enabled = False
+
+    @classmethod
+    def validate_params(cls, sampling_params):
+        return None
+
+    def is_argmax_invariant(self) -> bool:
+        return False
+
+    def update_state(self, batch_update: BatchUpdate | None) -> None:
+        if not self.enabled:
+            return
+
+        def _new_state(params, prompt_ids, output_ids):
+            return _RepReqState(output_ids=output_ids)
+
+        process_dict_updates(self.req_states, batch_update, _new_state)
+
+        for st in self.req_states.values():
+            out = st.output_ids
+            if not out:
+                continue
+            view = _strip_trailing_negative_token_ids(out)
+            if not view:
+                continue
+            # Clear cooldown and advance segment on a newly-generated <class_...> token.
+            if self.class_token_ids:
+                last_class_pos = -1
+                for i in range(len(view) - 1, -1, -1):
+                    if view[i] in self.class_token_ids:
+                        last_class_pos = i
+                        break
+                if last_class_pos != -1 and last_class_pos != st._last_class_pos:
+                    st._last_class_pos = last_class_pos
+                    st.in_cooldown = False
+                    st.segment_start = last_class_pos + 1
+
+    @staticmethod
+    def _max_consecutive_repetitions(seq: List[int], n: int) -> int:
+        if len(seq) < n:
+            return 0
+        max_consec = 1
+        cur = 1
+        prev = tuple(seq[0:n])
+        i = n
+        while i <= len(seq) - n:
+            cur_ng = tuple(seq[i:i + n])
+            if cur_ng == prev:
+                cur += 1
+                if cur > max_consec:
+                    max_consec = cur
+                i += n
+            else:
+                cur = 1
+                prev = cur_ng
+                i += 1
+        return max_consec
+
+    def _has_excessive_repetition(self, seq: List[int]) -> bool:
+        if not seq:
+            return False
+        check_seq = seq[-self.window_size:] if len(seq) > self.window_size else seq
+        for n in self.ngram_sizes:
+            if n <= 0:
+                continue
+            if self._max_consecutive_repetitions(check_seq, n) > self.max_repetitions:
+                return True
+        return False
+
+    def apply(self, logits: torch.Tensor) -> torch.Tensor:
+        if not self.enabled or not self.req_states:
+            return logits
+
+        for req_idx, st in self.req_states.items():
+            if st.in_cooldown:
+                continue
+            view = _strip_trailing_negative_token_ids(st.output_ids)
+            seg = view[st.segment_start:]
+            if not self._has_excessive_repetition(seg):
+                continue
+
+            st.in_cooldown = True
+            row = logits[req_idx]
+            original = row.clone()
+            row.fill_(-float("inf"))
+            for tid in self.x_token_ids:
+                row[tid] = original[tid]
+
+        return logits
+
+

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9d77822f43a04504619e4c3527e1371568b6e33dba94b7c3a0c0cf509a200cd4
+size 3745188184

postprocessing.py

@@ -0,0 +1,95 @@
+import re
+from latex2html import convert_html_tables_to_markdown, latex_table_to_html
+
+def extract_classes_bboxes(text: str):
+    _re_extract_class_bbox = re.compile(r'<x_(\d+(?:\.\d+)?)><y_(\d+(?:\.\d+)?)>(.*?)<x_(\d+(?:\.\d+)?)><y_(\d+(?:\.\d+)?)><class_([^>]+)>', re.DOTALL)
+    classes = []
+    bboxes = []
+    texts = []
+    for m in _re_extract_class_bbox.finditer(text):
+        x1, y1, text, x2, y2, cls = m.groups()
+        classes.append(cls)
+        bboxes.append((float(x1), float(y1), float(x2), float(y2)))
+        texts.append(text)
+
+    # TODO: Remove when fixed
+    classes = [
+        "Formula" if cls == "Inline-formula" else cls for cls in classes
+    ]
+    assert "Page-number" not in classes
+
+    return classes, bboxes, texts
+
+def transform_bbox_to_original(bbox, original_width, original_height, target_w=1664, target_h=2048):
+    # Replicate exact resize logic
+    aspect_ratio = original_width / original_height
+    new_height = original_height
+    new_width = original_width
+    
+    if original_height > target_h:
+        new_height = target_h
+        new_width = int(new_height * aspect_ratio)
+   
+    if new_width > target_w:
+        new_width = target_w
+        new_height = int(new_width / aspect_ratio)
+    
+    resized_width = new_width
+    resized_height = new_height
+    
+    # Calculate padding
+    pad_left = (target_w - resized_width) // 2
+    pad_top = (target_h - resized_height) // 2
+    
+#     # Transform: use the ACTUAL resized dimensions, not the scale
+#     # X coords
+    left = ((bbox[0] * target_w) - pad_left) * original_width / resized_width
+    right = ((bbox[2] * target_w) - pad_left) * original_width / resized_width
+    
+#     # Y coords - using original_height / resized_height directly
+    top = ((bbox[1] * target_h) - pad_top) * original_height / resized_height
+    bottom = ((bbox[3] * target_h) - pad_top) * original_height / resized_height
+    
+    return left, top, right, bottom
+
+def postprocess_text(text, cls = 'Text', text_format='markdown', table_format='latex', blank_text_in_figures=False):
+    assert text_format in ['markdown', 'plain'], 'Unknown text format. Supported: markdown | plain'
+    assert table_format in ['latex', 'HTML', 'markdown'], 'Unknown table format. Supported: latex | HTML | markdown'
+    if cls != 'Table':
+        if text_format == 'plain':
+            text = convert_mmd_to_plain_text_ours(text)
+    elif table_format == 'HTML':
+        text = latex_table_to_html(text)
+    elif table_format == 'markdown':
+        text = convert_html_tables_to_markdown(latex_table_to_html(text))
+    if blank_text_in_figures and cls == 'Picture':
+        text = ''
+    return text
+
+def remove_nemotron_formatting(text):
+   text = text.replace('<tbc>', '')
+   text = text.replace('\\<|unk|\\>', '')
+   text = text.replace('\\unknown', '')
+   return text
+def convert_mmd_to_plain_text_ours(mmd_text):
+    mmd_text = re.sub(r'<sup>(.*?)</sup>', r'^{\\1}', mmd_text, flags=re.DOTALL)
+    mmd_text = re.sub(r'<sub>(.*?)</sub>', r'_{\\1}', mmd_text, flags=re.DOTALL)
+    mmd_text = mmd_text.replace('<br>', '\n')
+
+    # Remove headers (e.g., ##)
+    mmd_text = re.sub(r'#+\s', '', mmd_text)
+    
+    # Remove bold (e.g., **)
+    mmd_text = re.sub(r'\*\*(.*?)\*\*', r'\1', mmd_text)
+    #mmd_text = mmd_text.replace("**","")
+    # Remove italic (e.g., *)
+    mmd_text = re.sub(r'\*(.*?)\*', r'\1', mmd_text)
+    # Remove emphasized text formatting (e.g., _)
+    mmd_text = re.sub(r'(?<!\w)_([^_]+)_', r'\1', mmd_text)
+    
+    # Remove formulas inside paragraphs (e.g., \(R_{ij}(P^{a})=0\))
+    #mmd_text = re.sub(r'\\\((.*?)\\\)', '', mmd_text)
+    
+    # Remove asterisk in lists
+    #mmd_text = re.sub(r'^\*\s', '', mmd_text, flags=re.MULTILINE)    
+    return mmd_text.strip()

preprocessor_config.json

@@ -0,0 +1,24 @@
+{
+  "feature_extractor_type": "NemotronParseImageProcessor",
+  "image_processor_type": "NemotronParseImageProcessor",
+  "processor_class": "NemotronParseProcessor",
+  "auto_map": {
+    "AutoImageProcessor": "hf_nemotron_parse_processor.NemotronParseImageProcessor",
+    "AutoProcessor": "hf_nemotron_parse_processor.NemotronParseProcessor"
+  },
+  "do_normalize": false,
+  "do_rescale": true,
+  "rescale_factor": 0.00392156862745098,
+  "size": {
+    "height": 2048,
+    "width": 1664,
+    "longest_edge": [
+      2048,
+      1664
+    ]
+  },
+  "final_size": [
+    2048,
+    1664
+  ]
+}

privacy.md

@@ -0,0 +1,10 @@
+Field | Response |
+| :---- | :---- |
+| Generatable or reverse engineerable personal data? | No |
+| Personal data used to create this model? | No |
+| How often is the dataset reviewed? | Before Release |
+| Was data from user interactions with the AI model (e.g. user input and prompts) used to train the model?| No|
+| Is there provenance for all datasets used in training? | Yes |
+| Does data labeling (annotation, metadata) comply with privacy laws? | Yes |
+| Is data compliant with data subject requests for data correction or removal, if such a request was made? | No, not possible with externally-sourced data. |
+| Applicable Privacy Policy | [NVIDIA Privacy Policy](https://www.nvidia.com/en-us/about-nvidia/privacy-policy/) |

processor_config.json

@@ -0,0 +1,8 @@
+{
+  "processor_class": "NemotronParseProcessor",
+  "image_processor_type": "NemotronParseImageProcessor",
+  "auto_map": {
+    "AutoProcessor": "hf_nemotron_parse_processor.NemotronParseProcessor",
+    "AutoImageProcessor": "hf_nemotron_parse_processor.NemotronParseImageProcessor"
+  }
+}

safety.md

@@ -0,0 +1,8 @@
+| Field | Response |
+| :---- | :---- |
+| Model Application Field(s): | Chat, Instruction Following, Chatbot Development, Code Generation, Reasoning, Customer Service |
+| Describe the life critical impact (if present). | Not Applicable |
+| Use Case Restrictions: | GOVERNING TERMS: [Product-Specific Terms for NVIDIA AI Products](https://www.nvidia.com/en-us/agreements/enterprise-software/product-specific-terms-for-ai-products/). Use of this model is governed by the [NVIDIA Community Model License](https://www.nvidia.com/en-us/agreements/enterprise-software/nvidia-community-models-license/). Use of the tokenizer included in this model is governed by the [CC-BY-4.0 license](https://creativecommons.org/licenses/by/4.0/)|
+| 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. |
+
+**You are responsible for ensuring that your use of NVIDIA provided models complies with all applicable laws.**

special_tokens_map.json

@@ -0,0 +1,39 @@
+{
+  "additional_special_tokens": [
+    {
+      "content": "<predict_no_text_in_pic>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    }
+  ],
+  "bos_token": {
+    "content": "<s>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "eos_token": {
+    "content": "</s>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "pad_token": {
+    "content": "<pad>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "unk_token": {
+    "content": "<unk>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  }
+}

vllm_example.py

@@ -0,0 +1,198 @@
+import argparse
+import base64
+import os
+from concurrent.futures import ThreadPoolExecutor, as_completed
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any, Dict, Iterable, List, Optional, Tuple
+
+from openai import OpenAI
+
+
+def _guess_mime(path: str) -> str:
+    ext = Path(path).suffix.lower().lstrip(".")
+    if ext in ("jpg", "jpeg"):
+        return "image/jpeg"
+    if ext in ("webp",):
+        return "image/webp"
+    # default
+    return "image/png"
+
+
+def _b64_image_data_url(path: str) -> str:
+    with open(path, "rb") as f:
+        img_b64 = base64.b64encode(f.read()).decode("utf-8")
+    mime = _guess_mime(path)
+    return f"data:{mime};base64,{img_b64}"
+
+
+def _iter_images(paths: List[str], image_dir: Optional[str]) -> List[str]:
+    out: List[str] = []
+    for p in paths:
+        out.append(p)
+    if image_dir:
+        for ext in ("*.png", "*.jpg", "*.jpeg", "*.webp"):
+            out.extend([str(x) for x in sorted(Path(image_dir).glob(ext))])
+    # De-dupe, keep order
+    seen = set()
+    deduped: List[str] = []
+    for p in out:
+        if p in seen:
+            continue
+        seen.add(p)
+        deduped.append(p)
+    return deduped
+
+
+@dataclass(frozen=True)
+class _ReqSpec:
+    image_path: str
+    request_idx: int
+
+
+def _make_client(base_url: str) -> OpenAI:
+    # openai>=1.x requires an API key; vLLM ignores it by default.
+    api_key = os.environ.get("OPENAI_API_KEY", "EMPTY")
+    return OpenAI(base_url=base_url, api_key=api_key)
+
+
+def _run_one(
+    req: _ReqSpec,
+    *,
+    base_url: str,
+    model: str,
+    prompt_text: str,
+    max_tokens: int,
+    temperature: float,
+    extra_body: Dict[str, Any],
+) -> Tuple[_ReqSpec, str]:
+    client = _make_client(base_url)
+    img_url = _b64_image_data_url(req.image_path)
+    resp = client.chat.completions.create(
+        model=model,
+        messages=[
+            {
+                "role": "user",
+                "content": [
+                    {"type": "text", "text": prompt_text},
+                    {"type": "image_url", "image_url": {"url": img_url}},
+                ],
+            }
+        ],
+        max_tokens=max_tokens,
+        temperature=temperature,
+        extra_body=extra_body,
+    )
+    text = resp.choices[0].message.content or ""
+    return req, text
+
+
+def _maybe_annotate(image_path: str, generated_text: str, out_image_path: str) -> None:
+    # Optional visualization (similar to example_with_table_processor.py).
+    from PIL import Image, ImageDraw  # local import so batching can run without pillow
+
+    from postprocessing import extract_classes_bboxes, postprocess_text, transform_bbox_to_original
+
+    image = Image.open(image_path).convert("RGB")
+
+    classes, bboxes, texts = extract_classes_bboxes(generated_text)
+    bboxes = [transform_bbox_to_original(bbox, image.width, image.height) for bbox in bboxes]
+
+    table_format = "HTML"  # latex | HTML | markdown
+    text_format = "markdown"  # markdown | plain
+    blank_text_in_figures = False
+
+    _ = [
+        postprocess_text(
+            text,
+            cls=cls,
+            table_format=table_format,
+            text_format=text_format,
+            blank_text_in_figures=blank_text_in_figures,
+        )
+        for text, cls in zip(texts, classes)
+    ]
+
+    draw = ImageDraw.Draw(image)
+    for bbox in bboxes:
+        draw.rectangle(
+            (bbox[0], bbox[1], max(bbox[0], bbox[2]), max(bbox[1], bbox[3])),
+            outline="red",
+            width=2,
+        )
+
+    image.save(out_image_path)
+
+
+def main() -> None:
+    ap = argparse.ArgumentParser(description="vLLM OpenAI-compatible example (batch + .txt outputs).")
+    ap.add_argument("--base-url", default="http://localhost:8000/v1")
+    ap.add_argument("--model", default="nvidia/NVIDIA-Nemotron-Parse-v1.2")
+    ap.add_argument("--image", action="append", default=[], help="Image path (repeatable).")
+    ap.add_argument("--image-dir", default=None, help="Directory of images to run (png/jpg/jpeg/webp).")
+    ap.add_argument("--out-dir", default="vllm_outputs", help="Where to write .txt outputs.")
+    ap.add_argument("--concurrency", type=int, default=4, help="How many concurrent requests to send.")
+    ap.add_argument("--max-tokens", type=int, default=8994)
+    ap.add_argument("--temperature", type=float, default=0.0)
+    ap.add_argument(
+        "--annotate",
+        action=argparse.BooleanOptionalAction,
+        default=True,
+        help="Write annotated images with boxes to --out-dir (default: enabled). Use --no-annotate to disable.",
+    )
+
+    args = ap.parse_args()
+
+    image_paths = _iter_images(args.image, args.image_dir)
+    if not image_paths:
+        raise SystemExit("No images provided. Use --image PATH (repeatable) or --image-dir DIR.")
+
+    out_dir = Path(args.out_dir)
+    out_dir.mkdir(parents=True, exist_ok=True)
+
+    prompt_text = "</s><s><predict_bbox><predict_classes><output_markdown><predict_no_text_in_pic>"
+    #prompt_text = "</s><s><predict_bbox><predict_classes><output_markdown><predict_text_in_pic>"
+
+    extra_body = {
+        "repetition_penalty": 1.1,
+        "top_k": 1,
+        "skip_special_tokens": False,
+    }
+
+    reqs: List[_ReqSpec] = []
+    for idx, img in enumerate(image_paths):
+        reqs.append(_ReqSpec(image_path=img, request_idx=idx))
+
+    # Concurrency is the simplest way to make sure vLLM batches requests internally.
+    summary_lines: List[str] = []
+    with ThreadPoolExecutor(max_workers=max(1, args.concurrency)) as ex:
+        futs = [
+            ex.submit(
+                _run_one,
+                r,
+                base_url=args.base_url,
+                model=args.model,
+                prompt_text=prompt_text,
+                max_tokens=args.max_tokens,
+                temperature=args.temperature,
+                extra_body=extra_body,
+            )
+            for r in reqs
+        ]
+        for fut in as_completed(futs):
+            req, text = fut.result()
+            base = Path(req.image_path).name
+            stem = f"{req.request_idx:04d}_{base}"
+            out_txt = out_dir / f"{stem}.txt"
+            out_txt.write_text(text, encoding="utf-8")
+            summary_lines.append(f"{req.image_path}\t{out_txt}")
+
+            if args.annotate:
+                out_img = out_dir / f"{stem}.annotated.jpg"
+                _maybe_annotate(req.image_path, text, str(out_img))
+
+    (out_dir / "summary.txt").write_text("\n".join(sorted(summary_lines)) + "\n", encoding="utf-8")
+
+
+if __name__ == "__main__":
+    main()