Add golden tests and latest runtime support for Nemotron Parse v1.1

#8
Dockerfile ADDED
@@ -0,0 +1,39 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ FROM nvcr.io/nvidia/pytorch:26.03-py3
2
+
3
+ # ---------------------------------------------------------------------------
4
+ # Install uv
5
+ # ---------------------------------------------------------------------------
6
+ COPY --from=ghcr.io/astral-sh/uv:latest /uv /uvx /bin/
7
+
8
+ WORKDIR /workspace
9
+
10
+ # ---------------------------------------------------------------------------
11
+ # Virtual environment
12
+ #
13
+ # --system-site-packages makes the venv inherit every package already
14
+ # installed in the NVIDIA base image, including torch, torchvision, and the
15
+ # matching CUDA libraries. uv will not reinstall those packages (they are
16
+ # declared in [tool.uv] exclude-dependencies), so the base-image builds are
17
+ # left untouched.
18
+ # ---------------------------------------------------------------------------
19
+ RUN uv venv /opt/venv --system-site-packages
20
+ ENV VIRTUAL_ENV=/opt/venv
21
+ ENV PATH="/opt/venv/bin:$PATH"
22
+
23
+ # ---------------------------------------------------------------------------
24
+ # Layer 1 - install dependencies (not the project itself)
25
+ #
26
+ # Only pyproject.toml is present at this point, so Docker re-runs this step
27
+ # only when the dependency list changes, not when source files change.
28
+ # ---------------------------------------------------------------------------
29
+ COPY pyproject.toml .
30
+ RUN uv sync --no-install-project
31
+
32
+ # ---------------------------------------------------------------------------
33
+ # Layer 2 - copy source and finish
34
+ #
35
+ # With package = false in pyproject.toml there is nothing extra to install;
36
+ # the second sync is a fast no-op that confirms the environment is complete.
37
+ # ---------------------------------------------------------------------------
38
+ COPY . .
39
+ RUN uv sync
README.md CHANGED
@@ -108,6 +108,8 @@ pip install transformers==4.51.3
108
  pip install timm==1.0.22
109
  ```
110
 
 
 
111
  ### Usage example
112
 
113
  ```python
@@ -169,6 +171,8 @@ for bbox in bboxes:
169
 
170
  **Update**: Nemotron-Parse-v1.1 is now [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).
171
 
 
 
172
  Note: when running on A100/A10 we recommend running vllm serve with _--attention-backend=TRITON_ATTN_
173
 
174
  You will need to install *albumentations* on top, and then follow the VLLM Inference example below:
@@ -376,4 +380,4 @@ Get access to knowledge base articles and support cases or [submit a ticket](htt
376
  primaryClass={cs.LG},
377
  url={https://arxiv.org/abs/2511.20478},
378
  }
379
- ```
 
108
  pip install timm==1.0.22
109
  ```
110
 
111
+ `transformers==4.51.3` remains the pinned reference environment for this model. The included remote code has also been tested against newer Transformers APIs, including Transformers 5.6.2, while preserving golden output compatibility with the pinned environment.
112
+
113
  ### Usage example
114
 
115
  ```python
 
171
 
172
  **Update**: Nemotron-Parse-v1.1 is now [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).
173
 
174
+ The model has also been validated with vLLM 0.20.1 using both implicit prompts and explicit encoder/decoder prompts.
175
+
176
  Note: when running on A100/A10 we recommend running vllm serve with _--attention-backend=TRITON_ATTN_
177
 
178
  You will need to install *albumentations* on top, and then follow the VLLM Inference example below:
 
380
  primaryClass={cs.LG},
381
  url={https://arxiv.org/abs/2511.20478},
382
  }
383
+ ```
config.json CHANGED
@@ -99,8 +99,9 @@
99
  },
100
  "decoder_start_token_id": 2,
101
  "encoder": {
102
- "_attn_implementation": "eager",
103
  "_name_or_path": "nvidia/C-RADIOv2-H",
 
104
  "adaptor_configs": {},
105
  "adaptor_names": null,
106
  "add_cross_attention": false,
@@ -382,7 +383,7 @@
382
  "torchscript": false,
383
  "transformers_version": "4.51.3",
384
  "typical_p": 1.0,
385
- "use_bfloat16": true,
386
  "version": "radio_v2.5-h",
387
  "vitdet_window_size": null
388
  },
 
99
  },
100
  "decoder_start_token_id": 2,
101
  "encoder": {
102
+ "attn_implementation": "eager",
103
  "_name_or_path": "nvidia/C-RADIOv2-H",
104
+ "image_processor_normalizes": true,
105
  "adaptor_configs": {},
106
  "adaptor_names": null,
107
  "add_cross_attention": false,
 
383
  "torchscript": false,
384
  "transformers_version": "4.51.3",
385
  "typical_p": 1.0,
386
+ "use_bfloat16": false,
387
  "version": "radio_v2.5-h",
388
  "vitdet_window_size": null
389
  },
docker-compose.yaml ADDED
@@ -0,0 +1,33 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ services:
2
+ nemotron-parse:
3
+ build: .
4
+ image: nemotron-parse:latest
5
+ working_dir: /workspace
6
+
7
+ deploy:
8
+ resources:
9
+ reservations:
10
+ devices:
11
+ - driver: nvidia
12
+ count: all
13
+ capabilities: [gpu]
14
+
15
+ volumes:
16
+ # Project source - live-edits on the host are reflected immediately.
17
+ - .:/workspace
18
+ # Shadow the host .venv (wrong Python version) so uv uses /opt/venv inside
19
+ # the container rather than the host-side virtual environment.
20
+ - /workspace/.venv
21
+ # HuggingFace model cache - persists the RADIO encoder weights across
22
+ # container restarts so they are not re-downloaded on every run.
23
+ - hf-cache:/root/.cache/huggingface
24
+
25
+ environment:
26
+ HF_HOME: /root/.cache/huggingface
27
+
28
+ # Keep the container alive for interactive use (exec, attach, etc.).
29
+ stdin_open: true
30
+ tty: true
31
+
32
+ volumes:
33
+ hf-cache:
golden_outputs.json ADDED
@@ -0,0 +1,133 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "metadata": {
3
+ "transformers_version": "4.51.3",
4
+ "torch_version": "2.11.0+cu130",
5
+ "device": "cuda:0",
6
+ "dtype": "torch.bfloat16",
7
+ "model_path": "/home/oholworthy/Code/NeMo/nemo-retriever-research/models/NVIDIA-Nemotron-Parse-v1.1"
8
+ },
9
+ "image_processing": {
10
+ "shape": [
11
+ 1,
12
+ 3,
13
+ 2048,
14
+ 1648
15
+ ],
16
+ "mean": 0.06600654870271683,
17
+ "std": 0.23970627784729004,
18
+ "first_20_values": [
19
+ 0.0,
20
+ 0.0,
21
+ 0.0,
22
+ 0.0,
23
+ 0.0,
24
+ 0.0,
25
+ 0.0,
26
+ 0.0,
27
+ 0.0,
28
+ 0.0,
29
+ 0.0,
30
+ 0.0,
31
+ 0.0,
32
+ 0.0,
33
+ 0.0,
34
+ 0.0,
35
+ 0.0,
36
+ 0.0,
37
+ 0.0,
38
+ 0.0
39
+ ]
40
+ },
41
+ "encoder_output": {
42
+ "shape": [
43
+ 1,
44
+ 3201,
45
+ 1024
46
+ ],
47
+ "mean": -0.0010558579815551639,
48
+ "std": 0.9514492154121399,
49
+ "token0_first16": [
50
+ 0.1474609375,
51
+ 0.84765625,
52
+ 2.09375,
53
+ -0.002960205078125,
54
+ 0.640625,
55
+ 0.494140625,
56
+ -0.2373046875,
57
+ -0.9921875,
58
+ -0.220703125,
59
+ 0.5078125,
60
+ 0.02734375,
61
+ 0.0615234375,
62
+ -0.267578125,
63
+ -1.796875,
64
+ -0.1962890625,
65
+ 0.17578125
66
+ ]
67
+ },
68
+ "forward_pass": {
69
+ "logits_shape": [
70
+ 1,
71
+ 1,
72
+ 52352
73
+ ],
74
+ "top_k_indices": [
75
+ 0,
76
+ 221,
77
+ 2,
78
+ 276,
79
+ 28,
80
+ 501,
81
+ 702,
82
+ 1004,
83
+ 123,
84
+ 23
85
+ ],
86
+ "top_k_values": [
87
+ 54.75,
88
+ 41.0,
89
+ 40.5,
90
+ 40.25,
91
+ 40.25,
92
+ 40.25,
93
+ 40.0,
94
+ 40.0,
95
+ 40.0,
96
+ 40.0
97
+ ]
98
+ },
99
+ "generation": {
100
+ "max_new_tokens": 50,
101
+ "token_ids": [
102
+ 2,
103
+ 0,
104
+ 50004,
105
+ 50008,
106
+ 50001,
107
+ 50409,
108
+ 51797,
109
+ 82,
110
+ 2722,
111
+ 113,
112
+ 18121,
113
+ 579,
114
+ 115,
115
+ 113,
116
+ 19321,
117
+ 89,
118
+ 115,
119
+ 221,
120
+ 82,
121
+ 493,
122
+ 113,
123
+ 18121,
124
+ 579,
125
+ 115,
126
+ 50632,
127
+ 51847,
128
+ 52325,
129
+ 2
130
+ ],
131
+ "decoded_text": "</s><s><predict_bbox><predict_classes><output_markdown><x_0.3906><y_0.5969>\\begin{tabular}{ccccc}\n\\end{tabular}<x_0.6084><y_0.6359><class_Table></s>"
132
+ }
133
+ }
hf_nemotron_parse_modeling.py CHANGED
@@ -23,6 +23,38 @@ from transformers.modeling_attn_mask_utils import (
23
  _prepare_4d_causal_attention_mask_for_sdpa,
24
  )
