executorch/examples/models/parakeet/tokenizer_utils.cpp at main · JakeStevens/executorch · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
#include "tokenizer_utils.h"

#include <exception>

#include <executorch/extension/llm/tokenizers/third-party/llama.cpp-unicode/include/unicode.h>
#include <executorch/runtime/platform/log.h>
#include <pytorch/tokenizers/tokenizer.h>

namespace {

// SentencePiece's word-boundary marker, spelled as UTF-8 bytes so this remains
// a const char[] literal when compiled as C++20.
constexpr char kSentencePieceWordBoundary[] = "\xE2\x96\x81";

bool is_whitespace_only(const std::string& token) {
  if (token.empty()) {
    return true;
  }

  try {
    const auto codepoints = unicode_cpts_from_utf8(token);
    for (const auto cp : codepoints) {
      if (!unicode_cpt_flags(cp).is_whitespace) {
        return false;
      }
    }
    return true;
  } catch (const std::exception&) {
    return false;
  }
}

bool is_special_token(const std::string& token) {
  if (token.size() >= 2 && token.front() == '[' && token.back() == ']') {
    return true;
  }
  if (token.size() >= 2 && token.front() == '<' && token.back() == '>') {
    return true;
  }
  if (token.rfind("##", 0) == 0) {
    return true;
  }
  if (token.rfind(kSentencePieceWordBoundary, 0) == 0) {
    return true;
  }
  if (is_whitespace_only(token)) {
    return true;
  }
  return false;
}

} // namespace

namespace parakeet::tokenizer_utils {

std::unordered_set<std::string> derive_supported_punctuation(
    const tokenizers::Tokenizer& tokenizer) {
  std::unordered_set<std::string> punctuation;

  const int32_t vocab_size = tokenizer.vocab_size();
  for (int32_t id = 0; id < vocab_size; id++) {
    const auto piece_result = tokenizer.id_to_piece(static_cast<TokenId>(id));
    if (!piece_result.ok()) {
      continue;
    }
    const std::string& piece = piece_result.get();
    if (is_special_token(piece)) {
      continue;
    }

    try {
      const auto codepoints = unicode_cpts_from_utf8(piece);
      for (const auto cp : codepoints) {
        if (unicode_cpt_flags(cp).is_punctuation) {
          punctuation.insert(unicode_cpt_to_utf8(cp));
        }
      }
    } catch (const std::exception&) {
      ET_LOG(
          Error,
          "Failed to decode token piece '%s' to codepoints",
          piece.c_str());
    }
  }

  return punctuation;
}

std::string decode_token_sequence(
    const std::vector<TokenId>& tokens,
    const tokenizers::Tokenizer& tokenizer) {
  std::string result;
  TokenId prev_token = tokenizer.bos_tok();
  for (const TokenId token : tokens) {
    auto decode_result = tokenizer.decode(prev_token, token);
    if (decode_result.ok()) {
      result += decode_result.get();
    }
    prev_token = token;
  }
  return result;
}

std::string decode_token_sequence(
    const std::vector<Token>& decoded_tokens,
    const tokenizers::Tokenizer& tokenizer) {
  std::vector<TokenId> token_ids;
  token_ids.reserve(decoded_tokens.size());
  for (const auto& tok : decoded_tokens) {
    token_ids.push_back(tok.id);
  }
  return decode_token_sequence(token_ids, tokenizer);
}

} // namespace parakeet::tokenizer_utils