react-native-executorch 0.4.7 → 0.5.1-rc.0

package/third-party/ios/ExecutorchLib/ExecutorchLib/tokenizer/base64.h

@@ -0,0 +1,202 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+// @lint-ignore-every LICENSELINT
+/**************************************************************************
+   Copyright (c) 2023 sewenew
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.
+ *************************************************************************/
+
+#pragma once
+
+#include <cassert>
+#include <executorch/runtime/core/error.h>
+#include <executorch/runtime/core/result.h>
+#include <executorch/runtime/platform/assert.h>
+#include <string>
+#include <string_view>
+
+namespace executorch {
+namespace extension {
+namespace llm {
+using Error = executorch::runtime::Error;
+template <typename T> using Result = executorch::runtime::Result<T>;
+
+namespace base64 {
+
+Result<std::string> decode(const std::string_view &input);
+
+namespace detail {
+
+constexpr uint32_t DECODE_TABLE[] = {
+    255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255,
+    255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255,
+    255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 62,  255,
+    255, 255, 63,  52,  53,  54,  55,  56,  57,  58,  59,  60,  61,  255, 255,
+    255, 255, 255, 255, 255, 0,   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,  255, 255, 255, 255, 255, 255, 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,  255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255,
+    255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255,
+    255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255,
+    255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255,
+    255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255,
+    255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255,
+    255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255,
+    255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255,
+    255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255,
+    255};
+
+inline Error validate(uint32_t v) {
+  ET_CHECK_OR_RETURN_ERROR(v != 255, InvalidArgument, "invalid char");
+  return Error::Ok;
+}
+
+inline Error decode(const std::string_view &input, std::string &output) {
+  ET_CHECK_OR_RETURN_ERROR(input.size() == 4, InvalidArgument,
+                           "input length must be 4, got %zu", input.size());
+
+  uint32_t val = 0;
+
+  uint8_t c = input[0];
+  auto v = DECODE_TABLE[c];
+  ET_CHECK_OK_OR_RETURN_ERROR(validate(v));
+  val = v;
+
+  c = input[1];
+  v = DECODE_TABLE[c];
+  ET_CHECK_OK_OR_RETURN_ERROR(validate(v));
+  val = (val << 6) | v;
+
+  c = input[2];
+  v = DECODE_TABLE[c];
+  ET_CHECK_OK_OR_RETURN_ERROR(validate(v));
+  val = (val << 6) | v;
+
+  c = input[3];
+  v = DECODE_TABLE[c];
+  ET_CHECK_OK_OR_RETURN_ERROR(validate(v));
+  val = (val << 6) | v;
+
+  output.push_back(static_cast<char>((val >> 16) & 0xFF));
+  output.push_back(static_cast<char>((val >> 8) & 0xFF));
+  output.push_back(static_cast<char>(val & 0xFF));
+  return Error::Ok;
+}
+
+inline Error decode_1_padding(const std::string_view &input,
+                              std::string &output) {
+  ET_CHECK_OR_RETURN_ERROR(input.size() == 3, InvalidArgument,
+                           "input length must be 3, got %zu", input.size());
+
+  uint32_t val = 0;
+
+  uint8_t c = input[0];
+  auto v = DECODE_TABLE[c];
+  ET_CHECK_OK_OR_RETURN_ERROR(validate(v));
+  val = v;
+
+  c = input[1];
+  v = DECODE_TABLE[c];
+  ET_CHECK_OK_OR_RETURN_ERROR(validate(v));
+  val = (val << 6) | v;
+
+  c = input[2];
+  v = DECODE_TABLE[c];
+  ET_CHECK_OK_OR_RETURN_ERROR(validate(v));
+  val = (val << 6) | v;
+
+  output.push_back(static_cast<char>((val >> 10) & 0xFF));
+  output.push_back(static_cast<char>((val >> 2) & 0xFF));
+  return Error::Ok;
+}
+
+inline Error decode_2_padding(const std::string_view &input,
+                              std::string &output) {
+  ET_CHECK_OR_RETURN_ERROR(input.size() == 2, InvalidArgument,
+                           "input length must be 2, got %zu", input.size());
+
+  uint32_t val = 0;
+
+  uint8_t c = input[0];
+  auto v = DECODE_TABLE[c];
+  ET_CHECK_OK_OR_RETURN_ERROR(validate(v));
+  val = v;
+
+  c = input[1];
+  v = DECODE_TABLE[c];
+  ET_CHECK_OK_OR_RETURN_ERROR(validate(v));
+  val = (val << 6) | v;
+
+  output.push_back(static_cast<char>((val >> 4) & 0xFF));
+  return Error::Ok;
+}
+
+} // namespace detail
+
+inline Result<std::string> decode(const std::string_view &input) {
+  ET_CHECK_OR_RETURN_ERROR(!input.empty(), InvalidArgument, "empty input");
+
+  // Faster than `input.size() % 4`.
+  ET_CHECK_OR_RETURN_ERROR(
+      (input.size() & 3) == 0 && input.size() >= 4, InvalidArgument,
+      "input length must be larger than 4 and is multiple of 4, got %zu",
+      input.size());
+
+  std::string output;
+  output.reserve(input.size() / 4 * 3);
+  auto idx = 0U;
+  for (; idx < input.size() - 4; idx += 4) {
+    ET_CHECK_OK_OR_RETURN_ERROR(detail::decode(input.substr(idx, 4), output));
+  }
+
+  // Last 4 bytes. Might contain paddings.
+  if (input[idx + 3] == '=') {
+    if (input[idx + 2] == '=') {
+      // Tow paddings.
+      ET_CHECK_OK_OR_RETURN_ERROR(
+          detail::decode_2_padding(input.substr(idx, 2), output));
+    } else {
+      // One padding.
+      ET_CHECK_OK_OR_RETURN_ERROR(
+          detail::decode_1_padding(input.substr(idx, 3), output));
+    }
+  } else {
+    // No padding.
+    ET_CHECK_OK_OR_RETURN_ERROR(detail::decode(input.substr(idx, 4), output));
+  }
+
+  return output;
+}
+
+} // namespace base64
+
+} // namespace llm
+} // namespace extension
+} // namespace executorch
+
+namespace torch {
+namespace executor {
+namespace base64 {
+// TODO(T197294990): Remove these deprecated aliases once all users have moved
+// to the new `::executorch` namespaces.
+using ::executorch::extension::llm::base64::decode;
+} // namespace base64
+} // namespace executor
+} // namespace torch

