mranzinger/fix_misaligned_address

Dockerfile

@@ -8,11 +8,11 @@ ENV TORCH_CUDA_ARCH_LIST=${TORCH_CUDA_ARCH_LIST}
 RUN --mount=type=cache,target=/root/.cache/pip \
     pip install -U pip hatchling "setuptools>=68" --root-user-action ignore
 
-COPY nemotron-ocr /workspace/nemotron-ocr
-WORKDIR /workspace/nemotron-ocr
+COPY nemo-retriever-ocr /workspace/nemo-retriever-ocr
+WORKDIR /workspace/nemo-retriever-ocr
 
 # Ensure no prebuilt binaries/artifacts from the host are present
-RUN rm -f src/nemotron_ocr_cpp/*.so || true \
+RUN rm -f src/nemo_retriever_ocr_cpp/*.so || true \
     && rm -rf build/ dist/
 
 RUN --mount=type=cache,target=/root/.cache/pip \

README.md

@@ -16,7 +16,7 @@ tags:
 - ingestion
 ---
 
-# Nemotron OCR v1
+# NeMo Retriever OCR v1
 
 ## **Model Overview**
 
@@ -27,15 +27,15 @@ tags:
 
 ### **Description**
 
-The Nemotron OCR v1 model is a state-of-the-art text recognition model designed for robust end-to-end optical character recognition (OCR) on complex real-world images. It integrates three core neural network modules: a detector for text region localization, a recognizer for transcription of detected regions, and a relational model for layout and structure analysis.
+The NeMo Retriever OCR v1 model is a state-of-the-art text recognition model designed for robust end-to-end optical character recognition (OCR) on complex real-world images. It integrates three core neural network modules: a detector for text region localization, a recognizer for transcription of detected regions, and a relational model for layout and structure analysis.
 
-This model is optimized for a wide variety of OCR tasks, including multi-line, multi-block, and natural scene text, and it supports advanced reading order analysis via its relational model component. Nemotron OCR v1 has been developed to be production-ready and commercially usable, with a focus on speed and accuracy on both document and natural scene images.
+This model is optimized for a wide variety of OCR tasks, including multi-line, multi-block, and natural scene text, and it supports advanced reading order analysis via its relational model component. NeMo Retriever OCR v1 has been developed to be production-ready and commercially usable, with a focus on speed and accuracy on both document and natural scene images.
 
-The Nemotron OCR v1 model is part of the NVIDIA NeMo Retriever collection of NIM microservices, which provides state-of-the-art, commercially-ready models and microservices optimized for the lowest latency and highest throughput. It features a production-ready information retrieval pipeline with enterprise support. The models that form the core of this solution have been trained using responsibly selected, auditable data sources. With multiple pre-trained models available as starting points, developers can readily customize them for domain-specific use cases, such as information technology, human resource help assistants, and research & development research assistants.
+The NeMo Retriever OCR v1 model is part of the NVIDIA NeMo Retriever collection of NIM microservices, which provides state-of-the-art, commercially-ready models and microservices optimized for the lowest latency and highest throughput. It features a production-ready information retrieval pipeline with enterprise support. The models that form the core of this solution have been trained using responsibly selected, auditable data sources. With multiple pre-trained models available as starting points, developers can readily customize them for domain-specific use cases, such as information technology, human resource help assistants, and research & development research assistants.
 
 This model is ready for commercial use.
 
-We are excited to announce the open sourcing of this commercial model. For users interested in deploying this model in production environments, it is also available via the model API in NVIDIA Inference Microservices (NIM) at [nemotron-ocr-v1](https://build.nvidia.com/nvidia/nemoretriever-ocr-v1).
+We are excited to announce the open sourcing of this commercial model. For users interested in deploying this model in production environments, it is also available via the model API in NVIDIA Inference Microservices (NIM) at [nemoretriever-ocr-v1](https://build.nvidia.com/nvidia/nemoretriever-ocr-v1).
 
 ### **License/Terms of use**
 
@@ -57,11 +57,11 @@ Global
 
 ### Use Case
 
-The **Nemotron OCR v1** model is designed for high-accuracy and high-speed extraction of textual information from images, making it ideal for powering multimodal retrieval systems, Retrieval-Augmented Generation (RAG) pipelines, and agentic applications that require seamless integration of visual and language understanding. Its robust performance and efficiency make it an excellent choice for next-generation AI systems that demand both precision and scalability across diverse real-world content.
+The **NeMo Retriever OCR v1** model is designed for high-accuracy and high-speed extraction of textual information from images, making it ideal for powering multimodal retrieval systems, Retrieval-Augmented Generation (RAG) pipelines, and agentic applications that require seamless integration of visual and language understanding. Its robust performance and efficiency make it an excellent choice for next-generation AI systems that demand both precision and scalability across diverse real-world content.
 
 ### Release Date
 
-10/23/2025 via https://huggingface.co/nvidia/nemotron-ocr-v1
+10/23/2025 via https://huggingface.co/nvidia/nemoretriever-ocr-v1
 
 ### References
 
@@ -71,7 +71,7 @@ The **Nemotron OCR v1** model is designed for high-accuracy and high-speed extra
 
 **Architecture Type:** Hybrid detector–recognizer with document-level relational modeling
 
-The Nemotron OCR v1 model integrates three specialized neural components:
+The NeMo Retriever OCR v1 model integrates three specialized neural components:
 
 - **Text Detector:** Utilizes a RegNetY-8GF convolutional backbone for high-accuracy localization of text regions within images.
 - **Text Recognizer:** Employs a Transformer-based sequence recognizer to transcribe text from detected regions, supporting variable word and line lengths.
@@ -163,11 +163,11 @@ git lfs install
 ```
 - Using https
 ```
-git clone https://huggingface.co/nvidia/nemotron-ocr-v1
+git clone https://huggingface.co/nvidia/nemoretriever-ocr-v1
 ```
 - Or using ssh
 ```
-git clone git@hf.co:nvidia/nemotron-ocr-v1
+git clone git@hf.co:nvidia/nemoretriever-ocr-v1
 ```
 
 2. Installation
@@ -179,7 +179,7 @@ git clone git@hf.co:nvidia/nemotron-ocr-v1
 - Run the following command to install the package:
 
 ```bash
-cd nemotron-ocr
+cd nemo-retriever-ocr
 pip install hatchling
 pip install -v .
 ```
@@ -197,7 +197,7 @@ docker run --rm --gpus all nvcr.io/nvidia/pytorch:25.09-py3 nvidia-smi
 - From the repo root, bring up the service to run the example against the provided image `ocr-example-image.png`:
 
 ```bash
-docker compose run --rm nemotron-ocr \
+docker compose run --rm nemo-retriever-ocr \
   bash -lc "python example.py ocr-example-input-1.png --merge-level paragraph"
 ```
 
@@ -212,9 +212,9 @@ Output is saved next to your input image as `<name>-annotated.<ext>` on the host
 3. Run the model using the following code:
 
 ```python
-from nemotron_ocr.inference.pipeline import NemotronOCR
+from nemo_retriever_ocr.inference.pipeline import NemoRetrieverOCR
 
-ocr = NemotronOCR()
+ocr = NemoRetrieverOCR()
 
 predictions = ocr("ocr-example-input-1.png")
 
@@ -230,7 +230,7 @@ for pred in predictions:
 ### Software Integration
 
 **Runtime Engine(s):**
-- **NeMo Nemotron OCR V1** NIM
+- **NeMo Retriever Page Elements v3** NIM
 
 
 **Supported Hardware Microarchitecture Compatibility [List in Alphabetic Order]:**
@@ -247,7 +247,7 @@ This AI model can be embedded as an Application Programming Interface (API) call
 
 ## Model Version(s):
 
-* `nemotron-ocr-v1`
+* `nemoretriever-ocr-v1`
 
 ## **Training and Evaluation Datasets:**
 
@@ -267,7 +267,7 @@ The model is trained on a large-scale, curated mix of public and proprietary OCR
 
 ### **Evaluation Datasets**
 
-The Nemotron OCR v1 model is evaluated on several NVIDIA internal datasets for various tasks, such as pure OCR, table content extraction, and document retrieval.
+The NeMo Retriever OCR v1 model is evaluated on several NVIDIA internal datasets for various tasks, such as pure OCR, table content extraction, and document retrieval.
 
 **Data Collection Method:** Hybrid (Automated, Human, Synthetic)<br>
 **Labeling Method:** Hybrid (Automated, Human, Synthetic)<br>
@@ -275,9 +275,9 @@ The Nemotron OCR v1 model is evaluated on several NVIDIA internal datasets for v
 
 ### **Evaluation Results**
 
-We benchmarked Nemotron OCR v1 on internal evaluation datasets against PaddleOCR on various tasks, such as pure OCR (Character Error Rate), table content extraction (TEDS), and document retrieval (Recall@5).
+We benchmarked NeMo Retriever OCR v1 on internal evaluation datasets against PaddleOCR on various tasks, such as pure OCR (Character Error Rate), table content extraction (TEDS), and document retrieval (Recall@5).
 
-| Metric                                   | Nemotron OCR v1 | PaddleOCR | Net change |
+| Metric                                   | NeMo Retriever OCR v1 | PaddleOCR | Net change |
 |-------------------------------------------|--------------------|-----------|-----------------|
 | Character Error Rate                      | 0.1633             | 0.2029    | -19.5% ✔️         |
 | Bag-of-character Error Rate               | 0.0453             | 0.0512    | -11.5% ✔️         |

THIRD_PARTY_NOTICES.md

@@ -1,519 +1,15 @@
-Copyright "Angus Johnson" - Boost Software License 1.0
-License Text([https://sourceforge.net/p/polyclipping/code/HEAD/tree/tags/6.2.0/License.txt](https://sourceforge.net/p/polyclipping/code/HEAD/tree/tags/6.2.0/License.txt))
-
-This notice applies to **clipper**.
-
-Copyright (c) 2010-2014 Angus Johnson
-
-Boost Software License - Version 1.0 - August 17th, 2003
-
-Permission is hereby granted, free of charge, to any person or organization
-obtaining a copy of the software and accompanying documentation covered by
-this license (the "Software") to use, reproduce, display, distribute,
-execute, and transmit the Software, and to prepare derivative works of the
-Software, and to permit third-parties to whom the Software is furnished to
-do so, all subject to the following:
-
-The copyright notices in the Software and this entire statement, including
-the above license grant, this restriction and the following disclaimer,
-must be included in all copies of the Software, in whole or in part, and
-all derivative works of the Software, unless such copies or derivative
-works are solely in the form of machine-executable object code generated by
-a source language processor.
-
-THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
-IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
-FITNESS FOR A PARTICULAR PURPOSE, TITLE AND NON-INFRINGEMENT. IN NO EVENT
-SHALL THE COPYRIGHT HOLDERS OR ANYONE DISTRIBUTING THE SOFTWARE BE LIABLE
-FOR ANY DAMAGES OR OTHER LIABILITY, WHETHER IN CONTRACT, TORT OR OTHERWISE,
-ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
-DEALINGS IN THE SOFTWARE.
-
------
-
-Copyright "Ofek Lev" - MIT License
-License Text([https://github.com/pypa/hatch/blob/master/LICENSE.txt](https://github.com/pypa/hatch/blob/master/LICENSE.txt))
-
-This notice applies to **hatchling**.
-
-Copyright (c) 2017-present Ofek Lev <ofekmeister@gmail.com>
-
-Permission is hereby granted, free of charge, to any person obtaining a copy
-of this software and associated documentation files (the "Software"), to deal
-in the Software without restriction, including without limitation the rights
-to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
-copies of the Software, and to permit persons to whom the Software is
-furnished to do so, subject to the following conditions:
-
-The above copyright notice and this permission notice shall be included in all
-copies or substantial portions of the Software.
-
-THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
-IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
-FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
-AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
-LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
-OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
-SOFTWARE.
-
------
-
-Copyright "NumPy Developers" - BSD 3-Clause License
-License Text([https://github.com/numpy/numpy/blob/main/LICENSE.txt](https://github.com/numpy/numpy/blob/main/LICENSE.txt))
-
-This notice applies to **numpy**.
-
-Copyright (c) 2005-2023, NumPy Developers.
-All rights reserved.
-
-Redistribution and use in source and binary forms, with or without
-modification, are permitted provided that the following conditions are
-met:
-
-  * Redistributions of source code must retain the above copyright
-    notice, this list of conditions and the following disclaimer.
-
-  * Redistributions in binary form must reproduce the above
-    copyright notice, this list of conditions and the following
-    disclaimer in the documentation and/or other materials provided
-    with the distribution.
-
-  * Neither the name of the NumPy Developers nor the names of any
-    contributors may be used to endorse or promote products derived
-    from this software without specific prior written permission.
-
-
-THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
-"AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
-LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
-A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
-OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
-SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
-LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
-DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
-THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
-(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
-OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
-
------
-
-Copyright "pandas Developers" - BSD 3-Clause License
-License Text([https://github.com/pandas-dev/pandas/blob/main/LICENSE](https://github.com/pandas-dev/pandas/blob/main/LICENSE))
-
-This notice applies to **pandas**.
-
-Copyright (c) 2008-2011, AQR Capital Management, LLC, Lambda Foundry, Inc. and PyData Development Team
-All rights reserved.
-
-Copyright (c) 2011-2023, The PyData Development Team
-All rights reserved.
-
-Redistribution and use in source and binary forms, with or without
-modification, are permitted provided that the following conditions are
-met:
-
-
-  * Redistributions of source code must retain the above copyright
-    notice, this list of conditions and the following disclaimer.
-
-  * Redistributions in binary form must reproduce the above
-    copyright notice, this list of conditions and the following
-    disclaimer in the documentation and/or other materials provided
-    with the distribution.
-
-  * Neither the name of the pandas development team nor the names of
-    any contributors may be used to endorse or promote products
-    derived from this software without specific prior written
-    permission.
-
-
-THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
-"AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
-LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
-A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
-OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
-SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
-LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
-DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
-THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
-(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
-OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
-
------
-
-Copyright "Secret Labs AB, Fredrik Lundh, Alex Clark and contributors" - Pillow License
-License Text([https://github.com/python-pillow/Pillow/blob/main/LICENSE](https://github.com/python-pillow/Pillow/blob/main/LICENSE))
-
-This notice applies to **PIL (Pillow)**.
-
-The Python Imaging Library (PIL) is
-Copyright (c) 1997-2011 by Secret Labs AB
-Copyright (c) 1995-2011 by Fredrik Lundh
-Copyright (c) 2010-2023 by Alex Clark and contributors
-
-Like PIL, Pillow is licensed under the open source HPND License:
-
-By obtaining, using, and/or copying this software and/or its
-associated documentation, you agree that you have read, understood,
-and will comply with the following terms and conditions:
-
-Permission to use, copy, modify, and distribute this software and
-its associated documentation for any purpose and without fee is
-hereby granted, provided that the above copyright notice appears in
-all copies, and that both that copyright notice and this permission
-notice appear in supporting documentation, and that the name of
-Secret Labs AB or the author not be used in advertising or publicity
-pertaining to distribution of the software without specific, written
-prior permission.
-
-SECRET LABS AB AND THE AUTHOR DISCLAIMS ALL WARRANTIES WITH REGARD
-TO THIS SOFTWARE, INCLUDING ALL IMPLIED WARRANTIES OF MERCHANT-
-ABILITY AND FITNESS. IN NO EVENT SHALL SECRET LABS AB OR THE AUTHOR
-BE LIABLE FOR ANY SPECIAL, INDIRECT OR CONSEQUENTIAL DAMAGES OR ANY
-DAMAGES WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS,
-WHETHER IN AN ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS
-ACTION, ARISING OUT OF OR IN CONNECTION WITH THE USE OR PERFORMANCE
-OF THIS SOFTWARE.
-
------
-
-Copyright "The scikit-learn developers" - BSD 3-Clause License
-License Text([https://github.com/scikit-learn/scikit-learn/blob/main/COPYING](https://github.com/scikit-learn/scikit-learn/blob/main/COPYING))
-
-This notice applies to **scikit-learn**.
-
-Copyright (c) 2007-2024 The scikit-learn developers.
-All rights reserved.
-
-Redistribution and use in source and binary forms, with or without
-modification, are permitted provided that the following conditions are
-met:
-
-
-  * Redistributions of source code must retain the above copyright
-    notice, this list of conditions and the following disclaimer.
-
-  * Redistributions in binary form must reproduce the above
-    copyright notice, this list of conditions and the following
-    disclaimer in the documentation and/or other materials provided
-    with the distribution.
-
-  * Neither the name of the scikit-learn developers nor the names of
-    any contributors may be used to endorse or promote products
-    derived from this software without specific prior written
-    permission.
-
-
-THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
-"AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
-LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
-A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
-OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
-SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
-LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
-DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
-THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
-(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
-OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
-
------
-
-Copyright "Jason R. Coombs" - MIT License
-License Text([https://github.com/pypa/setuptools/blob/main/LICENSE](https://github.com/pypa/setuptools/blob/main/LICENSE))
-
-This notice applies to **setuptools**.
-
-Copyright (c) 2016 Jason R. Coombs <jaraco@jaraco.com>
-
-Permission is hereby granted, free of charge, to any person obtaining a copy
-of this software and associated documentation files (the "Software"), to deal
-in the Software without restriction, including without limitation the rights
-to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
-copies of the Software, and to permit persons to whom the Software is
-furnished to do so, subject to the following conditions:
-
-The above copyright notice and this permission notice shall be included in all
-copies or substantial portions of the Software.
-
-THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
-IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
-FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
-AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
-LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
-OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
-SOFTWARE.
-
------
-
-Copyright "Sean Gillies" - BSD 3-Clause License
-License Text([https://github.com/shapely/shapely/blob/main/LICENSE.txt](https://github.com/shapely/shapely/blob/main/LICENSE.txt))
-
-This notice applies to **Shapely**.
-
-Copyright (c) 2007, Sean Gillies.
-All rights reserved.
-
-Redistribution and use in source and binary forms, with or without
-modification, are permitted provided that the following conditions are met:
-
-
-  * Redistributions of source code must retain the above copyright
-    notice, this list of conditions and the following disclaimer.
-
-  * Redistributions in binary form must reproduce the above copyright
-    notice, this list of conditions and the following disclaimer in the
-    documentation and/or other materials provided with the distribution.
-
-  * Neither the name of Sean Gillies nor the names of
-    its contributors may be used to endorse or promote products derived from
-    this software without specific prior written permission.
-
-
-THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
-AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
-IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
-ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE
-LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
-CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
-SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
-INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
-CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
-ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
-POSSIBILITY OF SUCH DAMAGE.
-
------
-
-Copyright "PyTorch Contributors" - BSD-style License
-License Text([https://github.com/pytorch/pytorch/blob/main/LICENSE](https://github.com/pytorch/pytorch/blob/main/LICENSE))
-
-This notice applies to **torch** and **torchvision**.
-
-Copyright (c) 2016-     Facebook, Inc. (Adam Paszke)
-Copyright (c) 2014-     Facebook, Inc. (Soumith Chintala)
-Copyright (c) 2011-2014 Idiap Research Institute (Ronan Collobert)
-Copyright (c) 2012-2014 DeepMind Technologies    (Koray Kavukcuoglu)
-Copyright (c) 2011-2012 NEC Laboratories America (Clement Farabet)
-Copyright (c) 2011-2013 New York University    (Antoine Bordes)
-Copyright (c) 2012-2013 University of Montreal   (Pascal Vincent)
-Copyright (c) 2014-    Google Inc.
-Copyright (c) 2015-    Twitter, Inc.
-Copyright (c) 2015-    Intel Corporation
-Copyright (c) 2015-    AMD Inc.
-Copyright (c) 2016-    Baidu, Inc.
-Copyright (c) 2016-    Microsoft Corporation
-Copyright (c) 2017-    Amazon.com, Inc.
-Copyright (c) 2018-    Facebook AI Research
-Copyright (c) 2019-    fast.ai, Inc.
-Copyright (c) 2022-     PyTorch Contributors
-All rights reserved.
-
-Redistribution and use in source and binary forms, with or without
-modification, are permitted provided that the following conditions are met:
-
-  * Redistributions of source code must retain the above copyright notice, this
-    list of conditions and the following disclaimer.
-
-  * Redistributions in binary form must reproduce the above copyright notice,
-    this list of conditions and the following disclaimer in the documentation
-    and/or other materials provided with the distribution.
-
-  * Neither the name of Facebook Inc. nor the names of its contributors may be
-    used to endorse or promote products derived from this software without
-    specific prior written permission.
-
-THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
-AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
-IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
-DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
-FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
-DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
-SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
-CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
-OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
-OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
-
------
-
-Copyright "Baidu USA LLC" - Apache License 2.0
-License Text([https://github.com/bryancatanzaro/trove/blob/master/LICENSE](https://github.com/bryancatanzaro/trove/blob/master/LICENSE))
-
-This notice applies to **trove**.
-
-Copyright 2015-2016 Baidu USA LLC. All rights reserved.
-
-                                 Apache License
-                           Version 2.0, January 2004
-                        http://www.apache.org/licenses/
-
-   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
-
-   1. Definitions.
-
-      "License" shall mean the terms and conditions for use, reproduction,
-      and distribution as defined by Sections 1 through 9 of this document.
-
-      "Licensor" shall mean the copyright owner or entity authorized by
-      the copyright owner that is granting the License.
-
-      "Legal Entity" shall mean the union of the acting entity and all
-      other entities that control, are controlled by, or are under common
-      control with that entity. For the purposes of this definition,
-      "control" means (i) the power, direct or indirect, to cause the
-      direction or management of such entity, whether by contract or
-      otherwise, or (ii) ownership of fifty percent (50%) or more of the
-      outstanding shares, or (iii) beneficial ownership of such entity.
-
-      "You" (or "Your") shall mean an individual or Legal Entity
-      exercising permissions granted by this License.
-
-      "Source" form shall mean the preferred form for making modifications,
-      including but not limited to software source code, documentation
-      source, and configuration files.
-
-      "Object" form shall mean any form resulting from mechanical
-      transformation or translation of a Source form, including but
-      not limited to compiled object code, generated documentation,
-      and conversions to other media types.
-
-      "Work" shall mean the work of authorship, whether in Source or
-      Object form, made available under the License, as indicated by a
-      copyright notice that is included in or attached to the work
-      (an example is provided in the Appendix below).
-
-      "Derivative Works" shall mean any work, whether in Source or Object
-      form, that is based on (or derived from) the Work and for which the
-      editorial revisions, annotations, elaborations, or other modifications
-      represent, as a whole, an original work of authorship. For the purposes
-      of this License, Derivative Works shall not include works that remain
-      separable from, or merely link (or bind by name) to the interfaces of,
-      the Work and Derivative Works thereof.
-
-      "Contribution" shall mean any work of authorship, including
-      the original version of the Work and any modifications or additions
-      to that Work or Derivative Works thereof, that is intentionally
-      submitted to Licensor for inclusion in the Work by the copyright owner
-      or by an individual or Legal Entity authorized to submit on behalf of
-      the copyright owner. For the purposes of this definition, "submitted"
-      means any form of electronic, verbal, or written communication sent
-      to the Licensor or its representatives, including but not limited to
-      communication on electronic mailing lists, source code control systems,
-      and issue tracking systems that are managed by, or on behalf of, the
-      Licensor for the purpose of discussing and improving the Work, but
-      excluding communication that is conspicuously marked or otherwise
-      designated in writing by the copyright owner as "Not a Contribution."
-
-      "Contributor" shall mean Licensor and any individual or Legal Entity
-      on behalf of whom a Contribution has been received by Licensor and
-      subsequently incorporated within the Work.
-
-   2. Grant of Copyright License. Subject to the terms and conditions of
-      this License, each Contributor hereby grants to You a perpetual,
-      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
-      copyright license to reproduce, prepare Derivative Works of,
-      publicly display, publicly perform, sublicense, and distribute the
-      Work and such Derivative Works in Source or Object form.
-
-   3. Grant of Patent License. Subject to the terms and conditions of
-      this License, each Contributor hereby grants to You a perpetual,
-      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
-      (except as stated in this section) patent license to make, have made,
-      use, offer to sell, sell, import, and otherwise transfer the Work,
-      where such license applies only to those patent claims licensable
-      by such Contributor that are necessarily infringed by their
-      Contribution(s) alone or by combination of their Contribution(s)
-      with the Work to which such Contribution(s) was submitted. If You
-      institute patent litigation against any entity (including a
-      cross-claim or counterclaim in a lawsuit) alleging that the Work
-      or a Contribution incorporated within the Work constitutes direct
-      or contributory patent infringement, then any patent licenses
-      granted to You under this License for that Work shall terminate
-      as of the date such litigation is filed.
-
-   4. Redistribution. You may reproduce and distribute copies of the
-      Work or Derivative Works thereof in any medium, with or without
-      modifications, and in Source or Object form, provided that You
-      meet the following conditions:
-
-      (a) You must give any other recipients of the Work or
-          Derivative Works a copy of this License; and
-
-      (b) You must cause any modified files to carry prominent notices
-          stating that You changed the files; and
-
-      (c) You must retain, in the Source form of any Derivative Works
-          that You distribute, all copyright, patent, trademark, and
-          attribution notices from the Source form of the Work,
-          excluding those notices that do not pertain to any part of
-          the Derivative Works; and
-
-      (d) If the Work includes a "NOTICE" text file as part of its
-          distribution, then any Derivative Works that You distribute must
-          include a readable copy of the attribution notices contained
-          within such NOTICE file, excluding those notices that do not
-          pertain to any part of the Derivative Works, in at least one
-          of the following places: within a NOTICE text file distributed
-          as part of the Derivative Works; within the Source form or
-          documentation, if provided along with the Derivative Works; or,
-          within a display generated by the Derivative Works, if and
-          wherever such third-party notices normally appear. The contents
-          of the NOTICE file are for informational purposes only and
-          do not modify the License. You may add Your own attribution
-          notices within Derivative Works that You distribute, alongside
-          or as an addendum to the NOTICE text from the Work, provided
-          that such additional attribution notices cannot be construed
-          as modifying the License.
-
-      You may add Your own copyright statement to Your modifications and
-      may provide additional or different license terms and conditions
-      for use, reproduction, or distribution of Your modifications, or
-      for any such Derivative Works as a whole, provided Your use,
-      reproduction, and distribution of the Work otherwise complies with
-      the conditions stated in this License.
-
-   5. Submission of Contributions. Unless You explicitly state otherwise,
-      any Contribution intentionally submitted for inclusion in the Work
-      by You to the Licensor shall be under the terms and conditions of
-      this License, without any additional terms or conditions.
-      Notwithstanding the above, nothing herein shall supersede or modify
-      the terms of any separate license agreement you may have executed
-      with Licensor regarding such Contributions.
-
-   6. Trademarks. This License does not grant permission to use the trade
-      names, trademarks, service marks, or product names of the Licensor,
-      except as required for reasonable and customary use in describing the
-      origin of the Work and reproducing the content of the NOTICE file.
-
-   7. Disclaimer of Warranty. Unless required by applicable law or
-      agreed to in writing, Licensor provides the Work (and each
-      Contributor provides its Contributions) on an "AS IS" BASIS,
-      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
-      implied, including, without limitation, any warranties or conditions
-      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
-      PARTICULAR PURPOSE. You are solely responsible for determining the
-      appropriateness of using or redistributing the Work and assume any
-      risks associated with Your exercise of permissions under this License.
-
-   8. Limitation of Liability. In no event and under no legal theory,
-      whether in tort (including negligence), contract, or otherwise,
-      unless required by applicable law (such as deliberate and grossly
-      negligent acts) or agreed to in writing, shall any Contributor be
-      liable to You for damages, including any direct, indirect, special,
-      incidental, or consequential damages of any character arising as a
-      result of this License or out of the use or inability to use the
-      Work (including but not limited to damages for loss of goodwill,
-      work stoppage, computer failure or malfunction, or any and all
-      other commercial damages or losses), even if such Contributor
-      has been advised of the possibility of such damages.
-
-   9. Accepting Warranty or Additional Liability. While redistributing
-      the Work or Derivative Works thereof, You may choose to offer,
-      and charge a fee for, acceptance of support, warranty, indemnity,
-      or other liability obligations and/or rights consistent with this
-      License. However, in accepting such obligations, You may act only
-      on Your own behalf and on Your sole responsibility, not on behalf
-      of any other Contributor, and only if You agree to indemnify,
-      defend, and hold each Contributor harmless for any liability
-      incurred by, or claims asserted against, such Contributor by reason
-      of your accepting any such warranty or additional liability.
-
-   END OF TERMS AND CONDITIONS
+# Third Party Notices
+
+The scripts contained in this repository make use of the following third-party libraries:
+
+- [clipper](https://sourceforge.net/p/polyclipping/code/HEAD/tree/tags/6.2.0/License.txt)
+- [hatchling](https://github.com/pypa/hatch/blob/master/LICENSE.txt)
+- [numpy](https://github.com/numpy/numpy/blob/main/LICENSE.txt)
+- [pandas](https://github.com/pandas-dev/pandas/blob/main/LICENSE)
+- [Pillow](https://github.com/python-pillow/pillow/blob/main/LICENSE)
+- [scikit_learn](https://github.com/scikit-learn/scikit-learn/blob/main/COPYING)
+- [setuptools](https://github.com/pypa/setuptools/blob/main/LICENSE)
+- [Shapely](https://github.com/shapely/shapely/blob/main/LICENSE.txt)
+- [torch](https://github.com/pytorch/pytorch/blob/main/LICENSE)
+- [torchvision](https://github.com/pytorch/vision/blob/main/LICENSE)
+- [trove](https://github.com/bryancatanzaro/trove/blob/master/LICENSE)

docker-compose.yaml

@@ -1,5 +1,5 @@
 services:
-  nemotron-ocr:
+  nemo-retriever-ocr:
     build:
       context: .
       dockerfile: Dockerfile

example.py

@@ -4,11 +4,11 @@
 
 import argparse
 
-from nemotron_ocr.inference.pipeline import NemotronOCR
+from nemo_retriever_ocr.inference.pipeline import NemoRetrieverOCR
 
 
 def main(image_path, merge_level, no_visualize, model_dir):
-    ocr_pipeline = NemotronOCR(model_dir=model_dir)
+    ocr_pipeline = NemoRetrieverOCR()
 
     predictions = ocr_pipeline(image_path, merge_level=merge_level, visualize=not no_visualize)
 

{nemotron-ocr → nemo-retriever-ocr}/cpp/non_maximal_suppression/cuda_non_maximal_suppression.cu

@@ -157,8 +157,11 @@ void device_row_collapse(torch::PackedTensorAccessor64<T, 5> allQuads,
     torch::PackedTensorAccessor64<T, 3> allConfs,
     T confThreshold, T iouThreshold,
     torch::PackedTensorAccessor64<int32_t, 1> allOutCounts,
-    torch::PackedTensorAccessor64<T, 3> allOutEmbedQuads,
-    torch::PackedTensorAccessor64<int32_t, 2> allOutIds)
+    torch::PackedTensorAccessor64<T, 3> allOutEmbedQuads
+#ifdef NMS_VERIFY_CORRECTNESS
+    , torch::PackedTensorAccessor64<int32_t, 2> allOutIds
+#endif
+    )
 {
     typedef InPlaceQuad_<T> Quadf;
     static_assert(sizeof(Quadf) == sizeof(T) * 8, "Invalid QuadMem size!");
@@ -303,9 +306,11 @@ void device_row_collapse(torch::PackedTensorAccessor64<T, 5> allQuads,
         }
 
         write_embed_quad(outEmbedQuads, outQuad, storeOff + procLabel - 1);
+#ifdef NMS_VERIFY_CORRECTNESS
         if (threadRank == 0) {
             allOutIds[b][storeOff + procLabel - 1] = r * 32 + startIdx;
         }
+#endif
     }
 
     if (threadRank == 0) {
@@ -316,9 +321,9 @@ void device_row_collapse(torch::PackedTensorAccessor64<T, 5> allQuads,
 #undef threadRank
 }
 
-template<typename T>
+template<bool IsSingleExample, typename T>
 __global__
-void device_a2a_adjacency_sparse(const int32_t *ptrQuadCts,
+void device_a2a_adjacency_sparse(const uint64_t punCounts,
                                  T iouThreshold,
                                  torch::PackedTensorAccessor64<T, 3> embedQuads,
                                  torch::PackedTensorAccessor64<bool, 2> outIsStart,
@@ -327,11 +332,7 @@ void device_a2a_adjacency_sparse(const int32_t *ptrQuadCts,
 {
     const uint32_t b = blockIdx.y;
 
-    const int32_t quadCt = ptrQuadCts[b];
-
-    if (quadCt == 0) {
-        return;
-    }
+    const int32_t quadCt = IsSingleExample ? punCounts : reinterpret_cast<const int32_t*>(punCounts)[b];
 
     const int32_t jobIdx = blockIdx.x * blockDim.x + threadIdx.x;
     const int32_t row = jobIdx / quadCt;
@@ -342,7 +343,7 @@ void device_a2a_adjacency_sparse(const int32_t *ptrQuadCts,
         return;
     }
 
-    T* exData = embedQuads[b].data();
+    T* exData = IsSingleExample ? embedQuads.data() : embedQuads[b].data();
 
     const auto qRow = StridedEmbedQuad_<T>{ exData + row * embedQuads.stride(2), embedQuads.stride(1) }.Bounds(),
                qCol = StridedEmbedQuad_<T>{ exData + col * embedQuads.stride(2), embedQuads.stride(1) }.Bounds();
@@ -404,9 +405,9 @@ void device_a2a_adjacency_sparse(const int32_t *ptrQuadCts,
     }
 }
 
-template<uint32_t NumWarps, typename T, int32_t I_CELL_SIZE>
+template<uint32_t NumWarps, bool IsSingleExample, typename T, int32_t I_CELL_SIZE>
 __global__
-void device_a2a_adjacency_build_grid(const int32_t *ptrQuadCts,
+void device_a2a_adjacency_build_grid(const uint64_t punCounts,
                                      torch::PackedTensorAccessor64<T, 3> embedQuads,
                                      torch::PackedTensorAccessor64<int32_t, 4> outGridCells,
                                      torch::PackedTensorAccessor64<int32_t, 3> outQuadCells)
@@ -422,10 +423,10 @@ void device_a2a_adjacency_build_grid(const int32_t *ptrQuadCts,
 
     const uint32_t b = blockIdx.z;
 
-    const uint32_t quadCt = ptrQuadCts[b];
+    const uint32_t quadCt = IsSingleExample ? punCounts : reinterpret_cast<const int32_t*>(punCounts)[b];
     const uint32_t quadIdx = blockIdx.y;
 
-    if (quadIdx >= quadCt) {
+    if (!IsSingleExample && quadIdx >= quadCt) {
         return;
     }
 
@@ -484,9 +485,9 @@ void device_a2a_adjacency_build_grid(const int32_t *ptrQuadCts,
 
 typedef uint8_t visit_mask_t;
 
-template<uint32_t NumWarps, typename T>
+template<uint32_t NumWarps, bool IsSingleExample, typename T>
 __global__
-void device_a2a_adjacency_with_grid(const int32_t *ptrQuadCts,
+void device_a2a_adjacency_with_grid(const uint64_t punCounts,
                                     T iouThreshold,
                                     torch::PackedTensorAccessor64<T, 3> allEmbedQuads,
                                     torch::PackedTensorAccessor64<int32_t, 4> allCells,
@@ -502,10 +503,10 @@ void device_a2a_adjacency_with_grid(const int32_t *ptrQuadCts,
 
     const uint32_t b = blockIdx.z;
 
-    const uint32_t quadCt = ptrQuadCts[b];
+    const uint32_t quadCt = IsSingleExample ? punCounts : reinterpret_cast<const int32_t*>(punCounts)[b];
     const uint32_t quadIdx = blockIdx.y;
 
-    if (quadIdx >= quadCt) {
+    if (!IsSingleExample && quadIdx >= quadCt) {
         return;
     }
 
@@ -534,7 +535,7 @@ void device_a2a_adjacency_with_grid(const int32_t *ptrQuadCts,
     auto exAdjCounts = reinterpret_cast<uint32_t*>(outAdjCounts[b].data());
     auto exAdjValues = outSparseAdj[b][quadIdx].data();
 
-    T *exData = allEmbedQuads[b].data();
+    T *exData = IsSingleExample ? allEmbedQuads.data() : allEmbedQuads[b].data();
 
     const auto bdsAnchor = Quad_<T>{ s_quadVerts }.Bounds();
 
@@ -598,8 +599,9 @@ void device_a2a_adjacency_with_grid(const int32_t *ptrQuadCts,
     }
 }
 
+template<bool IsSingleExample>
 __global__
-void device_flatten_graph_iterative(const int32_t *ptrQuadCts,
+void device_flatten_graph_iterative(const uint64_t punCounts,
                                     torch::PackedTensorAccessor64<bool, 2> allIsStart,
                                     volatile uint32_t *allAdjCounts,
                                     volatile uint32_t *allAdjValues
@@ -620,12 +622,14 @@ void device_flatten_graph_iterative(const int32_t *ptrQuadCts,
     const uint32_t b = blockIdx.z;
     const uint32_t anchorRow = blockIdx.y;
 
-    const uint32_t quadCt = ptrQuadCts[b];
+    const uint32_t quadCt = IsSingleExample ? punCounts : reinterpret_cast<const int32_t*>(punCounts)[b];
 
     // Only need to check this if there are multiple examples, since in the case of a single example,
     // the grid is precisely sized to that quadCt
-    if (anchorRow >= quadCt) {
-        return;
+    if constexpr (!IsSingleExample) {
+        if (anchorRow >= quadCt) {
+            return;
+        }
     }
 
     auto isStart = allIsStart[b].data();
@@ -686,13 +690,12 @@ void device_flatten_graph_iterative(const int32_t *ptrQuadCts,
     visitStack[1] = anchorRow;
 #ifndef NDEBUG
     for (uint32_t i = 2; i < VISIT_STACK_SIZE; ++i) {
-        visitStack[i] = TERM_VALUE;
+        visitStack[i] = -2;
     }
 #endif
     int32_t visitPtr = 1;
 
-    // NOTE: This loop is actually terminated by the `if (warpNextCol == TERM_VALUE)` check below
-    for (uint32_t dfsIter = 0; true; ++dfsIter) {
+    while (true) {
 #ifdef NMS_VERIFY_CORRECTNESS
         assert(visitPtr >= 0 && visitPtr < VISIT_STACK_SIZE);
 #endif
@@ -704,7 +707,7 @@ void device_flatten_graph_iterative(const int32_t *ptrQuadCts,
         if (threadNextCol == warpNextCol) {
 #ifndef NDEBUG
             // This makes it easier to debug where the pointer is
-            visitStack[visitPtr] = TERM_VALUE;
+            visitStack[visitPtr] = -2;
 #endif
             --visitPtr;
         }
@@ -728,15 +731,12 @@ void device_flatten_graph_iterative(const int32_t *ptrQuadCts,
         const uint32_t procAdjCount = adjCounts[procRow];
         auto procAdjValues = adjValues + (procRow * maxExCount);
 
+        // Offsetting by the iteration number will help balance out the maximum depth of any stack in the warp.
+        // The reason behind this is due to how otherwise, warp-0 will always get a new element, warp-1 iff the adj graph
+        // has more than one element, warp-2 iff the adj graph has more than two elements, and so on. Basically,
+        // the warps have decreasing pressure. With the rotation mechanism, it helps to balance out stack usage.
         for (uint32_t i = threadRank; i < procAdjCount; i += WARP_SIZE) {
-            uint32_t adjCol = procAdjValues[i];
-
-            auto group = cg::coalesced_threads();
-            // Offsetting by the iteration number will help balance out the maximum depth of any stack in the warp.
-            // The reason behind this is due to how otherwise, warp-0 will always get a new element, warp-1 iff the adj graph
-            // has more than one element, warp-2 iff the adj graph has more than two elements, and so on. Basically,
-            // the warps have decreasing pressure. With the rotation mechanism, it helps to balance out stack usage.
-            adjCol = group.shfl(adjCol, (group.thread_rank() + dfsIter) % group.size());
+            const uint32_t adjCol = procAdjValues[i];
 
             // This will set the queued flag for this column, if it's not already set.
             // It also returns the old state. In our case, we only want to add this value to the
@@ -748,6 +748,7 @@ void device_flatten_graph_iterative(const int32_t *ptrQuadCts,
 
             bool alreadyAdded = oldMask & ADDED_MASK;
 
+            auto group = cg::coalesced_threads();
             const uint32_t gThreadRank = group.thread_rank();
             uint32_t notAddedBallot = group.ballot(!alreadyAdded);
             if (notAddedBallot) {
@@ -824,7 +825,8 @@ void add_to_set(const torch::TensorAccessor<int32_t, 1>& adjCounts,
     }
 }
 
-void cpu_flatten_graph(const int32_t *ptrQuadCts,
+template<bool IsSingleExample>
+void cpu_flatten_graph(const uint64_t punCounts,
                        torch::Tensor isStartTensorGPU,
                        torch::Tensor adjCountsTensorGPU,
                        torch::Tensor adjValuesTensorGPU)
@@ -838,7 +840,7 @@ void cpu_flatten_graph(const int32_t *ptrQuadCts,
     auto allAdjValues = adjValuesTensor.accessor<int32_t, 3>();
 
     for (int32_t b = 0; b < allAdjCounts.size(0); ++b) {
-        const int32_t quadCt = ptrQuadCts[b];
+        const int32_t quadCt = IsSingleExample ? punCounts : reinterpret_cast<const int32_t*>(punCounts)[b];
 
         for (int32_t row = 0; row < quadCt; ++row) {
             std::unordered_set<int32_t> fullAdjSet;
@@ -893,9 +895,9 @@ void device_a2a_adj_cleanup(const int32_t *counts,
     }
 }
 
-template<uint32_t NumWarps, typename T>
+template<uint32_t NumWarps, typename T, bool IsSingleExample>
 __global__
-void device_a2a_collapse(torch::PackedTensorAccessor64<int32_t, 1> quadCounts,
+void device_a2a_collapse(const uint64_t punCounts,
                          torch::PackedTensorAccessor64<T, 3> allEmbedQuads,
                          torch::PackedTensorAccessor64<bool, 2> allIsLeadRow,
                          const int64_t *regionCounts,
@@ -915,14 +917,16 @@ void device_a2a_collapse(torch::PackedTensorAccessor64<int32_t, 1> quadCounts,
     const uint32_t b = blockIdx.z;
     const uint32_t row = blockIdx.y;
 
-    const int32_t quadCt = quadCounts[b];
+    const int32_t quadCt = IsSingleExample ? punCounts : reinterpret_cast<const int32_t*>(punCounts)[b];
 
-    if (row >= quadCt) {
-        return;
+    if constexpr (!IsSingleExample) {
+        if (row >= quadCt) {
+            return;
+        }
     }
 
     // Only process the lead rows
-    const auto isLeadRow = allIsLeadRow[b].data();
+    const auto isLeadRow = IsSingleExample ? allIsLeadRow.data() : allIsLeadRow[b].data();
     if (!isLeadRow[row]) {
         return;
     }
@@ -941,7 +945,7 @@ void device_a2a_collapse(torch::PackedTensorAccessor64<int32_t, 1> quadCounts,
         __syncthreads();
     }
 
-    T *exData = allEmbedQuads[b].data();
+    T *exData = IsSingleExample ? allEmbedQuads.data() : allEmbedQuads[b].data();
 
     const int32_t adjCount = allAdjCounts[b][row];
     const int32_t *adjIdxs = allAdjValues[b][row].data();
@@ -982,12 +986,20 @@ void device_a2a_collapse(torch::PackedTensorAccessor64<int32_t, 1> quadCounts,
 
     // Figure out the output position
     uint32_t writePosition = 0;
-    for (int32_t i = threadRank; i < b; i += BLOCK_WIDTH) {
-        writePosition += regionCounts[i];
+    if constexpr (!IsSingleExample) {
+        for (int32_t i = threadRank; i < b; i += BLOCK_WIDTH) {
+            writePosition += regionCounts[i];
+        }
     }
 
+    const int32_t numLongs = row >> 3; // Divide by 8
     const uint8_t *pCurrIsLeadRow = reinterpret_cast<const uint8_t*>(isLeadRow);
-    for (int32_t i = threadRank; i < row; i += BLOCK_WIDTH) {
+    const uint64_t *lpCurrIsLeadRow = reinterpret_cast<const uint64_t*>(pCurrIsLeadRow);
+
+    for (int32_t i = threadRank; i < numLongs; i += BLOCK_WIDTH) {
+        writePosition += __popcll(lpCurrIsLeadRow[i]);
+    }
+    for (int32_t i = (numLongs * 8) + threadRank; i < row; i += BLOCK_WIDTH) {
         if (pCurrIsLeadRow[i]) {
             ++writePosition;
         }
@@ -1063,9 +1075,13 @@ CollapseRowsResult collapse_rows(
     int64_t embedSize = sizeof(EmbedQuad_<scalar_t>) / sizeof(scalar_t);
     auto rowMergeTensor = torch::empty({ quads.size(0), embedSize, quads.size(1) * quads.size(2) }, quads.options());
 
+#ifdef NMS_VERIFY_CORRECTNESS
     auto idsTensor = torch::full({ quads.size(0), quads.size(1) * quads.size(2) },
                                  std::numeric_limits<int32_t>::max(),
                                  counts.options().dtype(torch::kInt32));
+#else
+    torch::Tensor idsTensor;
+#endif
 
     dim3 blockSize(32, 3, 1);
     dim3 gridSize(1,
@@ -1077,8 +1093,10 @@ CollapseRowsResult collapse_rows(
         probs.packed_accessor64<scalar_t, 3>(),
         probThreshold, iouThreshold,
         counts.packed_accessor64<int32_t, 1>(),
-        rowMergeTensor.packed_accessor64<scalar_t, 3>(),
-        idsTensor.packed_accessor64<int32_t, 2>()
+        rowMergeTensor.packed_accessor64<scalar_t, 3>()
+#ifdef NMS_VERIFY_CORRECTNESS
+        , idsTensor.packed_accessor64<int32_t, 2>()
+#endif
     );
 
 #ifdef NMS_VERIFY_CORRECTNESS
@@ -1101,6 +1119,7 @@ CollapseRowsResult collapse_rows(
 
     counts = counts.slice(/*dim=*/ 0, 0, counts.size(0) - 1);
 