25
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
26
 
27
  class NemotronParseDecoder(MBartPreTrainedModel):
28
  """
@@ -47,7 +79,11 @@ class NemotronParseDecoder(MBartPreTrainedModel):
47
  if embed_tokens is not None:
48
  self.embed_tokens.weight = embed_tokens.weight
49
 
50
- self.layers = nn.ModuleList([MBartDecoderLayer(config) for _ in range(config.decoder_layers)])
 
 
 
 
51
  self.config = config
52
 
53
  self.layernorm_embedding = nn.LayerNorm(config.d_model)
@@ -163,8 +199,8 @@ class NemotronParseDecoder(MBartPreTrainedModel):
163
  else:
164
  raise ValueError("You have to specify either decoder_input_ids or decoder_inputs_embeds")
165
 
166
- # past_key_values_length
167
- past_key_values_length = past_key_values[0][0].shape[2] if past_key_values is not None else 0
168
 
169
  if inputs_embeds is None:
170
  inputs_embeds = self.embed_tokens(input_ids)
@@ -221,7 +257,22 @@ class NemotronParseDecoder(MBartPreTrainedModel):
221
  all_hidden_states = () if output_hidden_states else None
222
  all_self_attns = () if output_attentions else None
223
  all_cross_attentions = () if (output_attentions and encoder_hidden_states is not None) else None
224
- next_decoder_cache = () if use_cache else None
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
225
 
226
  # check if head_mask/cross_attn_head_mask has a correct number of layers specified if desired
227
  for attn_mask, mask_name in zip([head_mask, cross_attn_head_mask], ["head_mask", "cross_attn_head_mask"]):
@@ -240,45 +291,68 @@ class NemotronParseDecoder(MBartPreTrainedModel):
240
  if dropout_probability < self.layerdrop:
241
  continue
242
 
243
- past_key_value = past_key_values[idx] if past_key_values is not None else None
244
-
245
- if self.gradient_checkpointing and self.training:
246
- layer_outputs = self._gradient_checkpointing_func(
247
- decoder_layer.__call__,
248
- hidden_states,
249
- attention_mask,
250
- encoder_hidden_states,
251
- encoder_attention_mask,
252
- head_mask[idx] if head_mask is not None else None,
253
- cross_attn_head_mask[idx] if cross_attn_head_mask is not None else None,
254
- None,
255
- output_attentions,
256
- use_cache,
257
- )
258
- else:
259
  layer_outputs = decoder_layer(
260
  hidden_states,
261
  attention_mask=attention_mask,
262
  encoder_hidden_states=encoder_hidden_states,
263
  encoder_attention_mask=encoder_attention_mask,
264
- layer_head_mask=(head_mask[idx] if head_mask is not None else None),
265
- cross_attn_layer_head_mask=(
266
- cross_attn_head_mask[idx] if cross_attn_head_mask is not None else None
267
- ),
268
- past_key_value=past_key_value,
269
- output_attentions=output_attentions,
270
  use_cache=use_cache,
271
  )
272
- hidden_states = layer_outputs[0]
273
-
274
- if use_cache:
275
- next_decoder_cache += (layer_outputs[3 if output_attentions else 1],)
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
276
 
277
- if output_attentions:
278
- all_self_attns += (layer_outputs[1],)
 
 
 
279
 
280
- if encoder_hidden_states is not None:
281
- all_cross_attentions += (layer_outputs[2],)
 
 
282
 
283
  hidden_states = self.layer_norm(hidden_states)
284
 
@@ -308,6 +382,8 @@ class RadioWithNeck(nn.Module):
308
  def __init__(self, config):
309
  super().__init__()
310
  self.config = config
 
 
311
 
312
  self.model_encoder = AutoModel.from_config(config, trust_remote_code=True)
313
 
@@ -319,8 +395,26 @@ class RadioWithNeck(nn.Module):
319
  self.layer_norm2 = nn.LayerNorm(last_hidden_state, eps=1e-06, elementwise_affine=True)
320
  self.sum_proj = nn.Linear(3840, last_hidden_state)
321
  self.layer_norm3 = nn.LayerNorm(last_hidden_state, eps=1e-06, elementwise_affine=True)
 
 
 
 
 
 
 
 
 
 
 
 
 
 
322
 
323
  def forward(self, pixel_values, output_attentions=False, output_hidden_states=False, return_dict=False, **kwargs):
 
 
 
 
324
  radio_output = self.model_encoder(pixel_values)
325
  summary, feature = radio_output
326
 
@@ -533,6 +627,30 @@ class NemotronParseForConditionalGeneration(NemotronParsePreTrainedModel, Genera
533
  encoder_attentions=encoder_outputs.attentions,
534
  )
535
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
536
  def prepare_decoder_input_ids_from_labels(self, labels: torch.Tensor):
537
  return shift_tokens_right(labels, self.config.pad_token_id, self.config.decoder_start_token_id)
538
 
 
23
  _prepare_4d_causal_attention_mask_for_sdpa,
24
  )
25
 
26
+ # ---------------------------------------------------------------------------
27
+ # Cache compatibility (transformers 5.x introduced Cache objects;
28
+ # 4.x used plain tuple-of-tuples for past_key_values)
29
+ # ---------------------------------------------------------------------------
30
+ import inspect
31
+ try:
32
+ from transformers.cache_utils import Cache as _CacheBase
33
+ def _is_cache_object(obj) -> bool:
34
+ return isinstance(obj, _CacheBase)
35
+ except ImportError:
36
+ def _is_cache_object(obj) -> bool:
37
+ return False
38
+
39
+ def _past_key_values_length(past_key_values) -> int:
40
+ """Return the number of already-decoded tokens regardless of cache format."""
41
+ if past_key_values is None:
42
+ return 0
43
+ if _is_cache_object(past_key_values):
44
+ return past_key_values.get_seq_length()
45
+ return past_key_values[0][0].shape[2]
46
+
47
+ # ---------------------------------------------------------------------------
48
+ # MBartDecoderLayer API detection
49
+ #
50
+ # transformers <~4.57: forward() takes `past_key_value` (singular), returns a
51
+ # tuple (hidden_states, [attentions], [present_key_value])
52
+ # transformers >=~4.57: forward() takes `past_key_values` (plural, Cache).
53
+ # True 5.x returns a single torch.Tensor (cache updated in-place);
54
+ # intermediate versions (e.g. 4.57.x) still return a tuple.
55
+ # ---------------------------------------------------------------------------
56
+ _layer_takes_plural_past_kv = 'past_key_values' in inspect.signature(MBartDecoderLayer.forward).parameters
57
+
58
 
59
  class NemotronParseDecoder(MBartPreTrainedModel):
60
  """
 