package/third-party/ios/ExecutorchLib/ExecutorchLib/tokenizer/bpe_tokenizer.cpp

@@ -0,0 +1,313 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#include "bpe_tokenizer.h"
+
+#include <cstring>
+
+using ::executorch::runtime::Error;
+using ::executorch::runtime::Result;
+
+namespace executorch {
+namespace extension {
+namespace llm {
+
+static int compare_tokens(const void *a, const void *b) {
+  if (((TokenIndex *)a)->str == nullptr) {
+    return -1;
+  }
+  if (((TokenIndex *)b)->str == nullptr) {
+    return 1;
+  }
+  return strcmp(((TokenIndex *)a)->str, ((TokenIndex *)b)->str);
+}
+
+BPETokenizer::BPETokenizer() : Tokenizer() {
+  for (int i = 0; i < 256; i++) {
+    byte_pieces_[i * 2] = (unsigned char)i;
+    byte_pieces_[i * 2 + 1] = '\0';
+  }
+}
+
+/**
+ * @brief Load the tokenizer from a file. The tokenizer file contains the
+ * vocabulary and scores. The format is: the first integer is the maximum
+ * token length, followed by a list of (word_len, word) pairs. Here we
+ * are reading all the vocabulary into memory and keep it sorted for fast
+ * lookup.
+ *
+ * @param tokenizer_path The path to the tokenizer file.
+ * @return Error
+ */
+Error BPETokenizer::load(const std::string &tokenizer_path) {
+  if (initialized_) {
+    ET_LOG(Info, "Tokenizer already initialized");
+    return Error::Ok;
+  }
+  // read in the file
+  FILE *file = fopen(tokenizer_path.c_str(), "rb");
+  if (!file) {
+    ET_LOG(Error, "couldn't load %s", tokenizer_path.c_str());
+    return Error::InvalidArgument;
+  }
+  int32_t metadata[4];
+  for (int i = 0; i < 4; i++) {
+    if (fread(metadata + i, sizeof(int32_t), 1, file) != 1) {
+      ET_LOG(Error,
+             "Failed to read the metadata at position %d, the tokenizer file "
+             "is not valid!",
+             i);
+      return Error::InvalidArgument;
+    }
+  }
+
+  // now we have two vocab_sizes one from the model and another from the
+  // tokenizer file.
+  int32_t tokenizer_vocab_size = metadata[0];
+  vocab_size_ = tokenizer_vocab_size;
+  bos_tok_ = metadata[1];
+  eos_tok_ = metadata[2];
+  max_token_length_ = metadata[3];
+
+  // allocate space for the vocabulary
+  vocab_ = std::make_unique<char *[]>(vocab_size_);
+  vocab_scores_ = std::make_unique<float[]>(vocab_size_);
+  sorted_vocab_ = std::make_unique<TokenIndex[]>(vocab_size_);
+
+  // read in the vocabulary
+  for (int i = 0; i < vocab_size_; i++) {
+    if (fread(vocab_scores_.get() + i, sizeof(float), 1, file) != 1) {
+      // This is allowed, we just pad the rest of the vocab with <pad> strings
+      std::string padding = "<pad>";
+      vocab_[i] = new char[padding.length() + 1];
+      strcpy(vocab_[i], padding.c_str());
+      vocab_[i][padding.length()] = '\0';
+      continue;
+    }
+    int32_t len;
+    if (fread(&len, sizeof(int32_t), 1, file) != 1) {
+      ET_LOG(Error, "Failed to read the length of the word at index %d", i);
+      return Error::InvalidArgument;
+    }
+    vocab_[i] = new char[len + 1];
+    if (fread(vocab_[i], len, 1, file) != 1) {
+      ET_LOG(Error, "Failed to read the word, total length %d, index %d\n", len,
+             i);
+      return Error::InvalidArgument;
+    }
+    vocab_[i][len] = '\0'; // add the string terminating token
+  }
+  fclose(file);
+
+  for (int32_t i = 0; i < vocab_size_; i++) {
+    sorted_vocab_[i].str = vocab_[i];
+    sorted_vocab_[i].id = i;
+  }
+  qsort(sorted_vocab_.get(), vocab_size_, sizeof(TokenIndex), compare_tokens);
+
+  initialized_ = true;
+  return Error::Ok;
+}
+
+BPETokenizer::~BPETokenizer() {
+  for (int i = 0; i < vocab_size_; i++) {
+    delete[] vocab_[i];
+  }
+}
+
+/**
+ * @brief Decode a token into string.
+ *
+ * @param prev_token The previous token.
+ * @param token The current token.
+ * @return Result<std::string> A pointer to the string representation of the
+ * token.
+ */
+Result<std::string> BPETokenizer::decode(uint64_t prev_token,
+                                         uint64_t token) const {
+  ET_CHECK_OK_OR_RETURN_ERROR(Tokenizer::decode_verify(token));
+  const char *piece = vocab_[token];
+  // following BOS token, sentencepiece decoder strips any leading
+  // whitespace
+  if (prev_token == bos_tok_ && piece[0] == ' ') {
+    piece++;
+  }
+  // careful, some tokens designate raw bytes, and look like e.g. '<0x01>'
+  // parse this and convert and return the actual byte
+  unsigned char byte_val;
+  if (sscanf(piece, "<0x%02hhX>", &byte_val) == 1) {
+    piece = (char *)byte_pieces_ + byte_val * 2;
+  }
+  std::string res(piece);
+  return res;
+}
+
+static int32_t str_lookup(const char *str, TokenIndex *sorted_vocab,
+                          int32_t vocab_size) {
+  // efficiently find the perfect match for str in vocab, return its index or -1
+  // if not found
+  TokenIndex tok = {.str = str}; // acts as the key to search for
+  TokenIndex *res = (TokenIndex *)bsearch(&tok, sorted_vocab, vocab_size,
+                                          sizeof(TokenIndex), compare_tokens);
+  return res != nullptr ? res->id : -1;
+}
+
+/**
+ * @brief Encode a string into a sequence of tokens.
+ *
+ * @param text The string to be encoded.
+ * @param bos The number of BOS to prepend to the token list.
+ * @param eos The number of EOS to append to the token list.
+ * @param tokens The output tokens.
+ * @param n_tokens The number of tokens.
+ * @return Result<std::vector<uint64_t>>
+ */
+Result<std::vector<uint64_t>>
+BPETokenizer::encode(const std::string &text, int8_t bos, int8_t eos) const {
+  if (!initialized_) {
+    ET_LOG(Error, "Tokenizer not initialized");
+    return Error::NotSupported;
+  }
+  // encode the string text (input) into an upper-bound preallocated tokens[]
+  // array bos != 0 means prepend the BOS token (=1), eos != 0 means append the
+  // EOS token (=2)
+  if (text.empty()) {
+    ET_LOG(Error, "cannot encode empty text");
+    return Error::InvalidArgument;
+  }
+
+  // create a temporary buffer that will store merge candidates of always two
+  // consecutive tokens *2 for concat, +1 for null terminator +2 for UTF8 (in
+  // case max_token_length is 1)
+  char *str_buffer = new char[max_token_length_ * 2 + 1 + 2];
+  size_t str_len = 0;
+
+  // start at 0 tokens
+  std::vector<uint64_t> tokens;
+
+  // add optional BOS token, if desired
+  if (bos >= 0) {
+    while (bos--) {
+      tokens.push_back(bos_tok_);
+    }
+  } else {
+    ET_LOG(Error, "bos %d should be >= 0", bos);
+    return Error::InvalidArgument;
+  }
+
+  // add_dummy_prefix is true by default
+  // so prepend a dummy prefix token to the input string, but only if text != ""
+  // TODO: pretty sure this isn't correct in the general case but I don't have
+  // the energy to read more of the sentencepiece code to figure out what it's
+  // doing
+  const char *space = " ";
+  if (text[0] != '\0') {
+    int dummy_prefix = str_lookup(space, sorted_vocab_.get(), vocab_size_);
+    tokens.push_back(dummy_prefix);
+  }
+
+  // Okay UTF-8 time. This will get messy. Here is the reference from Wikipedia:
+  // Code point ↔ UTF-8 conversion
+  // First code point  Last code point  Byte 1  Byte 2  Byte 3  Byte 4
+  // U+0000  U+007F      0xxxxxxx
+  // U+0080  U+07FF      110xxxxx  10xxxxxx
+  // U+0800  U+FFFF      1110xxxx  10xxxxxx  10xxxxxx
+  // U+10000  U+10FFFF    11110xxx  10xxxxxx  10xxxxxx  10xxxxxx
+
+  // process the raw (UTF-8) byte sequence of the input string
+  for (const char *c = text.c_str(); *c != '\0'; c++) {
+    // reset buffer if the current byte is ASCII or a leading byte
+    // 0xC0 is 11000000, so (*c & 0xC0) keeps the first 2 bits and zeros the
+    // rest 0x80 is 10000000 in UTF-8, all continuation bytes start with "10" in
+    // first two bits so in English this is: "if this byte is not a continuation
+    // byte"
+    if ((*c & 0xC0) != 0x80) {
+      // this byte must be either a leading byte (11...) or an ASCII char
+      // (0x...)
+      // => reset our location, as we're starting a new UTF-8 codepoint
+      str_len = 0;
+    }
+
+    // append the current byte to the buffer
+    str_buffer[str_len++] =
+        *c; // ++ is post-increment, incremented after this line
+    str_buffer[str_len] = '\0';
+
+    // while the next character is a continuation byte, continue appending
+    // but if there are too many of them, just stop to avoid overruning
+    // str_buffer size.
+    if ((*(c + 1) & 0xC0) == 0x80 && str_len < 4) {
+      continue;
+    }
+
+    // ok c+1 is not a continuation byte, so we've read in a full codepoint
+    int id = str_lookup(str_buffer, sorted_vocab_.get(), vocab_size_);
+    if (id != -1) {
+      // we found this codepoint in vocab, add it as a token
+      tokens.push_back(id);
+    } else {
+      // byte_fallback encoding: just encode each byte as a token
+      // +3 is here because the first 3 vocab elements are <unk>, <s>, </s>
+      // so the individual bytes only start at index 3
+      for (int i = 0; i < str_len; i++) {
+        tokens.push_back((unsigned char)str_buffer[i] + 3);
+      }
+    }
+    str_len = 0; // protect against a sequence of stray UTF8 continuation bytes
+  }
+
+  // merge the best consecutive pair each iteration, according the scores in
+  // vocab_scores
+  while (1) {
+    float best_score = -1e10;
+    int best_id = -1;
+    int best_idx = -1;
+
+    for (int i = 0; i < tokens.size() - 1; i++) {
+      // check if we can merge the pair (tokens[i], tokens[i+1])
+      snprintf(str_buffer, max_token_length_ * 2 + 3, "%s%s", vocab_[tokens[i]],
+               vocab_[tokens[i + 1]]);
+      int id = str_lookup(str_buffer, sorted_vocab_.get(), vocab_size_);
+      if (id != -1 && vocab_scores_[id] > best_score) {
+        // this merge pair exists in vocab! record its score and position
+        best_score = vocab_scores_[id];
+        best_id = id;
+        best_idx = i;
+      }
+    }
+
+    if (best_idx == -1) {
+      break; // we couldn't find any more pairs to merge, so we're done
+    }
+
+    // merge the consecutive pair (best_idx, best_idx+1) into new token best_id
+    tokens[best_idx] = best_id;
+    // delete token at position best_idx+1, shift the entire sequence back 1
+    for (int i = best_idx + 1; i < tokens.size() - 1; i++) {
+      tokens[i] = tokens[i + 1];
+    }
+    tokens.pop_back(); // token length decreased
+  }
+
+  // add optional EOS (=2) token, if desired
+  if (eos >= 0) {
+    while (eos--) {
+      tokens.push_back(eos_tok_);
+    }
+  } else {
+    ET_LOG(Error, "eos %d should be >= 0", eos);
+    return Error::InvalidArgument;
+  }
+
+  delete[] str_buffer;
+  return Result(tokens);
+}
+
+} // namespace llm
+} // namespace extension
+} // namespace executorch

