tokenizer.h

Go to the documentation of this file.

#pragma once
 
#include <limits>
#include <map>
#include <memory>
#include <string>
#include <unordered_map>
#include <unordered_set>
#include <utility>
#include <vector>
#include <queue>
#include <functional>
 
#include "gguf_structs.h"
#include "logger.h"
#include "model.h"
 
 
 
// Helper struct to represent segments during BPE tokenization
struct llm_symbol {
  using index = int;
  index prev;              // index of the previous symbol in the linked list
  index next;              // index of the next symbol in the linked list
  const char * text;       // pointer to the start of the symbol's text in the original string
  size_t n;                // length of the symbol's text
};
 
// Helper struct representing a potential byte pair merge
struct llm_bigram_bpe {
    // Comparator for the priority queue: higher rank first, then lower left index
    struct comparator {
        bool operator()(const llm_bigram_bpe & l, const llm_bigram_bpe & r) const {
            // Prioritize lower rank (higher priority in BPE merges)
            // If ranks are equal, prioritize the one starting earlier (lower left index)
            return l.rank > r.rank || (l.rank == r.rank && l.left > r.left);
        }
    };
 
    using queue_storage = std::vector<llm_bigram_bpe>;
    // Define a min-priority queue based on the comparator
    using queue = std::priority_queue<llm_bigram_bpe, queue_storage, comparator>; 
 
    llm_symbol::index left;  // index of the left symbol in the pair
    llm_symbol::index right; // index of the right symbol in the pair
    std::string text;        // the merged text of the pair (for checking against merges map)
    int rank;                // rank of the merge (lower is better)
    size_t size;             // size of the merged text (for validation)
};
 
 
 
class Tokenizer {
 public:
  enum PreTokenizeMethod {
    DEFAULT,     
    LLAMA_REGEX  
  };
 
  enum class Type {
    UNKNOWN,
    SENTENCEPIECE_BPE,
    TIKTOKEN_BPE
  };
 
  Tokenizer(const std::string& vocab_path, const std::string& model_path, const ModelConfig& config);
 
  explicit Tokenizer(const GGUFData& gguf_data, const ModelConfig& config);
 
  std::vector<std::string> tokenize(const std::string& text) const;
 
 
  std::vector<std::string> ids_to_tokens(const std::vector<int>& ids) const;
 
  std::string detokenize(const std::vector<std::string>& tokens) const;
 
  std::vector<int> encode(
      const std::string& text, bool add_bos = true, bool add_eos = false,
      PreTokenizeMethod pre_tok_override = PreTokenizeMethod::DEFAULT) const;
 
  std::string decode(const std::vector<int>& ids,
                     bool skip_special_tokens = true) const;
 
  std::string apply_chat_template(const std::string& user_prompt,
                                  const std::string& system_message,
                                  const ModelConfig& config) const;
 
  int vocab_size() const;
 
  bool is_added_token(int id) const;
 
  int bos_token_id() const { return bos_token_id_; }
 
  int eos_token_id() const { return eos_token_id_; }
 
  int pad_token_id() const { return pad_token_id_; }
 
  int unk_token_id() const { return unk_token_id_; }
 
  const std::string& get_gguf_chat_template() const;
 
 private:
  void load_vocab_from_json(const std::string& vocab_path,
                            std::unordered_map<std::string, int>& token_to_id,
                            std::vector<std::string>& id_to_token);
 
  void load_bpe_merges_from_json(const std::string& model_path);
 
  void load_sentencepiece_model(const std::string& model_path);
 
 
  // Token mappings
  std::unordered_map<std::string, int> token_to_id_;    
  std::vector<std::string> id_to_token_;                
  std::unordered_map<std::string, int> bpe_merges_;     
  std::vector<std::string> tiktoken_merges_list_;       
  std::vector<float> token_scores_;                      
  std::vector<int32_t> token_types_;                    
  ModelConfig::TokenizerFamily tokenizer_family_ = ModelConfig::TokenizerFamily::UNKNOWN;
  bool initialized_from_gguf_ = false;                   
  std::unordered_map<std::string, int> added_tokens_;   
  // Special tokens
  std::string unk_token_;    
  std::string bos_token_;    
  std::string eos_token_;    
  std::string pad_token_;    
  // Special token IDs
  int unk_token_id_ = 0;     
  int bos_token_id_ = -1;    
  int eos_token_id_ = -1;    
  int pad_token_id_ = -1;    
  bool sentencepiece_model_loaded_ = false;              
  std::string pre_tok_type_ = "unknown";                
  std::unordered_map<int, std::string> id_to_added_token_; 
  std::unordered_set<std::string> chat_template_special_tokens; 
  std::unordered_map<char, int> byte_char_to_id_;       
  Type type_ = Type::UNKNOWN;
  std::string gguf_chat_template_;                    
  // SentencePiece specific helper methods (reinstated)
  std::vector<std::string> bpe_tokenize(const std::string& text) const;
  std::vector<std::string> bpe_tokenize_from_scores(const std::string& text) const;
  std::vector<int> tokens_to_ids(const std::vector<std::string>& tokens) const;
  std::string decode_sentencepiece(const std::vector<int>& ids, bool skip_special_tokens) const;
  std::string capitalize_first_letter(std::string s) const;
 
 
  // New BPE tokenization method for TikToken (Llama 3) path
  std::vector<int> bpe_tokenize_to_ids(const std::string& text,
                                       bool add_bos_token_param,
                                       bool add_eos_token_param,
                                       bool ignore_merges_param) const;
 
  // Helper for the new BPE tokenization path
  void add_bigram_to_queue_refactored(const char* text_data_base,
                                      const std::vector<llm_symbol>& symbols,
                                      llm_symbol::index first_symbol_idx,
                                      std::priority_queue<std::pair<int, int>,
                                                          std::vector<std::pair<int, int>>,
                                                          std::greater<std::pair<int, int>>>& work_queue) const;
 
  int find_bpe_rank(const std::string & token_left, const std::string & token_right) const;
  
};