vocabulary.hpp

Go to the documentation of this file.

 
#pragma once
 
#include <cstdint>
#include <string>
#include <unordered_map>
#include <vector>
 
namespace pulsatrix {
 
class Vocabulary {
public:
    static constexpr int64_t kUnkIndex = 0;
    static constexpr const char* kUnkToken = "<unk>";
 
    explicit Vocabulary(std::vector<std::string> ranked_tokens);
 
    [[nodiscard]] int64_t IndexOf(const std::string& token) const;
 
    [[nodiscard]] const std::string& TokenAt(int64_t index) const;
 
    [[nodiscard]] int64_t size() const { return static_cast<int64_t>(tokens_.size()); }
 
private:
    std::vector<std::string> tokens_;
    std::unordered_map<std::string, int64_t> token_to_index_;
};
 
[[nodiscard]] Vocabulary BuildVocabulary(const std::vector<std::vector<std::string>>& tokenized_corpus,
                                          int64_t max_vocab_size = -1);
 
}  // namespace pulsatrix