package/third-party/ios/ExecutorchLib/ExecutorchLib/tokenizer/bpe_tokenizer.h

@@ -0,0 +1,57 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#pragma once
+
+#include "tokenizer.h"
+#include <memory>
+
+namespace executorch {
+namespace extension {
+namespace llm {
+
+struct TokenIndex {
+  const char *str;
+  int32_t id;
+};
+
+// A simple Byte Pair Encoding (BPE) Tokenizer. Note that the current C++ code
+// won't work with this class, it needs to go through tokenizer.py first.
+class BPETokenizer : public Tokenizer {
+public:
+  explicit BPETokenizer();
+  ~BPETokenizer() override;
+
+  ::executorch::runtime::Error load(const std::string &tokenizer_path) override;
+
+  ::executorch::runtime::Result<std::vector<uint64_t>>
+  encode(const std::string &input, int8_t bos, int8_t eos) const override;
+
+  ::executorch::runtime::Result<std::string>
+  decode(uint64_t prev_token, uint64_t token) const override;
+
+private:
+  std::unique_ptr<char *[]> vocab_ = nullptr;
+  std::unique_ptr<float[]> vocab_scores_ = nullptr;
+  std::unique_ptr<TokenIndex[]> sorted_vocab_ = nullptr;
+  unsigned int max_token_length_ = 0;
+  unsigned char byte_pieces_[512]; // stores all single-byte strings
+};
+
+} // namespace llm
+} // namespace extension
+} // namespace executorch
+
+namespace torch {
+namespace executor {
+// TODO(T197294990): Remove these deprecated aliases once all users have moved
+// to the new `::executorch` namespaces.
+using ::executorch::extension::llm::BPETokenizer;
+using ::executorch::extension::llm::TokenIndex;
+} // namespace executor
+} // namespace torch

