Method Tokenizer

src/tokenizer.cpp:3–42

#include "tokenizer.h"

Tokenizer::Tokenizer(const YALMData& data) {
  this->bos_id = std::stoi(data.metadata.at("bos_token_id").get<std::string>());
  this->eos_id = std::stoi(data.metadata.at("eos_token_id").get<std::string>());
  // TODO: figure out edge cases:
  // Q: should `vocab` include byte fallback tokens?
  // Q: should `vocab` include special tokens, e.g. '<unk>', '<s>', '</s>'?
  // TODO: avoid copy by using std::string_view
  const Tensor& tokens_tensor = data.tensors.at("tokenizer.tokens");
  char* tokens_tensor_end = (char*)tokens_tensor.data + tokens_tensor.size;
  for (char* ptr = (char*)tokens_tensor.data; ptr < tokens_tensor_end; ptr++) {
    char* s = ptr;
    while (*ptr != '\0' && ptr < tokens_tensor_end) {
      ptr++;
    }
    vocab.emplace_back(s, ptr - s);
  }
  for (size_t i = 0; i < vocab.size(); i++) {
    if (vocab[i] == "<0x00>") {
      byte_fallback_start = i;
    } else if (vocab[i] == "<|eot_id|>" || vocab[i] == "<|end|>" || vocab[i] == "<|im_end|>") {
      eot_id = i;
    }
  }
  // init byte_pieces
  for (size_t i = 0; i < 256; i++) {
    byte_pieces[i] = (char)i;
  }
  // init vocab trie
  for (size_t i = 0; i < vocab.size(); i++) {
    const std::string& word = vocab[i];
    TokenTrie* p = &vocab_trie;
    for (char c : word) {
      if (p->children.count(c) == 0) {
        p->children[c] = std::make_shared<TokenTrie>();
      }
      p = p->children[c].get();
    }
    p->token_id = i;
  }
}

std::string Tokenizer::decode_one(int prev_token, int token) const {
  const std::string& piece = vocab[token];