79
  if embed_tokens is not None:
80
  self.embed_tokens.weight = embed_tokens.weight
81
 
82
+ _layer_supports_idx = 'layer_idx' in inspect.signature(MBartDecoderLayer.__init__).parameters
83
+ self.layers = nn.ModuleList([
84
+ MBartDecoderLayer(config, layer_idx=i) if _layer_supports_idx else MBartDecoderLayer(config)
85
+ for i in range(config.decoder_layers)
86
+ ])
87
  self.config = config
88
 
89
  self.layernorm_embedding = nn.LayerNorm(config.d_model)
 
199
  else:
200
  raise ValueError("You have to specify either decoder_input_ids or decoder_inputs_embeds")
201
 
202
+ # past_key_values_length works with tuple-of-tuples (4.x) and Cache objects (5.x).
203
+ past_key_values_length = _past_key_values_length(past_key_values)
204
 
205
  if inputs_embeds is None:
206
  inputs_embeds = self.embed_tokens(input_ids)
 
257
  all_hidden_states = () if output_hidden_states else None
258
  all_self_attns = () if output_attentions else None
259
  all_cross_attentions = () if (output_attentions and encoder_hidden_states is not None) else None
260
+ # In 5.x the Cache object is updated in-place by each layer, so we just
261
+ # carry the same object through. In 4.x we collect per-layer tuples.
262
+ _using_cache_obj = _is_cache_object(past_key_values)
263
+ next_decoder_cache = past_key_values if (_using_cache_obj and use_cache) else (() if use_cache else None)
264
+
265
+ # 5.x: on the first call (past_key_values=None), create an EncoderDecoderCache
266
+ # so each MBartAttention layer can populate cross-/self-attention KV states
267
+ # in-place. This enables proper KV caching during multi-step generation.
268
+ if _layer_takes_plural_past_kv and use_cache and past_key_values is None:
269
+ try:
270
+ from transformers.cache_utils import EncoderDecoderCache, DynamicCache
271
+ past_key_values = EncoderDecoderCache(DynamicCache(), DynamicCache())
272
+ _using_cache_obj = True
273
+ next_decoder_cache = past_key_values
274
+ except (ImportError, AttributeError, TypeError):
275
+ pass # fallback: layers recompute KV each step (correct but slower)
276
 
277
  # check if head_mask/cross_attn_head_mask has a correct number of layers specified if desired
278
  for attn_mask, mask_name in zip([head_mask, cross_attn_head_mask], ["head_mask", "cross_attn_head_mask"]):
 
291
  if dropout_probability < self.layerdrop:
292
  continue
293
 
294
+ if _layer_takes_plural_past_kv:
295
+ # Plural-param API: cache updated in-place, nothing to collect.
 
 
 
 
 
 
 
 
 
 
 
 
 
 
296
  layer_outputs = decoder_layer(
297
  hidden_states,
298
  attention_mask=attention_mask,
299
  encoder_hidden_states=encoder_hidden_states,
300
  encoder_attention_mask=encoder_attention_mask,
301
+ past_key_values=past_key_values if use_cache else None,
 
 
 
 
 
302
  use_cache=use_cache,
303
  )
304
+ # True 5.x returns a single Tensor; intermediate versions
305
+ # (e.g. 4.57.x) have the renamed parameter but still return
306
+ # a tuple, so handle both.
307
+ hidden_states = layer_outputs if isinstance(layer_outputs, torch.Tensor) else layer_outputs[0]
308
+ else:
309
+ # Singular-param API: returns a tuple, collect cache per-layer.
310
+ if past_key_values is None:
311
+ past_key_value = None
312
+ elif _using_cache_obj:
313
+ past_key_value = past_key_values # full Cache object
314
+ else:
315
+ past_key_value = past_key_values[idx] # per-layer tuple
316
+
317
+ if self.gradient_checkpointing and self.training:
318
+ layer_outputs = self._gradient_checkpointing_func(
319
+ decoder_layer.__call__,
320
+ hidden_states,
321
+ attention_mask,
322
+ encoder_hidden_states,
323
+ encoder_attention_mask,
324
+ head_mask[idx] if head_mask is not None else None,
325
+ cross_attn_head_mask[idx] if cross_attn_head_mask is not None else None,
326
+ None,
327
+ output_attentions,
328
+ use_cache,
329
+ )
330
+ else:
331
+ layer_outputs = decoder_layer(
332
+ hidden_states,
333
+ attention_mask=attention_mask,
334
+ encoder_hidden_states=encoder_hidden_states,
335
+ encoder_attention_mask=encoder_attention_mask,
336
+ layer_head_mask=(head_mask[idx] if head_mask is not None else None),
337
+ cross_attn_layer_head_mask=(
338
+ cross_attn_head_mask[idx] if cross_attn_head_mask is not None else None
339
+ ),
340
+ past_key_value=past_key_value,
341
+ output_attentions=output_attentions,
342
+ use_cache=use_cache,
343
+ )
344
+ hidden_states = layer_outputs[0]
345
 
346
+ if use_cache and not _using_cache_obj:
347
+ # 4.x: cache is the last element of layer_outputs.
348
+ cache_idx = 3 if output_attentions else 1
349
+ if len(layer_outputs) > cache_idx:
350
+ next_decoder_cache += (layer_outputs[cache_idx],)
351
 
352
+ if output_attentions:
353
+ all_self_attns += (layer_outputs[1],)
354
+ if encoder_hidden_states is not None:
355
+ all_cross_attentions += (layer_outputs[2],)
356
 
357
  hidden_states = self.layer_norm(hidden_states)
358
 
 
382
  def __init__(self, config):
383
  super().__init__()
384
  self.config = config
385
+ self.image_processor_normalizes = bool(getattr(config, "image_processor_normalizes", False))
386
+ self._radio_input_conditioner_externalized = False
387
 
388
  self.model_encoder = AutoModel.from_config(config, trust_remote_code=True)
389
 
 
395
  self.layer_norm2 = nn.LayerNorm(last_hidden_state, eps=1e-06, elementwise_affine=True)
396
  self.sum_proj = nn.Linear(3840, last_hidden_state)
397
  self.layer_norm3 = nn.LayerNorm(last_hidden_state, eps=1e-06, elementwise_affine=True)
398
+
399
+ def _externalize_radio_input_conditioner(self):
400
+ if not self.image_processor_normalizes or self._radio_input_conditioner_externalized:
401
+ return
402
+
403
+ make_external = getattr(self.model_encoder, "make_preprocessor_external", None)
404
+ if make_external is None:
405
+ raise ValueError(
406
+ "image_processor_normalizes=True requires a RADIO encoder with "
407
+ "make_preprocessor_external()."
408
+ )
409
+
410
+ make_external()
411
+ self._radio_input_conditioner_externalized = True
412
 
413
  def forward(self, pixel_values, output_attentions=False, output_hidden_states=False, return_dict=False, **kwargs):
414
+ self._externalize_radio_input_conditioner()
415
+ if self.image_processor_normalizes:
416
+ dtype = next(self.model_encoder.parameters()).dtype
417
+ pixel_values = pixel_values.to(dtype=dtype)
418
  radio_output = self.model_encoder(pixel_values)
419
  summary, feature = radio_output
420
 
 
627
  encoder_attentions=encoder_outputs.attentions,
628
  )
629
 
