Upload folder using huggingface_hub

README.md

@@ -34,7 +34,7 @@ Transformer-based vision-encoder-decoder model
 
 ### Network Architecture 
 * 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 (13184 tokens to 3201 tokens)<br>
+* Adapter Layer: 1D convolutions & norms to compress dimensionality and sequence length of the latent space (13184 tokens to 833 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>
@@ -51,7 +51,7 @@ Carbon Emissions: 3.21 tCO2e <br>
 * Input Type(s): Red, Green, Blue (RGB) + Prompt (String)
 * Input Parameters: 2D, 1D
 - Other Properties Related to Input:
-  - Max Input Resolution (Width, Height): 1648, 2048
+  - Max Input Resolution (Width, Height): 1668, 2048
   - Min Input Resolution (Width, Height): 1024, 1280
 - Channel Count: 3
 
@@ -78,6 +78,7 @@ The integration of foundation and fine-tuned models into AI systems requires add
 ## Model Version:
 
 V1.1-Light
+This version preserves reading order of Tables, Captions, Pictures, and other elements as well as offers 20% speed improvement compared to Nemotron-Parse-1.1
 
 ## Quick Start
 
@@ -95,15 +96,6 @@ 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
 
-from hf_nemotron_parse_config import NemotronParseConfig
-from hf_nemotron_parse_modeling import NemotronParseForConditionalGeneration
-from hf_nemotron_parse_processor import NemotronParseProcessor, NemotronParseImageProcessor
-
-AutoConfig.register("nemotron_parse", NemotronParseConfig)
-AutoModel.register(NemotronParseConfig, NemotronParseForConditionalGeneration)
-AutoProcessor.register("nemotron_parse", NemotronParseProcessor)
-AutoImageProcessor.register("nemotron_parse", NemotronParseImageProcessor)
-
 # Load model and processor
 model_path = "nvidia/NVIDIA-Nemotron-Parse-v1.1-Light"  # Or use a local path
 device = "cuda:0"
@@ -131,6 +123,15 @@ 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
+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
 classes, bboxes, texts = extract_classes_bboxes(generated_text)
 bboxes = [transform_bbox_to_original(bbox, image.width, image.height) for bbox in bboxes]
 
@@ -148,6 +149,7 @@ for bbox in bboxes:
   draw.rectangle((bbox[0], bbox[1], bbox[2], bbox[3]), outline="red")
 ```
 
+
 ## Training, Testing, and Evaluation Datasets:
 
 

example.py

@@ -0,0 +1,49 @@
+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.1-Light"  # 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>"
+
+# Process image
+inputs = processor(images=[image], text=task_prompt, return_tensors="pt").to(device)
+prompt_ids = processor.tokenizer.encode(task_prompt, return_tensors="pt", add_special_tokens=False).cuda()
+
+
+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]
+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)
+
+# OPTIONAL - Draw bounding boxes
+draw = ImageDraw.Draw(image)
+for bbox in bboxes:
+  draw.rectangle((bbox[0], bbox[1], bbox[2], bbox[3]), outline="red")
+

hf_nemotron_parse_modeling.py

@@ -13,7 +13,7 @@ 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 NemotronParseLightConfig
+from .hf_nemotron_parse_config import NemotronParseLightConfig
 from transformers import AutoModel
 import time
 from transformers.modeling_attn_mask_utils import (

hf_nemotron_parse_processor.py

@@ -252,7 +252,7 @@ class NemotronParseLightImageProcessor(BaseImageProcessor, ImageProcessingMixin)
 class NemotronParseLightProcessor(ProcessorMixin):
     
     attributes = ["image_processor", "tokenizer"]
-    image_processor_class = "NemotronParseLightImageProcessor"
+    image_processor_class = "AutoImageProcessor"
     tokenizer_class = ("PreTrainedTokenizer", "PreTrainedTokenizerFast")
     
     def __init__(self, image_processor=None, tokenizer=None, **kwargs):
@@ -350,8 +350,24 @@ class NemotronParseLightProcessor(ProcessorMixin):
         
         This method is compatible with AutoProcessor.from_pretrained().
         """
-        # Use the parent class's from_pretrained method which handles auto-loading
-        return super().from_pretrained(pretrained_model_name_or_path, **kwargs)
+        # 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):
         """

preprocessor_config.json

@@ -2,6 +2,10 @@
   "feature_extractor_type": "NemotronParseLightImageProcessor",
   "image_processor_type": "NemotronParseLightImageProcessor",
   "processor_class": "NemotronParseLightProcessor",
+  "auto_map": {
+	  "AutoImageProcessor": "hf_nemotron_parse_processor.NemotronParseLightImageProcessor",
+	  "AutoProcessor": "hf_nemotron_parse_processor.NemotronParseLightProcessor"
+  },
   "do_normalize": false,
   "do_rescale": true,
   "rescale_factor": 0.00392156862745098,