package/third-party/ios/ExecutorchLib/ExecutorchLib/tokenizer/llama_tiktoken.cpp

@@ -0,0 +1,78 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#include "llama_tiktoken.h"
+
+namespace example {
+
+using ::executorch::extension::llm::Tiktoken;
+
+namespace {
+static constexpr int32_t kSpecialTokensSize = 256;
+static constexpr size_t kBOSTokenIndex = 0;
+static constexpr size_t kEOSTokenIndex = 1;
+
+static inline std::unique_ptr<std::vector<std::string>>
+_get_default_special_tokens() {
+  auto special_tokens =
+      std::make_unique<std::vector<std::string>>(std::vector<std::string>{
+          "<|begin_of_text|>", "<|end_of_text|>",
+          "<|reserved_special_token_0|>", "<|reserved_special_token_1|>",
+          "<|finetune_right_pad_id|>", "<|step_id|>", "<|start_header_id|>",
+          "<|end_header_id|>", "<|eom_id|>", "<|eot_id|>", "<|python_tag|>"});
+  // pad the rest of the special tokens with reserved tokens
+  ssize_t reserved_special_token_num = 2;
+  while (special_tokens->size() < kSpecialTokensSize) {
+    special_tokens->emplace_back("<|reserved_special_token_" +
+                                 std::to_string(reserved_special_token_num++) +
+                                 "|>");
+  }
+  return special_tokens;
+}
+
+static inline std::unique_ptr<std::vector<std::string>>
+_get_multimodal_special_tokens() {
+  auto special_tokens =
+      std::make_unique<std::vector<std::string>>(std::vector<std::string>{
+          "<|begin_of_text|>", "<|end_of_text|>",
+          "<|reserved_special_token_0|>", "<|reserved_special_token_1|>",
+          "<|reserved_special_token_2|>", "<|reserved_special_token_3|>",
+          "<|start_header_id|>", "<|end_header_id|>", "<|eom_id|>",
+          "<|eot_id|>", "<|image|>"});
+
+  // pad the rest of the special tokens with reserved tokens except the last
+  // one
+  ssize_t reserved_special_token_num = 4;
+  while (special_tokens->size() < kSpecialTokensSize - 1) {
+    special_tokens->emplace_back("<|reserved_special_token_" +
+                                 std::to_string(reserved_special_token_num++) +
+                                 "|>");
+  }
+
+  special_tokens->emplace_back("<|python_tag|>");
+
+  return special_tokens;
+}
+
+std::unique_ptr<std::vector<std::string>> _get_special_tokens(Version version) {
+  switch (version) {
+  case Version::Multimodal:
+    return _get_multimodal_special_tokens();
+  default:
+    return _get_default_special_tokens();
+  }
+}
+
+} // namespace
+
+std::unique_ptr<Tiktoken> get_tiktoken_for_llama(Version version) {
+  return std::make_unique<Tiktoken>(_get_special_tokens(version),
+                                    kBOSTokenIndex, kEOSTokenIndex);
+}
+
+} // namespace example

package/third-party/ios/ExecutorchLib/ExecutorchLib/tokenizer/llama_tiktoken.h

@@ -0,0 +1,23 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#pragma once
+
+#include "tiktoken.h"
+
+namespace example {
+
+enum class Version {
+  Default,
+  Multimodal,
+};
+
+std::unique_ptr<::executorch::extension::llm::Tiktoken>
+get_tiktoken_for_llama(Version version = Version::Default);
+
+} // namespace example