630
+ def prepare_inputs_for_generation(
631
+ self,
632
+ input_ids,
633
+ past_key_values=None,
634
+ attention_mask=None,
635
+ use_cache=None,
636
+ encoder_outputs=None,
637
+ **kwargs,
638
+ ):
639
+ if past_key_values is not None:
640
+ past_length = _past_key_values_length(past_key_values)
641
+ if input_ids.shape[1] > past_length:
642
+ input_ids = input_ids[:, past_length:]
643
+ else:
644
+ input_ids = input_ids[:, -1:]
645
+ return {
646
+ "pixel_values": None, # encoder_outputs carries the image features
647
+ "encoder_outputs": encoder_outputs,
648
+ "past_key_values": past_key_values,
649
+ "decoder_input_ids": input_ids,
650
+ "decoder_attention_mask": attention_mask,
651
+ "use_cache": use_cache,
652
+ }
653
+
654
  def prepare_decoder_input_ids_from_labels(self, labels: torch.Tensor):
655
  return shift_tokens_right(labels, self.config.pad_token_id, self.config.decoder_start_token_id)
656
 
hf_nemotron_parse_processor.py CHANGED
@@ -7,12 +7,15 @@ import albumentations as A
7
  import cv2
8
  import json
9
 
10
- from transformers import ProcessorMixin, BaseImageProcessor, ImageProcessingMixin
11
- from transformers.tokenization_utils_base import BatchEncoding
12
  from transformers.image_utils import ChannelDimension, ImageInput, PILImageResampling, infer_channel_dimension_format
13
  from transformers.utils import TensorType
14
 
15
 
 
 
 
 
16
  class NemotronParseImageProcessor(BaseImageProcessor, ImageProcessingMixin):
17
  """
18
  Image processor for NemotronParse model.
@@ -25,6 +28,9 @@ class NemotronParseImageProcessor(BaseImageProcessor, ImageProcessingMixin):
25
  def __init__(
26
  self,
27
  final_size: tuple = (2048, 1648),
 
 
 
28
  **kwargs,
29
  ):
30
  clean_kwargs = {}
@@ -45,6 +51,9 @@ class NemotronParseImageProcessor(BaseImageProcessor, ImageProcessingMixin):
45
  final_size = (int(size_config['height']), int(size_config['width']))
46
 
47
  super().__init__(**clean_kwargs)
 
 
 
48
 
49
  if isinstance(final_size, (list, tuple)) and len(final_size) >= 2:
50
  self.final_size = (int(final_size[0]), int(final_size[1]))
@@ -74,7 +83,6 @@ class NemotronParseImageProcessor(BaseImageProcessor, ImageProcessingMixin):
74
 
75
  self.torch_transform = T.Compose([
76
  T.ToTensor(),
77
- # Note: Normalization is done within RADIO model
78
  ])
79
 
80
  def to_dict(self):
@@ -101,6 +109,10 @@ class NemotronParseImageProcessor(BaseImageProcessor, ImageProcessingMixin):
101
  final_size = config_dict['final_size']
102
  if isinstance(final_size, (list, tuple)):
103
  config_dict['final_size'] = tuple(int(x) for x in final_size)
 
 
 
 
104
 
105
  try:
106
  return cls(**config_dict, **kwargs)
@@ -120,13 +132,16 @@ class NemotronParseImageProcessor(BaseImageProcessor, ImageProcessingMixin):
120
  config = {
121
  "feature_extractor_type": "NemotronParseImageProcessor",
122
  "image_processor_type": "NemotronParseImageProcessor",
123
- "processor_class": "NemotronParseImageProcessor",
124
  "size": {
125
  "height": self.final_size[0],
126
  "width": self.final_size[1],
127
  "longest_edge": self.final_size
128
  },
129
  "final_size": self.final_size,
 
 
 
130
  }
131
 
132
  config_path = os.path.join(save_directory, "preprocessor_config.json")
@@ -231,6 +246,11 @@ class NemotronParseImageProcessor(BaseImageProcessor, ImageProcessingMixin):
231
  pixel_values.append(pixel_values_tensor)
232
 
233
  pixel_values = torch.stack(pixel_values)
 
 
 
 
 
234
 
235
  data = {"pixel_values": pixel_values}
236
 
@@ -285,7 +305,7 @@ class NemotronParseProcessor(ProcessorMixin):
285
  verbose: bool = True,
286
  return_tensors: Optional[Union[str, "TensorType"]] = None,
287
  **kwargs
288
- ) -> BatchEncoding:
289
  """
290
  Main method to prepare for the model one or several text(s) and image(s).
291
  """
@@ -320,7 +340,7 @@ class NemotronParseProcessor(ProcessorMixin):
320
  text_inputs = {}
321
 
322
  # Combine inputs
323
- return BatchEncoding({**image_inputs, **text_inputs})
324
 
325
  def decode(self, *args, **kwargs):
326
  """Decode token ids to strings."""
 
7
  import cv2
8
  import json
9
 
10
+ from transformers import BatchFeature, ProcessorMixin, BaseImageProcessor, ImageProcessingMixin
 
11
  from transformers.image_utils import ChannelDimension, ImageInput, PILImageResampling, infer_channel_dimension_format
12
  from transformers.utils import TensorType
13
 
14
 
15
+ OPENAI_CLIP_MEAN = (0.48145466, 0.4578275, 0.40821073)
16
+ OPENAI_CLIP_STD = (0.26862954, 0.26130258, 0.27577711)
17
+
18
+
19
  class NemotronParseImageProcessor(BaseImageProcessor, ImageProcessingMixin):
20
  """
21
  Image processor for NemotronParse model.
 
28
  def __init__(
29
  self,
30
  final_size: tuple = (2048, 1648),
31
+ do_normalize: bool = False,
32
+ image_mean: Optional[Union[List[float], tuple]] = None,
33
+ image_std: Optional[Union[List[float], tuple]] = None,
34
  **kwargs,
35
  ):
36
  clean_kwargs = {}
 
51
  final_size = (int(size_config['height']), int(size_config['width']))
52
 
53
  super().__init__(**clean_kwargs)
54
+ self.do_normalize = bool(do_normalize)
55
+ self.image_mean = list(image_mean or OPENAI_CLIP_MEAN)
56
+ self.image_std = list(image_std or OPENAI_CLIP_STD)
57
 
58
  if isinstance(final_size, (list, tuple)) and len(final_size) >= 2:
59
  self.final_size = (int(final_size[0]), int(final_size[1]))
 
83
 
84
  self.torch_transform = T.Compose([
85
  T.ToTensor(),
 
86
  ])
87
 
88
  def to_dict(self):
 
109
  final_size = config_dict['final_size']
110
  if isinstance(final_size, (list, tuple)):
111
  config_dict['final_size'] = tuple(int(x) for x in final_size)
112
+ if 'image_mean' in config_dict:
113
+ config_dict['image_mean'] = [float(x) for x in config_dict['image_mean']]
114
+ if 'image_std' in config_dict:
115
+ config_dict['image_std'] = [float(x) for x in config_dict['image_std']]
116
 
117
  try:
118
  return cls(**config_dict, **kwargs)
 
132
  config = {
133
  "feature_extractor_type": "NemotronParseImageProcessor",
134
  "image_processor_type": "NemotronParseImageProcessor",
135
+ "processor_class": "NemotronParseProcessor",
136
  "size": {
137
  "height": self.final_size[0],
138
  "width": self.final_size[1],
139
  "longest_edge": self.final_size
140
  },
141
  "final_size": self.final_size,
142
+ "do_normalize": self.do_normalize,
143
+ "image_mean": self.image_mean,
144
+ "image_std": self.image_std,
145
  }
146
 
147
  config_path = os.path.join(save_directory, "preprocessor_config.json")
 
246
  pixel_values.append(pixel_values_tensor)
247
 
248
  pixel_values = torch.stack(pixel_values)
249
+
250
+ if self.do_normalize:
251
+ mean = pixel_values.new_tensor(self.image_mean).view(1, -1, 1, 1)
252
+ std = pixel_values.new_tensor(self.image_std).view(1, -1, 1, 1)
253
+ pixel_values = (pixel_values - mean) / std
254
 
255
  data = {"pixel_values": pixel_values}
256
 
 
305
  verbose: bool = True,
306
  return_tensors: Optional[Union[str, "TensorType"]] = None,
307
  **kwargs
308
+ ) -> BatchFeature:
309
  """
310
  Main method to prepare for the model one or several text(s) and image(s).
311
  """
 
340
  text_inputs = {}
341
 
342
  # Combine inputs
343
+ return BatchFeature(data={**image_inputs, **text_inputs})
344
 
