nemotron-ocr 2.0.0__tar.gz

nemotron_ocr-2.0.0/.gitignore

@@ -0,0 +1,9 @@
+# Built C++/CUDA extension (produced by: pip install -v .)
+src/nemotron_ocr_cpp/*.so
+src/nemotron_ocr_cpp/*.pyd
+__pycache__/
+*.py[cod]
+.pytest_cache/
+.venv/
+build/
+*.egg-info/

nemotron_ocr-2.0.0/PKG-INFO

@@ -0,0 +1,13 @@
+Metadata-Version: 2.4
+Name: nemotron-ocr
+Version: 2.0.0
+Summary: Nemoton OCR
+Project-URL: Homepage, https://nvidia.com/
+Author: NVIDIA Nemotron
+Requires-Python: <3.13,>=3.12
+Requires-Dist: huggingface-hub>=0.20.0
+Requires-Dist: numpy>=1.26.0
+Requires-Dist: pillow>=12.0.0
+Requires-Dist: shapely<3,>=2.1.2
+Requires-Dist: torchvision~=0.26.0
+Requires-Dist: torch~=2.11.0

nemotron_ocr-2.0.0/cpp/.gitattributes

@@ -0,0 +1 @@
+load_png/wuffs-v0.3.c filter=lfs diff=lfs merge=lfs -text

nemotron_ocr-2.0.0/cpp/.gitignore

@@ -0,0 +1,6 @@
+__pycache__
+.vscode
+build
+*.egg-info
+dist
+.vs

nemotron_ocr-2.0.0/cpp/.gitmodules

@@ -0,0 +1,3 @@
+[submodule "trove"]
+	path = trove
+	url = https://github.com/bryancatanzaro/trove.git

nemotron_ocr-2.0.0/cpp/README.md

@@ -0,0 +1,15 @@
+# Optimized Image Operations for PyTorch
+
+## Installation
+
+```
+python setup.py install
+```
+
+## Usage
+
+```
+# It's important that you do this first
+import torch
+from pytorch_image_ops import color_transform, spatial_transform
+```

nemotron_ocr-2.0.0/cpp/beam_decode/beam_decode.cpp

@@ -0,0 +1,459 @@
+// SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+// SPDX-License-Identifier: Apache-2.0
+
+#include "beam_decode.h"
+
+#include <vector>
+#include <deque>
+#include <limits>
+#include <memory>
+#include <unordered_set>
+#include <set>
+#include <algorithm>
+#include <chrono>
+
+#include "../common.h"
+#include "prefix.h"
+#include "log_sum_exp.h"
+#include "sbo_lm.h"
+
+using namespace std;
+
+template<typename scalar_t>
+using pred_seq_t = torch::TensorAccessor<scalar_t, 2>;
+
+struct PrefixScore
+{
+    float_t lProbBlank;
+    float_t lProbChar;
+    // float_t raw_lProbBlank;
+    // float_t raw_lProbChar;
+    mutable float_t _lProb;
+
+    PrefixScore(float_t lProbBlank = NEG_INF /* log P(0) */, float_t lProbChar = NEG_INF /* log P(0) */)
+        : lProbBlank(lProbBlank), lProbChar(lProbChar), _lProb(NEG_INF)
+        //   , raw_lProbBlank(lProbBlank), raw_lProbChar(lProbChar)
+    {}
+
+    float_t get_lScore() const {
+        if (_lProb == NEG_INF) {
+            _lProb = log_sum_exp(lProbBlank, lProbChar);
+        }
+        return _lProb;
+    }
+
+    // float_t get_raw_lScore() const {
+    //     return log_sum_exp(raw_lProbBlank, raw_lProbChar);
+    // }
+};
+
+typedef std::unordered_map<Prefix*, PrefixScore> PrefixMap;
+typedef std::pair<Prefix*, PrefixScore> BeamItem;
+typedef std::vector<BeamItem> Beam;
+
+/*
+    Allows us to get an estimate of the vision model confidence, irrespective of how the language
+    model guided the decoding. NOTE: This scoring could follow an entirely different path than
+    the returned decoded sequence.
+*/
+template<typename scalar_t>
+scalar_t get_vision_confidence(const pred_seq_t<scalar_t> &logProbs, scalar_t minProb)
+{
+    const int64_t T = logProbs.size(0);
+    const int64_t S = logProbs.size(1);
+
+    scalar_t ret = 0; // log(1)
+
+    for (size_t t = 0; t < T; ++t) {
+        float_t maxP = logProbs[t][0];
+        int64_t maxC = 0;
+        for (int64_t c = 1; c < S; ++c) {
+            float_t p = logProbs[t][c];
+            if (p > maxP) {
+                maxP = p;
+                maxC = c;
+            }
+        }
+        ret += maxP;
+        // Ignore everything past the sequence terminator
+        if (maxC == 1) {
+            break;
+        }
+
+        if (ret < minProb) {
+            break;
+        }
+    }
+
+    return ret;
+}
+
+
+template<typename scalar_t>
+pair<vector<token_t>, float_t>
+    ctc_beam_decode_impl(const pred_seq_t<scalar_t> &probs, const int64_t beamSize,
+                         const int64_t blank, scalar_t minProb,
+                         const LanguageModel &langModel, scalar_t lmWeight)
+{
+    if (blank != 0) {
+        throw runtime_error("Currently, only ordinal 0 supported for the blank prediction");
+    }
+
+    const int64_t T = probs.size(0);
+    const int64_t S = probs.size(1);
+
+    // NOTE: In log space, the following is true:
+    // 1. Adding two probabilities: log_sum_exp(l_p_a, l_p_b)
+    // 2. Multiplying two probabilities: l_p_a + l_p_b
+    // 3. log P(0) = -inf
+    // 4. log P(1) = 0
+
+    // Convert to log-space
+    if (minProb > 0) {
+        minProb = log(minProb);
+    } else {
+        minProb = NEG_INF;
+    }
+
+    auto retScore = get_vision_confidence(probs, minProb);
+
+    if (retScore < minProb) {
+        return { {}, NEG_INF };
+    }
+
+    PrefixAllocator prefixAlloc;
+
+    Beam beam;
+    beam.emplace_back(prefixAlloc.GetPrefix(), PrefixScore{0, NEG_INF}); // Add a dummy first node
+
+    Beam terminated;
+
+    typedef tuple<Prefix*, token_t> lm_cache_key_t;
+    unordered_map<lm_cache_key_t, float_t> lmScoreCache;
+
+    for (int64_t t = 0; t < T; ++t) {
+        PrefixMap nextBeam;
+
+        // Add all of the completed paths to the next beam.
+        // This allows us to accumulate new paths into these,
+        // but otherwise not process them
+        for (const BeamItem &prevNode : beam) {
+            if (prevNode.first->Token == 1) {
+                nextBeam.insert(prevNode);
+            }
+        }
+
+        // Loop over vocab
+        for (int64_t s = 0; s < S; ++s) {
+            float_t lpEmit = probs[t][s];
+
+            if (lpEmit < minProb) {
+                continue;
+            }
+
+            for (const BeamItem &prevNode : beam) {
+                Prefix *prevPrefix = prevNode.first;
+                const PrefixScore &prevScore = prevNode.second;
+
+                // Ignore already completed paths
+                if (prevPrefix->Token == 1) {
+                    continue;
+                }
+
+                // Ignore impossible paths
+                if (prevScore.lProbBlank == NEG_INF && prevScore.lProbChar == NEG_INF) {
+                    continue;
+                }
+
+                // If we propose a blank the prefix doesn't change.
+                // Only the probability of ending in blank gets updated.
+                if (s == blank) {
+                    PrefixScore &score = nextBeam[prevPrefix];
+                    score.lProbBlank     = log_sum_exp(score.lProbBlank    , prevScore.lProbBlank     + lpEmit, prevScore.lProbChar     + lpEmit);
+                    // score.raw_lProbBlank = log_sum_exp(score.raw_lProbBlank, prevScore.raw_lProbBlank + lpEmit, prevScore.raw_lProbChar + lpEmit);
+                    continue;
+                }
+
+                // Extend the prefix by the new character s and add it to the beam.
+                // Only the probability of not ending in blank gets updated.
+                token_t prevToken = prevPrefix->Token;
+
+                // NOTE: We always create a new prefix regardless of duplication because the PrefixScore
+                // is simultaneously tracking prefixes that do and don't end in a blank. And it's those
+                // that end in a blank that would cause the prefix to be extended.
+                auto extendPrefix = prefixAlloc.GetPrefix(s, prevPrefix);
+
+                // Evaluate the language model, but use the cache if we've already considered this string before
+                auto lmCacheItem = make_tuple(prevPrefix, s);
+                auto lmCacheIter = lmScoreCache.find(lmCacheItem);
+                float_t lpLang = 0;
+                if (lmCacheIter == lmScoreCache.end()) {
+                    lpLang = langModel.ScoreTransition(prevPrefix, s);
+                    lpLang *= lmWeight;
+                    lmCacheIter = lmScoreCache.emplace(lmCacheItem, lpLang).first;
+                }
+                lpLang = lmCacheIter->second;
+
+                PrefixScore &extendScore = nextBeam[extendPrefix];
+                // Remember, adding two log probabilities is equivalent to multiplying two probabilities
+                if (s != prevToken) {
+                    extendScore.lProbChar     = log_sum_exp(extendScore.lProbChar,     prevScore.lProbBlank     + lpEmit + lpLang, prevScore.lProbChar     + lpEmit + lpLang);
+                    // extendScore.raw_lProbChar = log_sum_exp(extendScore.raw_lProbChar, prevScore.raw_lProbBlank + lpEmit         , prevScore.raw_lProbChar + lpEmit         );
+                } else {
+                    // We don't include the previous probability of not ending in blank if s is repeated at the end. The CTC
+                    // algorithm merges characters not separated by a blank.
+                    extendScore.lProbChar     = log_sum_exp(extendScore.lProbChar    , prevScore.lProbBlank     + lpEmit + lpLang);
+                    // extendScore.raw_lProbChar = log_sum_exp(extendScore.raw_lProbChar, prevScore.raw_lProbBlank + lpEmit         );
+                }
+
+                // If the token is repeated, we also have to deal with the unchanged prefix since repeated characters are collapsed
+                if (s == prevToken) {
+                    PrefixScore &collapseScore = nextBeam[prevPrefix];
+                    collapseScore.lProbChar     = log_sum_exp(collapseScore.lProbChar    , prevScore.lProbChar     + lpEmit);
+                    // collapseScore.raw_lProbChar = log_sum_exp(collapseScore.raw_lProbChar, prevScore.raw_lProbChar + lpEmit);
+                }
+
+            }
+        }
+
+        Beam vecNextBeam(begin(nextBeam), end(nextBeam));
+
+        if (vecNextBeam.size() > beamSize) {
+            partial_sort(begin(vecNextBeam), begin(vecNextBeam) + beamSize, end(vecNextBeam),
+                [] (const BeamItem &a, const BeamItem &b) {
+                    return a.second.get_lScore() > b.second.get_lScore();
+                }
+            );
+            vecNextBeam.resize(beamSize);
+        }
+
+        beam = move(vecNextBeam);
+    }
+
+    // Find the best raw score
+    const BeamItem *bestItem = nullptr;
+    // for (const BeamItem &b : beam) {
+    //     if (bestItem == nullptr or b.second.get_raw_lScore() > bestItem->second.get_raw_lScore()) {
+    //         bestItem = &b;
+    //     }
+    // }
+    if (! beam.empty()) {
+        bestItem = &beam[0];
+    }
+
+    if (bestItem != nullptr) {
+        auto retList = bestItem->first->ToList();
+
+        return { move(retList), retScore };
+    } else {
+        return { {}, NEG_INF };
+    }
+}
+
+typedef std::pair<Prefix*, float_t> RegBeamItem;
+
+bool operator<(const RegBeamItem &a, const RegBeamItem &b) {
+    return a.second > b.second;
+}
+
+template<typename scalar_t>
+pair<vector<token_t>, float_t>
+    reg_beam_decode_impl(const pred_seq_t<scalar_t> &logProbs, const int64_t beamSize,
+                         scalar_t minProb,
+                         const LanguageModel &langModel, scalar_t lmWeight)
+{
+    const int64_t T = logProbs.size(0);
+    const int64_t S = logProbs.size(1);
+
+    // NOTE: In log space, the following is true:
+    // 1. Adding two probabilities: log_sum_exp(l_p_a, l_p_b)
+    // 2. Multiplying two probabilities: l_p_a + l_p_b
+    // 3. log P(0) = -inf
+    // 4. log P(1) = 0
+
+    // Convert to log-space
+    if (minProb > 0) {
+        minProb = log(minProb);
+    } else {
+        minProb = NEG_INF;
+    }
+
+    auto retScore = get_vision_confidence(logProbs, minProb);
+
+    if (retScore < minProb) {
+        return { {}, NEG_INF };
+    }
+
+    PrefixAllocator prefixAlloc;
+
+    vector<RegBeamItem> beam, nextBeam;
+    beam.emplace_back(prefixAlloc.GetPrefix(), 0); // log(1) = 0
+
+    for (int64_t t = 0; t < T && !beam.empty(); ++t) {
+        nextBeam.clear();
+
+        auto addToBeam = [&nextBeam, beamSize] (const RegBeamItem &rbi) {
+            nextBeam.push_back(rbi);
+        };
+
+        // Expand each path in the beam
+        for (const RegBeamItem &prevNode : beam) {
+            if (prevNode.first->Token == 1) {
+                // Move completed paths along without processing further
+                addToBeam(prevNode);
+                continue;
+            }
+
+            Prefix *prevPrefix = prevNode.first;
+            float_t prevScore = prevNode.second;
+
+            // Loop over vocab
+            for (int64_t s = 0; s < S; ++s) {
+                float_t lpEmit = logProbs[t][s];
+
+                if (lpEmit < minProb) {
+                    // The probability dropped below threshold, so stop processing this path
+                    continue;
+                }
+
+                auto extendPrefix = prefixAlloc.GetPrefix(s, prevPrefix);
+
+                float_t lpLang = langModel.ScoreTransition(prevPrefix, s);
+
+                float_t lpNext = prevScore + lpLang + lpEmit;
+
+                addToBeam({extendPrefix, lpNext});
+            }
+        }
+
+        if (nextBeam.size() > beamSize) {
+            // Find the top-k items, and then truncate the rest
+            partial_sort(begin(nextBeam), begin(nextBeam) + beamSize, end(nextBeam));
+            nextBeam.resize(beamSize);
+        }
+
+        std::swap(beam, nextBeam);
+    }
+
+    if (!beam.empty()) {
+        // The highest probability element will always be in the back
+        RegBeamItem rbi{ nullptr, NEG_INF };
+        for (auto &rb : beam) {
+            if (rbi.first == nullptr || rb.second > rbi.second) {
+                rbi = rb;
+            }
+        }
+
+        auto retList = rbi.first->ToList();
+
+        return { move(retList), retScore };
+    } else {
+        return { {}, NEG_INF };
+    }
+}
+
+
+
+template<typename scalar_t>
+void dp_beam_decode_impl(const torch::TensorAccessor<scalar_t, 3> &probsAccess,
+                         torch::TensorAccessor<int64_t, 2> retAccess,
+                         torch::TensorAccessor<scalar_t, 1> confAccess,
+                         int64_t beamSize, int64_t blank,
+                         scalar_t minProb,
+                         const LanguageModel *langModel,
+                         scalar_t lmWeight,
+                         bool combineDuplicates)
+{
+    const int64_t N = probsAccess.size(0);
+
+    #pragma omp parallel for num_threads(8)
+    for (int64_t i = 0; i < N; ++i) {
+        vector<token_t> seq;
+        float_t lConf;
+        if (combineDuplicates) {
+            tie(seq, lConf) = ctc_beam_decode_impl(probsAccess[i], beamSize, blank,
+                                                   minProb,
+                                                   *langModel, lmWeight);
+        } else {
+            tie(seq, lConf) = reg_beam_decode_impl(probsAccess[i], beamSize,
+                                                   minProb,
+                                                   *langModel, lmWeight);
+        }
+
+        int64_t sz = min<int64_t>(seq.size(), retAccess.size(1));
+
+        for (int64_t k = 0; k < sz; ++k) {
+            retAccess[i][k] = seq[k];
+        }
+
+        confAccess[i] = exp(lConf);
+    }
+}
+
+std::tuple<torch::Tensor, torch::Tensor>
+    beam_decode(torch::Tensor probs, int64_t beamSize, int64_t blank,
+                float minProb,
+                const LanguageModel *langModel,
+                float lmWeight,
+                bool combineDuplicates)
+{
+    if (langModel == nullptr) {
+        langModel = &NullLanguageModel;
+    }
+
+    auto tStart = chrono::high_resolution_clock::now();
+
+    probs = probs.contiguous();
+
+    bool collapse = false;
+    if (probs.dim() == 2) {
+        // N,T,C
+        probs = probs.unsqueeze(0);
+        collapse = true;
+    }
+
+    probs = probs.log();
+
+    torch::Tensor ret = torch::ones({ probs.size(0), probs.size(1) }, torch::kInt64);
+    torch::Tensor conf = torch::zeros({ probs.size(0) }, probs.options());
+
+    auto retAccess = ret.accessor<int64_t, 2>();
+
+    AT_DISPATCH_FLOATING_TYPES(
+        probs.scalar_type(),
+        "cpu_beam_decode",
+        ([&] {
+            dp_beam_decode_impl(
+                probs.accessor<scalar_t, 3>(),
+                retAccess,
+                conf.accessor<scalar_t, 1>(),
+                beamSize, blank,
+                static_cast<scalar_t>(minProb),
+                langModel,
+                static_cast<scalar_t>(lmWeight),
+                combineDuplicates
+            );
+        })
+    );
+
+    if (collapse) {
+        ret = ret.squeeze(0);
+        conf = conf[0];
+    }
+
+    auto tEnd = chrono::high_resolution_clock::now();
+
+    typedef chrono::duration<double, std::milli> tp_t;
+    tp_t totalElapsed = tEnd - tStart;
+
+    cout << "Beam Decode " << probs.size(0) << " - "
+         << "Total: " << totalElapsed.count() << "ms"
+         << endl;
+
+    return { ret, conf };
+}
+
+std::unique_ptr<LanguageModel> create_sbo_lm(const std::string &dataFilePath, token_mapping_t tokenMapping, float_t backoffWeight)
+{
+    return make_unique<SBO_LanguageModel>(dataFilePath, move(tokenMapping), backoffWeight);
+}

