Image-Text-to-Text
Transformers
Safetensors
nemotron_parse
feature-extraction
nvidia
VLM
OCR
conversational
custom_code
Instructions to use nvidia/NVIDIA-Nemotron-Parse-v1.1 with libraries, inference providers, notebooks, and local apps. Follow these links to get started.
- Libraries
- Transformers
How to use nvidia/NVIDIA-Nemotron-Parse-v1.1 with Transformers:
# Use a pipeline as a high-level helper from transformers import pipeline pipe = pipeline("image-text-to-text", model="nvidia/NVIDIA-Nemotron-Parse-v1.1", trust_remote_code=True) messages = [ { "role": "user", "content": [ {"type": "image", "url": "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/p-blog/candy.JPG"}, {"type": "text", "text": "What animal is on the candy?"} ] }, ] pipe(text=messages)# Load model directly from transformers import AutoModel model = AutoModel.from_pretrained("nvidia/NVIDIA-Nemotron-Parse-v1.1", trust_remote_code=True, device_map="auto") - Notebooks
- Google Colab
- Kaggle
- Local Apps Settings
- vLLM
How to use nvidia/NVIDIA-Nemotron-Parse-v1.1 with vLLM:
Install from pip and serve model
# Install vLLM from pip: pip install vllm # Start the vLLM server: vllm serve "nvidia/NVIDIA-Nemotron-Parse-v1.1" # Call the server using curl (OpenAI-compatible API): curl -X POST "http://localhost:8000/v1/chat/completions" \ -H "Content-Type: application/json" \ --data '{ "model": "nvidia/NVIDIA-Nemotron-Parse-v1.1", "messages": [ { "role": "user", "content": [ { "type": "text", "text": "Describe this image in one sentence." }, { "type": "image_url", "image_url": { "url": "https://cdn.britannica.com/61/93061-050-99147DCE/Statue-of-Liberty-Island-New-York-Bay.jpg" } } ] } ] }'Use Docker
docker model run hf.co/nvidia/NVIDIA-Nemotron-Parse-v1.1
- SGLang
How to use nvidia/NVIDIA-Nemotron-Parse-v1.1 with SGLang:
Install from pip and serve model
# Install SGLang from pip: pip install sglang # Start the SGLang server: python3 -m sglang.launch_server \ --model-path "nvidia/NVIDIA-Nemotron-Parse-v1.1" \ --host 0.0.0.0 \ --port 30000 # Call the server using curl (OpenAI-compatible API): curl -X POST "http://localhost:30000/v1/chat/completions" \ -H "Content-Type: application/json" \ --data '{ "model": "nvidia/NVIDIA-Nemotron-Parse-v1.1", "messages": [ { "role": "user", "content": [ { "type": "text", "text": "Describe this image in one sentence." }, { "type": "image_url", "image_url": { "url": "https://cdn.britannica.com/61/93061-050-99147DCE/Statue-of-Liberty-Island-New-York-Bay.jpg" } } ] } ] }'Use Docker images
docker run --gpus all \ --shm-size 32g \ -p 30000:30000 \ -v ~/.cache/huggingface:/root/.cache/huggingface \ --env "HF_TOKEN=<secret>" \ --ipc=host \ lmsysorg/sglang:latest \ python3 -m sglang.launch_server \ --model-path "nvidia/NVIDIA-Nemotron-Parse-v1.1" \ --host 0.0.0.0 \ --port 30000 # Call the server using curl (OpenAI-compatible API): curl -X POST "http://localhost:30000/v1/chat/completions" \ -H "Content-Type: application/json" \ --data '{ "model": "nvidia/NVIDIA-Nemotron-Parse-v1.1", "messages": [ { "role": "user", "content": [ { "type": "text", "text": "Describe this image in one sentence." }, { "type": "image_url", "image_url": { "url": "https://cdn.britannica.com/61/93061-050-99147DCE/Statue-of-Liberty-Island-New-York-Bay.jpg" } } ] } ] }' - Docker Model Runner
How to use nvidia/NVIDIA-Nemotron-Parse-v1.1 with Docker Model Runner:
docker model run hf.co/nvidia/NVIDIA-Nemotron-Parse-v1.1
Add golden tests and latest runtime support for Nemotron Parse v1.1
#8
by oliverholworthy - opened
- Dockerfile +39 -0
- README.md +5 -1
- config.json +3 -2
- docker-compose.yaml +33 -0
- golden_outputs.json +133 -0
- hf_nemotron_parse_modeling.py +152 -34
- hf_nemotron_parse_processor.py +26 -6
- preprocessor_config.json +11 -1
- pyproject.toml +33 -0
- test_golden.py +550 -0
- test_vllm_golden.py +103 -0
- uv.lock +0 -0
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 |
-
"
|
| 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":
|
| 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 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
| 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
|
| 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 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 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 |
-
|
| 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 |
-
|
| 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 |
-
|
| 273 |
-
|
| 274 |
-
|
| 275 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 276 |
|
| 277 |
-
|
| 278 |
-
|
|
|
|
|
|
|
|
|
|
| 279 |
|
| 280 |
-
if
|
| 281 |
-
|
|
|
|
|
|
|
| 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": "
|
| 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 |
-
) ->
|
| 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
|
| 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":
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 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
|
|
|