345
  def decode(self, *args, **kwargs):
346
  """Decode token ids to strings."""
preprocessor_config.json CHANGED
@@ -6,7 +6,17 @@
6
  "AutoImageProcessor": "hf_nemotron_parse_processor.NemotronParseImageProcessor",
7
  "AutoProcessor": "hf_nemotron_parse_processor.NemotronParseProcessor"
8
  },
9
- "do_normalize": false,
 
 
 
 
 
 
 
 
 
 
10
  "do_rescale": true,
11
  "rescale_factor": 0.00392156862745098,
12
  "size": {
 
6
  "AutoImageProcessor": "hf_nemotron_parse_processor.NemotronParseImageProcessor",
7
  "AutoProcessor": "hf_nemotron_parse_processor.NemotronParseProcessor"
8
  },
9
+ "do_normalize": true,
10
+ "image_mean": [
11
+ 0.48145466,
12
+ 0.4578275,
13
+ 0.40821073
14
+ ],
15
+ "image_std": [
16
+ 0.26862954,
17
+ 0.26130258,
18
+ 0.27577711
19
+ ],
20
  "do_rescale": true,
21
  "rescale_factor": 0.00392156862745098,
22
  "size": {
pyproject.toml ADDED
@@ -0,0 +1,33 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ [project]
2
+ name = "nemotron-parse"
3
+ version = "1.1.0"
4
+ description = "NVIDIA Nemotron-Parse document parsing model"
5
+ requires-python = ">=3.10"
6
+ dependencies = [
7
+ "transformers>=4.51.3",
8
+ "accelerate==1.12.0",
9
+ "albumentations==2.0.8",
10
+ "timm==1.0.22",
11
+ "einops",
12
+ "Pillow",
13
+ "numpy",
14
+ "opencv-python-headless",
15
+ "beautifulsoup4",
16
+ "open-clip-torch>=3.3.0",
17
+ "pytest>=9.0.3",
18
+ ]
19
+
20
+ [project.optional-dependencies]
21
+ # vLLM serving (install separately in the serving container).
22
+ vllm = ["openai"]
23
+ # Development / testing.
24
+ dev = ["pytest"]
25
+
26
+ [tool.uv]
27
+ # This repo is a model directory loaded via trust_remote_code, not an
28
+ # installable Python package, so uv should only manage dependencies.
29
+ package = false
30
+
31
+ # torch and torchvision ship pre-compiled with CUDA support inside the
32
+ # NVIDIA base image (nvcr.io/nvidia/pytorch:*). uv must not overwrite them.
33
+ exclude-dependencies = ["torch", "torchvision"]
test_golden.py ADDED
@@ -0,0 +1,550 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """
2
+ Golden reference tests for NVIDIA-Nemotron-Parse-v1.1.
3
+
4
+ Captures reference outputs from the pinned dependency set, then verifies the
5
+ same outputs after dependency changes.
6
+
7
+ WORKFLOW
8
+ --------
9
+ Step 1 - capture (run once against pinned deps, e.g. transformers==4.51.3):
10
+
11
+ python test_golden.py --capture [--model-path /path/to/model]
12
+
13
+ This writes golden_outputs.json next to this file.
14
+
15
+ Step 2 - verify (run against new deps):
16
+
17
+ pytest test_golden.py -v
18
+
19
+ All tests skip automatically if golden_outputs.json is missing.
20
+
21
+ TEST LAYERS
22
+ -----------
23
+ 1. Image preprocessing - pixel value stats + first-N raw values (no GPU needed)
24
+ 2. Encoder output - hidden state shape, mean, std, and a fixed-position slice
25
+ 3. Decoder forward pass - top-k logit indices and values at a fixed decoder step
26
+ 4. Generation - exact token ID sequence for 50 greedy-decoded tokens
27
+ """
28
+
29
+ import json
30
+ import os
31
+ import sys
32
+ import pytest
33
+ import numpy as np
34
+ import torch
35
+ from pathlib import Path
36
+
37
+ # ---------------------------------------------------------------------------
38
+ # Paths / constants
39
+ # ---------------------------------------------------------------------------
40
+ MODEL_PATH = str(Path(__file__).parent)
41
+ GOLDEN_FILE = Path(__file__).parent / "golden_outputs.json"
42
+
43
+ TASK_PROMPT = "</s><s><predict_bbox><predict_classes><output_markdown>"
44
+ MAX_NEW_TOKENS_GOLDEN = 50 # short enough to be fast, long enough to be meaningful
45
+ TOP_K = 10 # number of top logit predictions to capture
46
+
47
+
48
+ # ---------------------------------------------------------------------------
49
+ # Deterministic test image (no external files required)
50
+ # ---------------------------------------------------------------------------
51
+ def make_test_image():
52
+ """Return a fully deterministic PIL image that loosely resembles a document."""
53
+ from PIL import Image, ImageDraw
54
+
55
+ img = Image.new("RGB", (400, 600), color=(255, 255, 255))
56
+ draw = ImageDraw.Draw(img)
57
+
58
+ # Title bar
59
+ draw.rectangle([20, 20, 380, 80], fill=(210, 210, 210))
60
+
61
+ # Body text area with ruled lines
62
+ draw.rectangle([20, 100, 380, 480], fill=(245, 245, 245))
63
+ for y in range(120, 470, 18):
64
+ draw.line([(40, y), (360, y)], fill=(170, 170, 170), width=1)
65
+
66
+ # Table-like grid at the bottom
67
+ draw.rectangle([20, 500, 380, 580], fill=(200, 220, 200))
68
+ for x in range(80, 380, 80):
69
+ draw.line([(x, 500), (x, 580)], fill=(100, 140, 100), width=1)
70
+ for y in range(520, 580, 20):
71
+ draw.line([(20, y), (380, y)], fill=(100, 140, 100), width=1)
72
+
73
+ return img
74
+
75
+
76
+ # ---------------------------------------------------------------------------
77
+ # Golden file helpers
78
+ # ---------------------------------------------------------------------------
79
+ def load_golden():
80
+ if GOLDEN_FILE.exists():
81
+ with open(GOLDEN_FILE) as f:
82
+ return json.load(f)
83
+ return None
84
+
85
+
86
+ def save_golden(data: dict):
87
+ with open(GOLDEN_FILE, "w") as f:
88
+ json.dump(data, f, indent=2)
89
+
90
+
91
+ def _pixel_values_for_golden(processor, pixel_values: torch.Tensor) -> torch.Tensor:
92
+ """Return the canonical raw pixel values used by the original golden file."""
93
+ pv = pixel_values.float()
94
+ image_processor = getattr(processor, "image_processor", None)
95
+ if not getattr(image_processor, "do_normalize", False):
96
+ return pv
97
+
98
+ mean = pv.new_tensor(image_processor.image_mean).view(1, -1, 1, 1)
99
+ std = pv.new_tensor(image_processor.image_std).view(1, -1, 1, 1)
100
+ return pv * std + mean
101
+
102
+
103
+ # ---------------------------------------------------------------------------
104
+ # Pytest fixtures (session-scoped so the model is loaded only once)
105
+ # ---------------------------------------------------------------------------
106
+ @pytest.fixture(scope="session")
107
+ def env():
108
+ """Load model, processor, and tokenizer once for the whole test session."""
109
+ import torch
110
+ from transformers import AutoModel, AutoProcessor, AutoTokenizer
111
+
112
+ device = "cuda:0" if torch.cuda.is_available() else "cpu"
113
+ dtype = torch.bfloat16 if torch.cuda.is_available() else torch.float32
114
+
115
+ print(f"\nLoading model from {MODEL_PATH} on {device} ({dtype})...")
116
+ model = AutoModel.from_pretrained(
117
+ MODEL_PATH,
118
+ trust_remote_code=True,
119
+ torch_dtype=dtype,
120
+ ).to(device).eval()
121
+
122
+ tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH, trust_remote_code=True)
123
+ processor = AutoProcessor.from_pretrained(MODEL_PATH, trust_remote_code=True)
124
+
125
+ return dict(model=model, tokenizer=tokenizer, processor=processor,
126
+ device=device, dtype=dtype)
127
+
128
+
129
+ @pytest.fixture(scope="session")
130
+ def processed_inputs(env):
131
+ """Preprocess the test image once for the whole session."""
132
+ import torch
133
+
134
+ image = make_test_image()
135
+ inputs = env["processor"](
136
+ images=[image],
137
+ text=TASK_PROMPT,
138
+ return_tensors="pt",
139
+ add_special_tokens=False,
140
+ ).to(env["device"])
141
+
142
+ return inputs, image
143
+
144
+
145
+ @pytest.fixture(scope="session")
146
+ def golden():
147
+ """Load golden data; tests that need it skip if the file is absent."""
148
+ data = load_golden()
149
+ if data is None:
150
+ pytest.skip("golden_outputs.json not found - run: python test_golden.py --capture")
151
+ return data
152
+
153
+
154
+ # ---------------------------------------------------------------------------
155
+ # Layer 1: Image preprocessing
156
+ # (Does not require a model or GPU - fast sanity check on the processor.)
157
+ # ---------------------------------------------------------------------------
158
+ class TestImageProcessing:
159
+ def test_pixel_values_shape(self, processed_inputs):
160
+ inputs, _ = processed_inputs
161
+ pv = inputs["pixel_values"]
162
+ # Model expects 2048x1648 images.
163
+ assert list(pv.shape) == [1, 3, 2048, 1648], f"Unexpected shape: {pv.shape}"
164
+
165
+ def test_pixel_values_dtype(self, processed_inputs):
166
+ inputs, _ = processed_inputs
167
+ # Image preprocessing keeps float32; the model casts internally.
168
+ assert inputs["pixel_values"].dtype == torch.float32
169
+
170
+ def test_pixel_value_range(self, processed_inputs):
171
+ """Values should be CLIP-normalized after image preprocessing."""
172
+ pv = processed_inputs[0]["pixel_values"].float()
173
+ assert pv.min() >= -2.0, f"Pixel values unexpectedly low: {pv.min()}"
174
+ assert pv.max() <= 2.5, f"Pixel values unexpectedly high: {pv.max()}"
175
+
176
+ def test_pixel_values_stats_match_golden(self, processed_inputs, env, golden):
177
+ pv = _pixel_values_for_golden(env["processor"], processed_inputs[0]["pixel_values"])
178
+ g = golden["image_processing"]
179
+
180
+ assert abs(pv.mean().item() - g["mean"]) < 1e-4, \
181
+ f"mean changed: {pv.mean().item():.6f} vs golden {g['mean']:.6f}"
182
+ assert abs(pv.std().item() - g["std"]) < 1e-4, \
183
+ f"std changed: {pv.std().item():.6f} vs golden {g['std']:.6f}"
184
+
185
+ def test_pixel_values_first_values_match_golden(self, processed_inputs, env, golden):
186
+ """Exact match on the first 20 float values (catches transform-order bugs)."""
187
+ pv = _pixel_values_for_golden(env["processor"], processed_inputs[0]["pixel_values"])
188
+ actual = pv.flatten()[:20].tolist()
189
+ expected = golden["image_processing"]["first_20_values"]
190
+
191
+ for i, (a, e) in enumerate(zip(actual, expected)):
192
+ assert abs(a - e) < 1e-5, f"pixel[{i}] changed: {a} vs {e}"
193
+
194
+
195
+ # ---------------------------------------------------------------------------
196
+ # Layer 2: Encoder output
197
+ # ---------------------------------------------------------------------------
198
+ class TestEncoderOutput:
199
+ @pytest.fixture(scope="class")
200
+ def encoder_out(self, env, processed_inputs):
201
+ import torch
202
+ with torch.no_grad():
203
+ out = env["model"].encoder(processed_inputs[0]["pixel_values"])
204
+ return out
205
+
206
+ def test_encoder_output_shape(self, encoder_out):
207
+ # RadioWithNeck outputs (batch, 3201, 1024): patch tokens plus summary token.
208
+ hs = encoder_out.last_hidden_state
209
+ assert hs.shape[0] == 1
210
+ assert hs.shape[2] == 1024, f"Unexpected hidden dim: {hs.shape[2]}"
211
+
212
+ def test_encoder_output_stats_match_golden(self, encoder_out, golden):
213
+ hs = encoder_out.last_hidden_state.float()
214
+ g = golden["encoder_output"]
215
+
216
+ assert abs(hs.mean().item() - g["mean"]) < 0.05, \
217
+ f"encoder mean changed: {hs.mean().item():.4f} vs {g['mean']:.4f}"
218
+ assert abs(hs.std().item() - g["std"]) < 0.05, \
219
+ f"encoder std changed: {hs.std().item():.4f} vs {g['std']:.4f}"
220
+
221
+ def test_encoder_output_slice_match_golden(self, encoder_out, golden):
222
+ """Fixed-position slice: token 0, first 16 hidden dims."""
223
+ hs = encoder_out.last_hidden_state.float()
224
+ actual = hs[0, 0, :16].tolist()
225
+ expected = golden["encoder_output"]["token0_first16"]
226
+
227
+ for i, (a, e) in enumerate(zip(actual, expected)):
228
+ assert abs(a - e) < 0.1, \
229
+ f"encoder hidden[0,0,{i}] changed: {a:.4f} vs {e:.4f}"
230
+
231
+
232
+ # ---------------------------------------------------------------------------
233
+ # Layer 3: Decoder forward pass (logits)
234
+ # ---------------------------------------------------------------------------
235
+ class TestForwardPass:
236
+ @pytest.fixture(scope="class")
237
+ def forward_out(self, env, processed_inputs):
238
+ import torch
239
+ # Minimal decoder input: just the decoder_start_token (EOS = 2 for mBART)
240
+ dec_ids = torch.tensor([[2]], device=env["device"])
241
+ with torch.no_grad():
242
+ out = env["model"](
243
+ pixel_values=processed_inputs[0]["pixel_values"],
244
+ decoder_input_ids=dec_ids,
245
+ return_dict=True,
246
+ )
247
+ return out
248
+
249
+ def test_logits_shape(self, forward_out, env):
250
+ logits = forward_out.logits
251
+ assert logits.shape[0] == 1
252
+ assert logits.shape[1] == 1 # one decoder step
253
+ assert logits.shape[2] == 52352, f"Unexpected vocab size: {logits.shape[2]}"
254
+
255
+ def test_top_k_indices_match_golden(self, forward_out, golden):
256
+ """The TOP_K predicted token IDs should be identical (order matters)."""
257
+ import torch
258
+ logits = forward_out.logits[0, -1, :].float()
259
+ top_k = torch.topk(logits, k=TOP_K)
260
+
261
+ actual = top_k.indices.tolist()
262
+ expected = golden["forward_pass"]["top_k_indices"]
263
+
264
+ assert actual == expected, \
265
+ f"Top-{TOP_K} predicted tokens changed.\n actual: {actual}\n expected: {expected}"
266
+
267
+ def test_top_k_values_match_golden(self, forward_out, golden):
268
+ """Logit magnitudes may drift slightly due to bf16; use a loose tolerance."""
269
+ import torch
270
+ logits = forward_out.logits[0, -1, :].float()
271
+ top_k = torch.topk(logits, k=TOP_K)
272
+
273
+ for i, (a, e) in enumerate(zip(top_k.values.tolist(),
274
+ golden["forward_pass"]["top_k_values"])):
275
+ assert abs(a - e) < 1.0, \
276
+ f"top-{i+1} logit value changed: {a:.3f} vs {e:.3f}"
277
+
278
+
279
+ # ---------------------------------------------------------------------------
280
+ # Layer 4: Generation (greedy, deterministic)
281
+ # ---------------------------------------------------------------------------
282
+ class TestGeneration:
283
+ @pytest.fixture(scope="class")
284
+ def gen_out(self, env, processed_inputs):
285
+ import torch
286
+ with torch.no_grad():
287
+ out = env["model"].generate(
288
+ **processed_inputs[0],
289
+ max_new_tokens=MAX_NEW_TOKENS_GOLDEN,
290
+ do_sample=False,
291
+ num_beams=1,
292
+ )
293
+ return out
294
+
295
+ def test_generated_token_ids_match_golden(self, gen_out, golden):
296
+ """Exact token-ID match - the most sensitive regression signal."""
297
+ actual = gen_out[0].cpu().tolist()
298
+ expected = golden["generation"]["token_ids"]
299
+
300
+ assert actual == expected, (
301
+ f"Generated token sequence differs from golden.\n"
302
+ f" first divergence at index "
303
+ f"{next((i for i,(a,e) in enumerate(zip(actual,expected)) if a!=e), '?')}\n"
304
+ f" actual: {actual}\n"
305
+ f" expected: {expected}"
306
+ )
307
+
308
+ def test_decoded_text_matches_golden(self, gen_out, env, golden):
309
+ text = env["tokenizer"].decode(gen_out[0], skip_special_tokens=False)
310
+ assert text == golden["generation"]["decoded_text"], \
311
+ f"Decoded text differs:\n actual: {text!r}\n expected: {golden['generation']['decoded_text']!r}"
312
+
313
+
314
+ # ---------------------------------------------------------------------------
315
+ # Layer 5: Processor (no model or GPU needed - pure preprocessing & text utils)
316
+ # ---------------------------------------------------------------------------
317
+ @pytest.fixture(scope="session")
318
+ def proc():
319
+ """Load processor + tokenizer only (no model weights, no GPU required)."""
320
+ from transformers import AutoProcessor, AutoTokenizer
321
+ tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH, trust_remote_code=True)
322
+ processor = AutoProcessor.from_pretrained(MODEL_PATH, trust_remote_code=True)
323
+ return dict(processor=processor, tokenizer=tokenizer)
324
+
325
+
326
+ class TestProcessor:
327
+ # ------------------------------------------------------------------
328
+ # post_process_generation
329
+ # ------------------------------------------------------------------
330
+ def test_post_process_generation_returns_string_for_string_input(self, proc, golden):
331
+ """String input -> string output."""
332
+ decoded = golden["generation"]["decoded_text"]
333
+ result = proc["processor"].post_process_generation(decoded)
334
+ assert isinstance(result, str)
335
+
336
+ def test_post_process_generation_removes_bos_eos(self, proc, golden):
337
+ """<s> and </s> tokens must be stripped from the output."""
338
+ decoded = golden["generation"]["decoded_text"]
339
+ result = proc["processor"].post_process_generation(decoded)
340
+ assert "<s>" not in result
341
+ assert "</s>" not in result
342
+
343
+ def test_post_process_generation_matches_manual_clean(self, proc, golden):
344
+ """Exact match against the expected cleaned string."""
345
+ decoded = golden["generation"]["decoded_text"]
346
+ expected = decoded.replace("<s>", "").replace("</s>", "").strip()
347
+ result = proc["processor"].post_process_generation(decoded)
348
+ assert result == expected
349
+
350
+ def test_post_process_generation_list_returns_list(self, proc, golden):
351
+ """Multi-element list input -> list output of the same length."""
352
+ decoded = golden["generation"]["decoded_text"]
353
+ result = proc["processor"].post_process_generation([decoded, decoded])
354
+ assert isinstance(result, list)
355
+ assert len(result) == 2
356
+ assert result[0] == result[1]
357
+
358
+ def test_post_process_generation_single_element_list_returns_string(self, proc, golden):
359
+ """Single-element list input -> scalar string (not a list)."""
360
+ decoded = golden["generation"]["decoded_text"]
361
+ result = proc["processor"].post_process_generation([decoded])
362
+ assert isinstance(result, str)
363
+
364
+ # ------------------------------------------------------------------
365
+ # decode / batch_decode via the processor
366
+ # ------------------------------------------------------------------
367
+ def test_decode_via_processor_matches_tokenizer(self, proc, golden):
368
+ """processor.decode() must give the same result as tokenizer.decode()."""
369
+ token_ids = golden["generation"]["token_ids"]
370
+ via_proc = proc["processor"].decode(token_ids, skip_special_tokens=False)
371
+ via_tok = proc["tokenizer"].decode(token_ids, skip_special_tokens=False)
372
+ assert via_proc == via_tok
373
+
374
+ def test_batch_decode_via_processor(self, proc, golden):
375
+ """processor.batch_decode() on repeated token lists matches golden decoded text."""
376
+ token_ids = golden["generation"]["token_ids"]
377
+ results = proc["processor"].batch_decode(
378
+ [token_ids, token_ids], skip_special_tokens=False
379
+ )
380
+ assert isinstance(results, list)
381
+ assert len(results) == 2
382
+ assert results[0] == results[1] == golden["generation"]["decoded_text"]
383
+
384
+ # ------------------------------------------------------------------
385
+ # Image processing edge cases
386
+ # ------------------------------------------------------------------
387
+ def test_large_image_resized_to_target(self, proc):
388
+ """Image larger than 2048x1648 is downscaled to exactly [1, 3, 2048, 1648]."""
389
+ from PIL import Image
390
+ large = Image.new("RGB", (4000, 5000), color=(128, 64, 32))
391
+ out = proc["processor"](images=[large], return_tensors="pt")
392
+ assert list(out["pixel_values"].shape) == [1, 3, 2048, 1648]
393
+
394
+ def test_grayscale_image_converted_to_rgb(self, proc):
395
+ """Grayscale (mode 'L') image is converted to RGB and produces 3 output channels."""
396
+ from PIL import Image
397
+ gray = Image.new("L", (400, 600), color=128)
398
+ out = proc["processor"](images=[gray], return_tensors="pt")
399
+ assert list(out["pixel_values"].shape) == [1, 3, 2048, 1648]
400
+
401
+ def test_multi_image_batch_first_dim(self, proc):
402
+ """A batch of N images produces pixel_values with first dimension N."""
403
+ from PIL import Image
404
+ imgs = [
405
+ Image.new("RGB", (400, 600), color=(i * 30, i * 20, i * 10))
406
+ for i in range(3)
407
+ ]
408
+ out = proc["processor"](images=imgs, return_tensors="pt")
409
+ assert list(out["pixel_values"].shape) == [3, 3, 2048, 1648]
410
+
411
+ def test_image_only_input_has_no_input_ids(self, proc):
412
+ """Passing images without text returns pixel_values and no input_ids key."""
413
+ from PIL import Image
414
+ img = Image.new("RGB", (400, 600))
415
+ out = proc["processor"](images=[img], return_tensors="pt")
416
+ assert "pixel_values" in out
417
+ assert "input_ids" not in out
418
+
419
+ def test_text_only_input_has_no_pixel_values(self, proc):
420
+ """Passing text without images returns input_ids and no pixel_values key."""
421
+ out = proc["processor"](text="hello world", return_tensors="pt")
422
+ assert "input_ids" in out
423
+ assert "pixel_values" not in out
424
+
425
+
426
+ # ---------------------------------------------------------------------------
427
+ # Capture helper (run as script: python test_golden.py --capture)
428
+ # ---------------------------------------------------------------------------
429
+ def capture(model_path: str = MODEL_PATH):
430
+ """
431
+ Run a full inference pass and write golden_outputs.json.
432
+ Intended to be run once against the pinned dependency set.
433
+ """
434
+ import torch
435
+ import transformers
436
+ from transformers import AutoModel, AutoProcessor, AutoTokenizer
437
+
438
+ device = "cuda:0" if torch.cuda.is_available() else "cpu"
439
+ dtype = torch.bfloat16 if torch.cuda.is_available() else torch.float32
440
+
441
+ print(f"Capturing golden outputs")
442
+ print(f" transformers : {transformers.__version__}")
443
+ print(f" torch : {torch.__version__}")
444
+ print(f" device : {device} dtype={dtype}")
445
+ print(f" model_path : {model_path}")
446
+
447
+ model = AutoModel.from_pretrained(
448
+ model_path, trust_remote_code=True, torch_dtype=dtype
449
+ ).to(device).eval()
450
+ tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
451
+ processor = AutoProcessor.from_pretrained(model_path, trust_remote_code=True)
452
+
453
+ image = make_test_image()
454
+ inputs = processor(
455
+ images=[image],
456
+ text=TASK_PROMPT,
457
+ return_tensors="pt",
458
+ add_special_tokens=False,
459
+ ).to(device)
460
+
461
+ # ---------- image processing ----------
462
+ pv = inputs["pixel_values"].float()
463
+ image_data = {
464
+ "shape": list(pv.shape),
465
+ "mean": pv.mean().item(),
466
+ "std": pv.std().item(),
467
+ "first_20_values": pv.flatten()[:20].tolist(),
468
+ }
469
+ print(f"\n[image] shape={image_data['shape']} mean={image_data['mean']:.4f} std={image_data['std']:.4f}")
470
+
471
+ # ---------- encoder output ----------
472
+ with torch.no_grad():
473
+ enc_out = model.encoder(inputs["pixel_values"])
474
+ hs = enc_out.last_hidden_state.float()
475
+ encoder_data = {
476
+ "shape": list(hs.shape),
477
+ "mean": hs.mean().item(),
478
+ "std": hs.std().item(),
479
+ "token0_first16": hs[0, 0, :16].tolist(),
480
+ }
481
+ print(f"[encoder] shape={encoder_data['shape']} mean={encoder_data['mean']:.4f} std={encoder_data['std']:.4f}")
482
+
483
+ # ---------- forward pass (logits) ----------
484
+ dec_ids = torch.tensor([[2]], device=device) # decoder_start_token_id
485
+ with torch.no_grad():
486
+ fwd_out = model(
487
+ pixel_values=inputs["pixel_values"],
488
+ decoder_input_ids=dec_ids,
489
+ return_dict=True,
490
+ )
491
+ logits = fwd_out.logits[0, -1, :].float()
492
+ top_k = torch.topk(logits, k=TOP_K)
493
+ forward_data = {
494
+ "logits_shape": list(fwd_out.logits.shape),
495
+ "top_k_indices": top_k.indices.tolist(),
496
+ "top_k_values": top_k.values.tolist(),
497
+ }
498
+ top_tokens = [tokenizer.decode([i]) for i in top_k.indices.tolist()]
499
+ print(f"[forward] top-{TOP_K} tokens: {top_tokens}")
500
+
501
+ # ---------- generation ----------
502
+ with torch.no_grad():
503
+ gen_out = model.generate(
504
+ **inputs,
505
+ max_new_tokens=MAX_NEW_TOKENS_GOLDEN,
506
+ do_sample=False,
507
+ num_beams=1,
508
+ )
509
+ token_ids = gen_out[0].cpu().tolist()
510
+ decoded_text = tokenizer.decode(gen_out[0], skip_special_tokens=False)
511
+ generation_data = {
512
+ "max_new_tokens": MAX_NEW_TOKENS_GOLDEN,
513
+ "token_ids": token_ids,
514
+ "decoded_text": decoded_text,
515
+ }
516
+ print(f"[generation] {len(token_ids)} tokens: {decoded_text!r}")
517
+
518
+ # ---------- save ----------
519
+ golden = {
520
+ "metadata": {
521
+ "transformers_version": transformers.__version__,
522
+ "torch_version": torch.__version__,
523
+ "device": str(device),
524
+ "dtype": str(dtype),
525
+ "model_path": model_path,
526
+ },
527
+ "image_processing": image_data,
528
+ "encoder_output": encoder_data,
529
+ "forward_pass": forward_data,
530
+ "generation": generation_data,
531
+ }
532
+ save_golden(golden)
533
+ print(f"\nGolden outputs written to {GOLDEN_FILE}")
534
+ return golden
535
+
536
+
537
+ if __name__ == "__main__":
538
+ import argparse
539
+
540
+ parser = argparse.ArgumentParser(description="Golden reference capture/verify for Nemotron-Parse")
541
+ parser.add_argument("--capture", action="store_true", help="Capture golden outputs")
542
+ parser.add_argument("--model-path", default=MODEL_PATH, help="Path to model directory")
543
+ args = parser.parse_args()
544
+
545
+ if args.capture:
546
+ capture(model_path=args.model_path)
547
+ else:
548
+ parser.print_help()
549
+ print("\nTo run tests: pytest test_golden.py -v")
550
+ print("To capture: python test_golden.py --capture")
test_vllm_golden.py ADDED
@@ -0,0 +1,103 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """
2
+ Golden reference test for NVIDIA-Nemotron-Parse-v1.1 through vLLM.
3
+
4
+ This mirrors the generation layer in test_golden.py, but exercises the vLLM
5
+ encoder/decoder interface. vLLM returns completion text, while Transformers
6
+ stores the full decoded decoder sequence in golden_outputs.json, so the
7
+ comparison normalizes vLLM output back to the full decoded form.
8
+ """
9
+
10
+ from __future__ import annotations
11
+
12
+ import json
13
+
14
+ import pytest
15
+
16
+ pytest.importorskip("vllm", reason="vLLM is not installed")
17
+
18
+ from test_golden import (
19
+ GOLDEN_FILE,
20
+ MAX_NEW_TOKENS_GOLDEN,
21
+ MODEL_PATH,
22
+ TASK_PROMPT,
23
+ make_test_image,
24
+ )
25
+
26
+
27
+ def _load_golden_generation() -> str:
28
+ if not GOLDEN_FILE.exists():
29
+ pytest.skip("golden_outputs.json not found - run: python test_golden.py --capture")
30
+ with open(GOLDEN_FILE) as f:
31
+ return json.load(f)["generation"]["decoded_text"]
32
+
33
+
34
+ def _as_full_decoded_text(vllm_text: str, expected_full_text: str) -> str:
35
+ if vllm_text.startswith(TASK_PROMPT):
36
+ full_text = vllm_text
37
+ elif expected_full_text.startswith(TASK_PROMPT):
38
+ full_text = TASK_PROMPT + vllm_text
39
+ else:
40
+ full_text = vllm_text
41
+
42
+ # vLLM stops on EOS but does not include that stop token in completion text.
43
+ if expected_full_text.endswith("</s>") and not full_text.endswith("</s>"):
44
+ with_eos = full_text + "</s>"
45
+ if with_eos == expected_full_text:
46
+ return with_eos
47
+
48
+ return full_text
49
+
50
+
51
+ def test_vllm_generation_matches_golden():
52
+ from vllm import LLM, SamplingParams
53
+
54
+ expected = _load_golden_generation()
55
+ image = make_test_image()
56
+
57
+ sampling_params = SamplingParams(
58
+ temperature=0.0,
59
+ top_k=1,
60
+ repetition_penalty=1.1,
61
+ max_tokens=MAX_NEW_TOKENS_GOLDEN,
62
+ skip_special_tokens=False,
63
+ )
64
+ llm = LLM(
65
+ model=MODEL_PATH,
66
+ max_num_seqs=2,
67
+ limit_mm_per_prompt={"image": 1},
68
+ dtype="bfloat16",
69
+ trust_remote_code=True,
70
+ attention_config={"backend": "TRITON_ATTN"},
71
+ )
72
+
73
+ request_names = ["implicit", "explicit_encoder_decoder"]
74
+ requests = [
75
+ {
76
+ "prompt": TASK_PROMPT,
77
+ "multi_modal_data": {"image": image.copy()},
78
+ },
79
+ {
80
+ "encoder_prompt": {
81
+ "prompt": "",
82
+ "multi_modal_data": {"image": image.copy()},
83
+ },
84
+ "decoder_prompt": TASK_PROMPT,
85
+ },
86
+ ]
87
+
88
+ outputs = llm.generate(requests, sampling_params)
89
+ assert len(outputs) == len(request_names)
90
+
91
+ mismatches = []
92
+ for name, output in zip(request_names, outputs):
93
+ text = output.outputs[0].text
94
+ full_text = _as_full_decoded_text(text, expected)
95
+ if full_text != expected:
96
+ mismatches.append(
97
+ f"vLLM {name} generation differs from golden.\n"
98
+ f" raw vLLM text: {text!r}\n"
99
+ f" normalized: {full_text!r}\n"
100
+ f" expected: {expected!r}"
101
+ )
102
+
103
+ assert not mismatches, "\n\n".join(mismatches)
uv.lock ADDED
The diff for this file is too large to render. See raw diff