Duplicate from nvidia/nemotron-ocr-v2

82f8ed3 2 days ago

3.2 kB

// SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
// SPDX-License-Identifier: Apache-2.0

#pragma once

#include <cstdlib>
#include <memory>
#include <vector>
#include <unordered_map>
#include <list>

typedef int32_t token_t;

class Prefix;

// typedef std::shared_ptr<Prefix> PrefixPtr;

class Prefix
{
public:
    token_t Token;
    Prefix *Parent;

    Prefix(token_t token = 0 /* blank */, Prefix *parent = nullptr)
        : Token(token), Parent(parent)
    {}

    std::vector<token_t> ToList() const;

    size_t size() const;
};


///// Borrowed from Boost libraries
template<typename T>
void hash_combine(size_t & seed, T const& v)
{
    seed ^= std::hash<T>()(v) + 0x9e3779b9 + (seed << 6) + (seed >> 2);
}
/////

namespace std {
template<>
struct hash<Prefix*>
{
    size_t operator()(const Prefix *p) const noexcept
    {
        size_t seed = 0;

        while (p) {
            if (p->Token != 0) {
                hash_combine(seed, p->Token);
            }
            p = p->Parent;
        }
        return seed;
    }
};

template<>
struct hash<tuple<Prefix*, token_t>>
{
    size_t operator()(const tuple<Prefix*, token_t> &t) const noexcept
    {
        size_t seed = 0;
        hash_combine(seed, get<0>(t));
        hash_combine(seed, get<1>(t));
        return seed;
    }
};

template<>
struct equal_to<Prefix*>
{
    bool operator()(const Prefix *a, const Prefix *b) const noexcept
    {
        while (a != nullptr && b != nullptr) {
            if (a->Token != b->Token) {
                return false;
            }
            a = a->Parent;
            b = b->Parent;
        }
        // If one chain is shorter than the other
        return a == b;
    }
};
}

inline size_t Prefix::size() const
{
    size_t ret = 0;
    auto p = this;
    while (p != nullptr) {
        ret += 1;
        p = p->Parent;
    }
    return ret;
}


class PrefixAllocator
{
public:
    PrefixAllocator() = default;
    ~PrefixAllocator();

    template<typename ...Args>
    Prefix *GetPrefix(Args&& ...ctorArgs);

private:
    void AllocateNextBuffer();

    std::list<Prefix*> m_buffers;
    size_t m_allocSize = 0;
    size_t m_currOff = 0;
};

inline PrefixAllocator::~PrefixAllocator()
{
    for (auto p : m_buffers) {
        // Prefix is a POD, and are allocated without initializing
        // to prevent redundant work upfront
        // delete[] p;
        free(p);
    }
}

inline void PrefixAllocator::AllocateNextBuffer()
{
    size_t nextSize = m_allocSize == 0 ? 1000 : 2 * m_allocSize;

    // Using malloc here to prevent the ctor of Prefix being called for each item.
    // Instead, the ctor will be called upon first access using GetPrefix
    auto pBuff = reinterpret_cast<Prefix*>(malloc(sizeof(Prefix) * nextSize));

    m_buffers.push_back(pBuff);

    m_allocSize = nextSize;
    m_currOff = 0;
}

template<typename ...Args>
Prefix *PrefixAllocator::GetPrefix(Args&& ...ctorArgs)
{
    if (m_currOff == m_allocSize) {
        AllocateNextBuffer();
    }

    auto buff = m_buffers.back() + m_currOff;

    auto ret = new (buff) Prefix(std::forward<Args>(ctorArgs)...);

    ++m_currOff;

    return ret;
}

	// SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
	// SPDX-License-Identifier: Apache-2.0

	#pragma once

	#include <cstdlib>
	#include <memory>
	#include <vector>
	#include <unordered_map>
	#include <list>

	typedef int32_t token_t;

	class Prefix;

	// typedef std::shared_ptr<Prefix> PrefixPtr;

	class Prefix
	{
	public:
	token_t Token;
	Prefix *Parent;

	Prefix(token_t token = 0 /* blank /, Prefix parent = nullptr)
	: Token(token), Parent(parent)
	{}

	std::vector<token_t> ToList() const;

	size_t size() const;
	};


	///// Borrowed from Boost libraries
	template<typename T>
	void hash_combine(size_t & seed, T const& v)
	{
	seed ^= std::hash<T>()(v) + 0x9e3779b9 + (seed << 6) + (seed >> 2);
	}
	/////

	namespace std {
	template<>
	struct hash<Prefix*>
	{
	size_t operator()(const Prefix *p) const noexcept
	{
	size_t seed = 0;

	while (p) {
	if (p->Token != 0) {
	hash_combine(seed, p->Token);
	}
	p = p->Parent;
	}
	return seed;
	}
	};

	template<>
	struct hash<tuple<Prefix*, token_t>>
	{
	size_t operator()(const tuple<Prefix*, token_t> &t) const noexcept
	{
	size_t seed = 0;
	hash_combine(seed, get<0>(t));
	hash_combine(seed, get<1>(t));
	return seed;
	}
	};

	template<>
	struct equal_to<Prefix*>
	{
	bool operator()(const Prefix a, const Prefix b) const noexcept
	{
	while (a != nullptr && b != nullptr) {
	if (a->Token != b->Token) {
	return false;
	}
	a = a->Parent;
	b = b->Parent;
	}
	// If one chain is shorter than the other
	return a == b;
	}
	};
	}

	inline size_t Prefix::size() const
	{
	size_t ret = 0;
	auto p = this;
	while (p != nullptr) {
	ret += 1;
	p = p->Parent;
	}
	return ret;
	}


	class PrefixAllocator
	{
	public:
	PrefixAllocator() = default;
	~PrefixAllocator();

	template<typename ...Args>
	Prefix *GetPrefix(Args&& ...ctorArgs);

	private:
	void AllocateNextBuffer();

	std::list<Prefix*> m_buffers;
	size_t m_allocSize = 0;
	size_t m_currOff = 0;
	};

	inline PrefixAllocator::~PrefixAllocator()
	{
	for (auto p : m_buffers) {
	// Prefix is a POD, and are allocated without initializing
	// to prevent redundant work upfront
	// delete[] p;
	free(p);
	}
	}

	inline void PrefixAllocator::AllocateNextBuffer()
	{
	size_t nextSize = m_allocSize == 0 ? 1000 : 2 * m_allocSize;

	// Using malloc here to prevent the ctor of Prefix being called for each item.
	// Instead, the ctor will be called upon first access using GetPrefix
	auto pBuff = reinterpret_cast<Prefix>(malloc(sizeof(Prefix) nextSize));

	m_buffers.push_back(pBuff);

	m_allocSize = nextSize;
	m_currOff = 0;
	}

	template<typename ...Args>
	Prefix *PrefixAllocator::GetPrefix(Args&& ...ctorArgs)
	{
	if (m_currOff == m_allocSize) {
	AllocateNextBuffer();
	}

	auto buff = m_buffers.back() + m_currOff;

	auto ret = new (buff) Prefix(std::forward<Args>(ctorArgs)...);

	++m_currOff;

	return ret;
	}