+#ifdef NMS_VERIFY_CORRECTNESS
     int64_t maxExCount;
     if (counts.size(0) > 1) {
         maxExCount = counts.max().item<int32_t>();
@@ -1112,12 +1131,13 @@ CollapseRowsResult collapse_rows(
 
     rowMergeTensor = rowMergeTensor.slice(2, 0, maxExCount);
     idsTensor = idsTensor.slice(1, 0, maxExCount);
-    auto order = torch::argsort(idsTensor, /*dim=*/ 1, s_sortOrder);
+    auto order = torch::argsort(idsTensor, /*dim=*/ 1, s_sortOrder); s_sortOrder = !s_sortOrder;
 
     auto embOrder = order.unsqueeze(1).expand_as(rowMergeTensor);
 
     rowMergeTensor = torch::gather(rowMergeTensor, /*dim=*/ 2, embOrder);
     idsTensor = torch::gather(idsTensor, /*dim=*/ 1, order);
+#endif
 
     return { counts, rowMergeTensor, totalQuads, idsTensor, imageWidth, imageHeight };
 }
@@ -1157,8 +1177,8 @@ struct AdjacencyResult {
     int64_t MaxExCount;
 };
 
-template<typename T>
-void cpu_a2a_adjacency_sparse(const int32_t *ptrQuadCts,
+template<bool IsSingleExample, typename T>
+void cpu_a2a_adjacency_sparse(const uint64_t punCounts,
                               const T iouThreshold,
                               torch::Tensor embedQuadsTensor,
                               torch::Tensor outIsStartTensorGPU,
@@ -1176,7 +1196,7 @@ void cpu_a2a_adjacency_sparse(const int32_t *ptrQuadCts,
     auto adjValues = outSparseAdjTensor.accessor<int32_t, 3>();
 
     for (int32_t b = 0; b < embedQuadsTensor.size(0); ++b) {
-        const int32_t quadCt = ptrQuadCts[b];
+        const int32_t quadCt = IsSingleExample ? punCounts : reinterpret_cast<const int32_t*>(punCounts)[b];
 
         T *exData = embedQuads[b].data();
 
@@ -1264,6 +1284,13 @@ AdjacencyResult compute_all_to_all_adjacency(
                                        counts.options().dtype(torch::kInt32));
 #endif
 
+    // If the batch is only a single example, instead of hitting global memory for the count, we can
+    // just encode the count into the pointer instead
+    uint64_t ptrCounts = reinterpret_cast<uint64_t>(counts.data_ptr<int32_t>());
+    if (counts.size(0) == 1) {
+        ptrCounts = maxExCount;
+    }
+
 #ifdef NMS_VERIFY_CORRECTNESS
     auto cpuAdjValuesTensor = adjValuesTensor.cpu();
     auto cpuAdjCountsTensor = adjCountsTensor.cpu();
@@ -1291,15 +1318,23 @@ AdjacencyResult compute_all_to_all_adjacency(
     //blockSize = dim3{ GRID_NUM_WARPS * 32, 1, 1 };
     //gridSize = dim3{ 1, static_cast<uint32_t>(maxExCount), static_cast<uint32_t>(counts.size(0)) };
 
-    //device_a2a_adjacency_build_grid<GRID_NUM_WARPS, scalar_t, CELL_SIZE> KERNEL_ARG2(gridSize, blockSize) (
-    //    counts.data_ptr<int32_t>(),
+    //auto buildGridFn = counts.size(0) == 1 ?
+    //    device_a2a_adjacency_build_grid<GRID_NUM_WARPS, true, scalar_t, CELL_SIZE> :
+    //    device_a2a_adjacency_build_grid<GRID_NUM_WARPS, false, scalar_t, CELL_SIZE>;
+
+    //buildGridFn KERNEL_ARG2(gridSize, blockSize) (
+    //    ptrCounts,
     //    collapseResult.StridedMergeQuads.packed_accessor64<scalar_t, 3>(),
     //    gridCellsTensor.packed_accessor64<int32_t, 4>(),
     //    quadCellExtentsTensor.packed_accessor64<int32_t, 3>()
     //);
 
-    //device_a2a_adjacency_with_grid<GRID_NUM_WARPS, scalar_t> KERNEL_ARG3(gridSize, blockSize, smemSize) (
-    //    counts.data_ptr<int32_t>(),
+    //auto adjGridFn = counts.size(0) == 1 ?
+    //    device_a2a_adjacency_with_grid<GRID_NUM_WARPS, true, scalar_t> :
+    //    device_a2a_adjacency_with_grid<GRID_NUM_WARPS, false, scalar_t>;
+
+    //adjGridFn KERNEL_ARG3(gridSize, blockSize, smemSize) (
+    //    ptrCounts,
     //    iouThreshold,
     //    collapseResult.StridedMergeQuads.packed_accessor64<scalar_t, 3>(),
     //    gridCellsTensor.packed_accessor64<int32_t, 4>(),
@@ -1316,9 +1351,11 @@ AdjacencyResult compute_all_to_all_adjacency(
     gridSize = dim3{div_up(totalWork, blockSize.x),
                     static_cast<uint32_t>(counts.size(0))};
 
+    auto adjFn = counts.size(0) == 1 ? device_a2a_adjacency_sparse<true, scalar_t> : device_a2a_adjacency_sparse<false, scalar_t>;
+
     // This algorithm is O(n^2) with n being the current number of quads
-    device_a2a_adjacency_sparse<scalar_t> KERNEL_ARG2(gridSize, blockSize) (
-        counts.data_ptr<int32_t>(),
+    adjFn KERNEL_ARG2(gridSize, blockSize) (
+        ptrCounts,
         iouThreshold,
         collapseResult.StridedMergeQuads.packed_accessor64<scalar_t, 3>(),
         isStartTensor.packed_accessor64<bool, 2>(),
@@ -1328,9 +1365,7 @@ AdjacencyResult compute_all_to_all_adjacency(
 
 
 #ifdef NMS_VERIFY_CORRECTNESS
-    auto cpuCounts = counts.cpu();
-
-    cpu_a2a_adjacency_sparse<scalar_t>(cpuCounts.data_ptr<int32_t>(), iouThreshold,
+    cpu_a2a_adjacency_sparse<true>(ptrCounts, iouThreshold,
         collapseResult.StridedMergeQuads, cpuIsStartTensor, cpuAdjCountsTensor, cpuAdjValuesTensor);
 
     adjValuesTensor = std::get<0>(torch::sort(adjValuesTensor, /*dim=*/ 2));
@@ -1345,12 +1380,16 @@ AdjacencyResult compute_all_to_all_adjacency(
     auto maxDepthTensor = torch::tensor(0, adjCountsTensor.options());
 #endif
 
+    auto traverseFn = counts.size(0) == 1 ?
+                        device_flatten_graph_iterative<true> :
+                        device_flatten_graph_iterative<false>;
+
     blockSize = dim3{ 128, 1, 1 };
     gridSize = dim3{ 1, static_cast<uint32_t>(maxExCount), static_cast<uint32_t>(counts.size(0)) };
     smemSize = div_up(maxExCount * sizeof(visit_mask_t), sizeof(uint32_t)) * sizeof(uint32_t);
 
-    device_flatten_graph_iterative KERNEL_ARG3(gridSize, blockSize, smemSize) (
-        counts.data_ptr<int32_t>(),
+    traverseFn KERNEL_ARG3(gridSize, blockSize, smemSize) (
+        ptrCounts,
         isStartTensor.packed_accessor64<bool, 2>(),
         reinterpret_cast<uint32_t*>(adjCountsTensor.data_ptr<int32_t>()),
         reinterpret_cast<uint32_t*>(adjValuesTensor.data_ptr<int32_t>())
@@ -1360,7 +1399,7 @@ AdjacencyResult compute_all_to_all_adjacency(
     );
 
 #ifdef NMS_VERIFY_CORRECTNESS
-    cpu_flatten_graph(cpuCounts.data_ptr<int32_t>(), cpuIsStartTensor, cpuAdjCountsTensor, cpuAdjValuesTensor);
+    cpu_flatten_graph<true>(ptrCounts, cpuIsStartTensor, cpuAdjCountsTensor, cpuAdjValuesTensor);
 
     cpuAdjValuesTensor = std::get<0>(torch::sort(cpuAdjValuesTensor, /*dim=*/ 2));
     adjValuesTensor = std::get<0>(torch::sort(adjValuesTensor, /*dim=*/ 2));
@@ -1398,6 +1437,7 @@ AdjacencyResult compute_all_to_all_adjacency(
     cpuIsStartTensor = isStartTensor.cpu();
     cpuAdjCountsTensor = adjCountsTensor.cpu();
     cpuAdjValuesTensor = adjValuesTensor.cpu();
+    auto cpuCounts = counts.cpu();
     auto cpuCollapseIds = collapseResult.QuadIds.cpu();
 
     static std::vector<std::unordered_set<int32_t>> s_knownGroups;
@@ -1549,11 +1589,22 @@ nms_result_t
     dim3 blockSize(NUM_WARPS * 32, 1, 1);
     dim3 gridSize(1, adjResult.MaxExCount, counts.size(0));
 
+    // If the batch is only a single example, instead of hitting global memory for the count, we can
+    // just encode the count into the pointer instead
+    uint64_t ptrCounts = reinterpret_cast<uint64_t>(counts.data_ptr<int32_t>());
+    if (counts.size(0) == 1) {
+        ptrCounts = adjResult.MaxExCount;
+    }
+
     torch::Tensor outQuads = torch::empty({ numOutQuads, 4, 2 }, embedQuads.options());
     torch::Tensor outConf = torch::empty({ numOutQuads }, embedQuads.options());
 
-    device_a2a_collapse<NUM_WARPS, scalar_t>  KERNEL_ARG2(gridSize, blockSize) (
-        counts.packed_accessor64<int32_t, 1>(),
+    auto collapseFn = counts.size(0) == 1 ?
+        device_a2a_collapse<NUM_WARPS, scalar_t, true> :
+        device_a2a_collapse<NUM_WARPS, scalar_t, false>;
+
+    collapseFn KERNEL_ARG2(gridSize, blockSize) (
+        ptrCounts,
         embedQuads.packed_accessor64<scalar_t, 3>(),
         isLeadRow.packed_accessor64<bool, 2>(),
         regionCounts.data_ptr<int64_t>(),