nemotron_ocr-2.0.0/cpp/beam_decode/beam_decode.h

@@ -0,0 +1,17 @@
+// SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+// SPDX-License-Identifier: Apache-2.0
+
+#pragma once
+
+#include <torch/torch.h>
+
+#include "language_model.h"
+
+std::tuple<torch::Tensor, torch::Tensor>
+    beam_decode(torch::Tensor probs, int64_t beamSize, int64_t blank,
+                float minProb,
+                const LanguageModel *langModel,
+                float lmWeight,
+                bool combineDuplicates);
+
+std::unique_ptr<LanguageModel> create_sbo_lm(const std::string &dataFilePath, token_mapping_t tokenMapping, float_t backoffWeight);

nemotron_ocr-2.0.0/cpp/beam_decode/kn_lm.cpp

@@ -0,0 +1,85 @@
+// SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+// SPDX-License-Identifier: Apache-2.0
+
+#include "kn_lm.h"
+
+using namespace std;
+
+
+KN_LanguageModel::KN_LanguageModel(const string &dataFilePath, token_mapping_t tokenMapping, float_t knDelta)
+    : NGramLMBase(dataFilePath, move(tokenMapping)), m_knDelta(knDelta)
+{
+}
+
+float KN_LanguageModel::ScoreTransitionImpl(const std::wstring &prefix, const std::wstring &suffix) const
+{
+    if (prefix.empty()) {
+        return ScoreUnigram(suffix);
+    } else {
+        return ScoreTransition(prefix, suffix);
+    }
+}
+
+float_t KN_LanguageModel::ScoreUnigram(const std::wstring &uni) const
+{
+    auto lIter = m_lookup[1].find(L""s);
+    if (lIter == m_lookup[1].end()) {
+        throw std::runtime_error("Unigrams not supported by this model!");
+    }
+
+    auto uniIter = lIter->second.find(uni);
+    float_t ctUni = 1e-8;
+    if (uniIter != lIter->second.end()) {
+        ctUni = uniIter->second;
+    }
+
+    float_t ctSuffixes = GetPrefixSum(L""s);
+
+    return ctUni / ctSuffixes;
+}
+
+float_t KN_LanguageModel::ScoreTransition(const std::wstring &prefix, const std::wstring &suffix) const
+{
+    if (prefix.empty()) {
+        // The number of distinct bigrams that end with this token
+        auto rlIter = m_reverseLookup.find(suffix);
+
+        float_t ctEndingBigrams = 0;
+        if (rlIter != m_reverseLookup.end()) {
+            ctEndingBigrams = rlIter->second[2].size();
+        }
+
+        float_t ctAllBigrams = m_lookup[2].size();
+
+        return ctEndingBigrams / ctAllBigrams;
+    }
+
+    auto lIter = m_lookup[prefix.size() + 1].find(prefix);
+    float_t ctUqSuffixes = 0;
+    float_t ctSuffixes = 0;
+    float_t ctSuffix = 0;
+    if (lIter != m_lookup[prefix.size() + 1].end()) {
+        ctUqSuffixes = lIter->second.size();
+
+        ctSuffixes = GetPrefixSum(prefix);
+
+        auto sIter = lIter->second.find(suffix);
+        if (sIter != lIter->second.end()) {
+            ctSuffix = sIter->second;
+        }
+    }
+
+    float_t factor = 0;
+    float_t main = 0;
+    if (ctSuffixes != 0) {
+        factor = m_knDelta * ctUqSuffixes / ctSuffixes;
+        // TODO: Figure out how to make this call without copying the string!
+        factor *= ScoreTransition({begin(prefix) + 1, end(prefix)}, suffix);
+
+        main = max<float_t>(ctSuffix - m_knDelta, 0) / ctSuffixes;
+    }
+
+    float_t total = main + factor;
+
+    return total;
+}

nemotron_ocr-2.0.0/cpp/beam_decode/kn_lm.h

@@ -0,0 +1,26 @@
+// SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+// SPDX-License-Identifier: Apache-2.0
+
+#pragma once
+
+#include <unordered_map>
+#include <vector>
+
+#include "ngram_lm_base.h"
+
+
+class KN_LanguageModel
+    : public NGramLMBase
+{
+public:
+    KN_LanguageModel(const std::string &dataFilePath, token_mapping_t tokenMapping, float_t knDelta);
+
+protected:
+    virtual float_t ScoreTransitionImpl(const std::wstring &prefix, const std::wstring &suffix) const override;
+
+private:
+    float_t ScoreUnigram(const std::wstring &uni) const;
+    float_t ScoreTransition(const std::wstring &prefix, const std::wstring &suffix) const;
+
+    float_t m_knDelta;
+};